摘要
随着 PyTorch 2.2 的发布，我们为现有的 PyTorch 库带来了许多改进。这些更新表明我们致力于在所有领域开发通用且可扩展的 API，以便社区更容易在 PyTorch 上构建生态系统项目。
| 最新稳定库版本 (完整列表)* | ||
| TorchArrow 0.1.0 | TorchRec 0.6.0 | TorchVision 0.17 |
| TorchAudio 2.2.0 | TorchServe 0.9.0 | TorchX 0.7.0 |
| TorchData 0.7.1 | TorchText 0.17.0 | PyTorch on XLA Devices 2.1 |
*要查看先前版本或（不稳定）每夜构建版本，请点击左上角‘搜索文档’上方的版本菜单。
TorchRL
特性：TorchRL 的离线强化学习数据中心
TorchRL 现在提供了离线强化学习和模仿学习的最大数据集中心之一，并且所有数据都采用单一数据格式（TED，即 TorchRL Episode Data 格式）。这使得在单个训练循环中可以轻松切换不同的数据源。现在还可以通过 ReplayBufferEnsemble 类轻松组合来自不同数据源的数据集。数据处理是完全可定制的。数据源包括模拟任务 (Minari, D4RL, VD4RL)、机器人数据集 (Roboset, OpenX Embodied dataset) 和游戏 (GenDGRL/ProcGen, Atari/DQN)。请在文档中查看这些内容。
除了这些更改之外，我们的回放缓冲区现在可以使用.dumps()方法转储到磁盘，该方法将使用 TensorDict API 将缓冲区序列化到磁盘，这比使用 torch.save 更快、更安全、更高效。
最后，现在可以在同一台机器上从独立进程中读取和写入回放缓冲区，而无需用户编写任何额外代码！
TorchRL2Gym 环境 API
为了方便 TorchRL 与现有代码库集成，并享受 TorchRL 环境 API 的所有特性（设备上执行、批量操作、转换等），我们提供了一个 TorchRL-to-gym API，允许用户在 gym 或 gymnasium 中注册他们想要的任何环境。这反过来又可以使 TorchRL 成为一个通用的 lib-to-gym 转换器，适用于无状态（例如 dm_control）和无状态（Brax, Jumanji）环境。该特性在文档中有详细介绍。info_dict 读取 API 也得到了改进。
环境加速
我们添加了在 ParallelEnv 中，在与用于传递数据的环境不同的环境上执行环境的选项。我们还对 GymLikeEnv 类进行了加速，使其现在可以与 gym 本身相媲美。
扩展目标函数
RLHF 和大规模训练中最流行的目标函数（PPO 和 A2C）现在已与 FSDP 和 DDP 模型兼容！
TensorDict
特性：MemoryMappedTensor 替代 MemmapTensor
我们为 TensorDict 提供了一个更高效的 mmap 后端：MemoryMappedTensor，它直接继承自 torch.Tensor。它附带了许多用于构建的工具，例如from_tensor、empty等等。MemoryMappedTensor 现在比其对应的类更安全、更快。该库与先前的类完全兼容，以便于过渡。
我们还引入了一套新的多线程序列化方法，使得 tensordict 序列化具有很强的竞争力，LLM 的序列化和反序列化速度比使用 torch.save 快 3 倍以上。
特性：TensorDict 中的非张量数据
现在可以通过NonTensorData tensorclass 携带非张量数据。这使得构建带有元数据的 tensordict 成为可能。memmap-API 完全兼容这些值，用户可以无缝地序列化和反序列化此类对象。要在 tensordict 中存储非张量数据，只需使用__setitem__方法进行赋值即可。
效率提升
unbind、split、map 甚至 TensorDict 实例化等多种方法的运行时性能都得到了改进。请查看我们的基准测试！
TorchRec/fbgemm_gpu
VBE
TorchRec 现在在EmbeddingBagCollection模块中原生支持 VBE（变长批量嵌入）。这允许每个特征具有可变的批量大小，从而解锁稀疏输入数据去重，可以极大地加快嵌入查找和 all-to-all 的时间。要启用此功能，只需使用stride_per_key_per_rank和inverse_indices字段初始化KeyedJaggedTensor，它们分别指定每个特征的批量大小和用于重新索引嵌入输出的逆索引。
除了 TorchRec 库的更改外，fbgemm_gpu 还添加了对 TBE 中每个特征可变批量大小的支持。VBE 已在加权和未加权情况下的分割 TBE 训练中启用。要使用 VBE，请确保使用最新的 fbgemm_gpu 版本。
嵌入卸载
这项技术指的是使用 CUDA UVM 缓存“热点”嵌入（即将嵌入表存储在主机内存中，并在 HBM 内存中缓存），以及预取缓存。嵌入卸载允许使用更少的 GPU 运行更大的模型，同时保持有竞争力的性能。使用预取 pipeline (PrefetchTrainPipelineSparseDist)，并通过 planner 中的约束条件传入每个表的缓存加载因子和prefetch_pipeline标志即可使用此功能。
Fbgemm_gpu 在v0.5.0中引入了UVM 缓存 pipeline 预取，以提高 TBE 性能。这使得缓存插入可以与 TBE 前向/后向计算并行执行。要启用此功能，请确保使用最新的 fbgemm_gpu 版本。
Trec.shard/shard_modules
这些 API 将嵌入子模块替换为其分片变体。shard API 应用于单个嵌入模块，而 shard_modules API 替换所有嵌入模块，并且不会触及其他非嵌入子模块。
嵌入分片的行为与之前的 TorchRec DistributedModuleParallel 行为类似，不同之处在于 ShardedModules 现在是可组合的，这意味着模块由TableBatchedEmbeddingSlices支持，这些 TableBatchedEmbeddingSlices 是对底层 TBE 的视图（包括 .grad）。这意味着现在使用 named_parameters() 会返回融合参数，包括在 DistributedModuleParallel 中。
TorchVision
V2 转换现在已稳定！
torchvision.transforms.v2 命名空间此前一直处于 BETA 阶段。现在已稳定！无论您是 Torchvision 转换的新手，还是经验丰富，我们都鼓励您从V2 转换入门开始，以了解新 V2 转换的功能。
浏览我们的主文档，获取一般信息和性能提示。可在API 参考中找到可用的转换和函数。在我们的示例库中也可以找到更多信息和教程，例如Transforms v2：端到端目标检测/分割示例或如何编写自己的 V2 转换。
朝着支持torch.compile()迈进
我们正在逐步为 torchvision 接口添加对torch.compile()的支持，减少图中断并允许动态形状。
torchvision 操作（nms、[ps_]roi_align、[ps_]roi_pool和deform_conv_2d）现在已与torch.compile和动态形状兼容。
在转换方面，大多数低级内核（如resize_image()或crop_image()）应该能够正确编译，没有图中断并支持动态形状。我们仍在解决剩余的边缘情况，朝着完全的功能支持和类迈进，预计下一版本在此方面会有更多进展。
