TorchRec大量使用Jagged Tensor

在线体验各类最新模型,更有模型 免费Token 额度领取!
立即体验
简介: Jagged Tensor(锯齿张量)是专为变长序列设计的紧凑存储格式,用values+lengths/offsets替代padding,显著节省内存与计算。广泛应用于推荐系统中用户行为、多值标签等不等长特征处理,如HSTU模型中的拼接、拆分与矩阵乘法操作。

什么是 Jagged Tensor

Jagged Tensor(锯齿张量)是一种变长嵌套张量,用于表示每个样本长度不同的特征。与普通的 padded 2D tensor 不同,它用 values(拼接的值)+ offsets/lengths(每个样本的长度)来紧凑存储,避免大量 padding 浪费。


例子 1:用户历史点击序列(长度不一致)

假设一个 batch 有 3 个用户,他们的历史点击 item_id 分别为:

plaintext

用户A: 点击了 [101, 202, 303]         → 长度 3
用户B: 点击了 [404]                    → 长度 1  
用户C: 点击了 [505, 606, 707, 808]    → 长度 4

Padded 方式(浪费内存):

python

# 需要 pad 到最长 4,大量 0 浪费
tensor([[101, 202, 303,   0],
        [404,   0,   0,   0],
        [505, 606, 707, 808]])   # shape: (3, 4)

Jagged Tensor 方式(紧凑):

python

values  = [101, 202, 303, 404, 505, 606, 707, 808]   # 所有值拼接,长度 8
lengths = [3, 1, 4]                                     # 每个用户的序列长度
offsets = [0, 3, 4, 8]                                  # 累加偏移量


例子 2:多值标签特征(tag)

用户打的标签数量不同:

plaintext

用户A: tags = ["科技", "体育"]           → 长度 2
用户B: tags = ["美食", "旅行", "摄影"]   → 长度 3
用户C: tags = ["音乐"]                   → 长度 1

python

values  = [科技, 体育, 美食, 旅行, 摄影, 音乐]   # 拼接
lengths = [2, 3, 1]


例子 3:jagged_tensors.py 中三个 op 的实际场景

这三个 op 来源于 generative-recommenders(HSTU 模型),用于处理用户行为序列 + 候选 item 拼接的场景:

concat_2D_jagged — 拼接用户历史和候选 item

plaintext

用户A历史: [e1, e2, e3](3个embedding)    候选item: [c1, c2](2个embedding)
用户B历史: [e4](1个embedding)             候选item: [c3, c4, c5](3个embedding)

拼接后:

plaintext

用户A: [e1, e2, e3, c1, c2]       → 长度 3+2=5
用户B: [e4, c3, c4, c5]           → 长度 1+3=4

用 padded tensor 需要 shape (2, 5, dim),用户B 有 1 行 padding。

用 jagged tensor 只存 values: (9, dim) + offsets,零浪费。

split_2D_jagged — 从拼接结果中拆回用户/候选

HSTU 的 Transformer 对拼接后的序列做完 self-attention 后,需要把用户部分和候选部分拆回来:

plaintext

拼接的attention输出: [a1, a2, a3, a4, a5, a6, a7, a8, a9]
offsets_left  = [0, 3, 4]      ← 用户历史的偏移
offsets_right = [0, 2, 5]      ← 候选item的偏移
→ left:  用户表征
→ right: 候选打分

jagged_dense_bmm_broadcast_add — 变长序列的矩阵乘法

对每个用户的变长序列做投影:out = jagged × dense + bias

plaintext

jagged: (sum_B(M_i), K)  ← 所有用户序列拼接,总共 sum 个 token
dense:  (B, K, N)         ← 每个用户一个投影矩阵
bias:   (B, N)            ← 每个用户一个偏置
out:    (sum_B(M_i), N)   ← 变长输出

传统做法需要 pad → bmm → unpad,jagged 版本直接在紧凑格式上计算,省内存且更快。


为什么推荐系统特别需要 Jagged Tensor?

推荐系统中到处都是变长特征:

  • 用户点击历史(有人点了 5 个,有人点了 500 个)
  • 多值 ID 特征(用户标签、商品类目)
  • 序列特征中的多值子特征(每个行为关联多个属性)

如果全部 pad 到最大长度,内存和计算浪费巨大。Jagged Tensor 是推荐系统处理这类数据的标准做法,torchrec 的 KeyedJaggedTensor 也是基于同样的理念

相关实践学习
使用PAI+LLaMA Factory微调Qwen2-VL模型,搭建文旅领域知识问答机器人
使用PAI和LLaMA Factory框架,基于全参方法微调 Qwen2-VL模型,使其能够进行文旅领域知识问答,同时通过人工测试验证了微调的效果。
机器学习概览及常见算法
机器学习(Machine Learning, ML)是人工智能的核心,专门研究计算机怎样模拟或实现人类的学习行为,以获取新的知识或技能,重新组织已有的知识结构使之不断改善自身的性能,它是使计算机具有智能的根本途径,其应用遍及人工智能的各个领域。 本课程将带你入门机器学习,掌握机器学习的概念和常用的算法。
相关文章
|
3月前
|
存储 搜索推荐 PyTorch
为什么使用 TorchRec 训练和推理更快
本文结合TorchEasyRec实践,从四大维度解析推荐系统加速:1)KeyedJaggedTensor统一变长特征,实现Embedding批量融合查找;2)自动分布式分片突破单卡显存瓶颈;3)TrainPipelineSparseDist流水线并行,重叠通信与计算;4)fbgemm-gpu融合优化器,减少显存访问。端到端提升训练效率与扩展性。
507 9
|
Python
VSCode运行Python教程
VSCode运行Python教程
3301 0
VSCode运行Python教程
|
2月前
|
安全 网络协议 JavaScript
F5 BIG-IP 21.1 LTS 发布 - 多云安全和应用交付
BIG-IP 是硬件平台和软件解决方案的集合,提供专注于安全性、可靠性和性能的服务
145 5
F5 BIG-IP 21.1 LTS 发布 - 多云安全和应用交付
|
3月前
|
存储 人工智能 JavaScript
Prompt、Context、Harness:AI Agent 工程的三层架构解析
2023年重“Prompt”(如何说),2025年重“Context”(看到什么),2026年跃升至“Harness”(系统级约束与验证)。三者非替代而是分层:Prompt优化表达,Context管理信息环境,Harness构建可信执行系统——模型是马,Harness才是缰绳、马鞍与路。
1139 10
Prompt、Context、Harness:AI Agent 工程的三层架构解析
|
2月前
|
存储 传感器 人工智能
《2026鸿蒙NEXT纯血开发与AI辅助》第六章:「微距」项目启动——工程创建与整体架构设计-卓伊凡
《2026鸿蒙NEXT纯血开发与AI辅助》第六章启动「微距」项目:基于HarmonyOS 6.0 AR Engine的智能测量工具,支持单/多设备协同测量、端侧AI材质识别与安全数据管理,完成工程创建、模块化架构设计及Tab导航骨架搭建。
137 1
|
3月前
|
人工智能 JSON 自然语言处理
Vibe Coding 老翻车?可能是你的 AI 根本读不懂产品文档
Vibe Coding(氛围感编程)火爆Reddit/HN,主打自然语言驱动开发。但PDF需求文档常因排版复杂(表格错位、流程图丢失、多栏混乱)导致AI生成代码出错。解法:接入上海AI实验室开源的MinerU MCP Server,自动无损解析PDF为结构化Markdown,让AI真正“读懂”PRD,提升代码准确率。(239字)
460 6
|
3月前
|
数据采集 编解码 人工智能
从自回归到掩码生成:拆解零样本 TTS 在长视频配音场景下的技术路径
本文系统拆解零样本TTS在长视频配音中的工程落地路径:剖析自回归与非自回归范式取舍,详解掩码生成(mask-and-predict)这一新兴技术如何兼顾自然度、鲁棒性与效率;深入探讨零样本克隆、跨语种声纹一致性、画面对齐等关键工程挑战,并分享数据预处理、提示工程、合规风控等生产级实践经验。(239字)
237 9
|
3月前
|
前端开发 Java 调度
并发模型的对决——C++的线程库与异步世界的博弈
多线程编程从来不是C++的核心领域,但C++11将线程库纳入标准后,情况发生了根本改变。这个改变不仅仅是语言层面的——它标志着C++正式进入并发编程的竞技场,与Java、C#等语言直接竞争。
187 6
|
3月前
|
安全 编译器 开发者
移动语义的五张面孔——从Rvalue引用到完美转发的完整进化史
C++11引入的移动语义被认为是该标准中最重要的特性之一,它的影响力甚至超过了lambda表达式和auto关键字。但在实践中,移动语义的复杂性和隐蔽细节远远超出了大多数开发者的预期
124 3