TorchRec大量使用Jagged Tensor

简介: 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)是人工智能的核心,专门研究计算机怎样模拟或实现人类的学习行为,以获取新的知识或技能,重新组织已有的知识结构使之不断改善自身的性能,它是使计算机具有智能的根本途径,其应用遍及人工智能的各个领域。 本课程将带你入门机器学习,掌握机器学习的概念和常用的算法。
相关文章
|
20天前
|
存储 搜索推荐 PyTorch
为什么使用 TorchRec 训练和推理更快
本文结合TorchEasyRec实践,从四大维度解析推荐系统加速:1)KeyedJaggedTensor统一变长特征,实现Embedding批量融合查找;2)自动分布式分片突破单卡显存瓶颈;3)TrainPipelineSparseDist流水线并行,重叠通信与计算;4)fbgemm-gpu融合优化器,减少显存访问。端到端提升训练效率与扩展性。
176 9
|
28天前
|
人工智能 弹性计算 数据可视化
部署OpenClaw有哪些成本?附OpenClaw低成本部署指南
OpenClaw(“养龙虾”)是一款开源AI代理框架,可自动化文件处理、工作流与消息管理。本文详解其部署成本:软件免费,云服务器低至68元/年,阿里云百炼新用户享7000万Token免费额度,并提供一键图形化部署指南。
690 32
|
22天前
|
机器学习/深度学习 搜索推荐 数据处理
PAI-Rec推荐开发平台:企业级智能推荐解决方案,驱动业务全域增长
PAI-Rec是阿里云一站式推荐系统平台,集成多路召回、多目标精排(如DBMTL)、GPU加速推理与灵活迭代能力,已助力电商、直播、音视频等多行业提升点击率、转化率与ROI,实现高效、低成本、可自主演进的智能推荐。
190 16
|
22天前
|
人工智能 Linux API
VS Code 1.113 发布:Agent 与 Chat 体验全面升级!
VS Code 1.113 正式发布!聚焦AI开发体验升级:全面增强Agent能力(支持CLI/Claude代理的MCP、会话分支、嵌套子代理、调试日志),优化Chat体验(统一自定义编辑器、模型推理努力直调、图像预览查看器),大幅提升智能编程效率。
352 11
|
24天前
|
人工智能 JavaScript Java
【SpringAIAlibaba新手村系列】(10)Text to Voice 文本转语音技术
本文围绕 Spring AI Alibaba 1.1.2.2 的文本转语音实现展开,记录了基于 DashScopeAudioSpeechModel 与 stream() 的可运行方案。文章重点说明了模型、音色、输出格式与流式拼接音频文件的关键细节。
259 6
|
11天前
|
数据可视化 应用服务中间件 API
告别古法编程!我用Trae+阿里百炼Coding Plan,5分钟搭建一个Nginx日志可视化看板
利用闲置的阿里百炼Coding Plan,快速开发出一款单文件PHP版Nginx日志分析器。支持日志解析、IP/UA/Referer等多维统计、原始日志查看、关键词搜索及昼夜双模式,功能实用、部署简便。(239字)
|
21天前
|
人工智能 数据可视化 机器人
OpenClaw一键部署攻略,手把手教你 “养龙虾”!
还在为部署OpenClaw踩坑发愁?“养龙虾”其实超简单!本文奉上阿里云一键云端部署攻略:全程可视化、零代码,仅两步——买预装服务器+填API密钥,5分钟即可拥有专属AI数字员工!支持微信/钉钉协同、文件处理、日程管理、代码辅助等,新手友好,成本低廉(新用户首月9.9元+7000万Token免费额度)。
432 25
|
23天前
|
人工智能 数据挖掘 程序员
藏在Claude Code里的小惊喜!187种Loading状态词,告别单调编程等待
本文揭秘Claude Code中187种趣味又专业的Loading状态词(如Analyzing、Baking、Crunching等),覆盖编译、推理、数据处理等全场景,附中文翻译与自定义配置教程,让等待过程更愉悦、更高效!
|
18天前
|
人工智能 API 开发者
阿里云百炼Coding Plan售罄,抢不到怎么办?附抢购攻略及备选方案
阿里云百炼Coding Plan Pro版因Lite停售而一票难求,每日9:30限量抢购,几分钟即罄。本文提供实操抢购技巧,并推荐百炼平台作为高性价比备选:支持同款大模型、新用户赠100万Tokens、按量付费,开通即用。
678 9