跳转至

DTorch 的优势与机遇

DTorch 是基于 Single-Controller 和 DTensor 构建的 PyTorch 分布式 API

前两篇博客分别介绍了 DTorch 的整体情况架构设计。本文在此基础上,系统讨论 DTorch 相对于现有分布式框架的优势与劣势,以及其面临的行业机会与展望。

1 DTensor 的优势

分布式的本质复杂度在于:一个 Tensor 被切分到多个设备后,"哪个设备持有哪一块数据"这一信息必须被记录在某处。如果它不被记录在程序里(PyTorch 的普通 Tensor),就必须由用户代码显式表达——手动计算局部 shape、按 rank 分支、在正确的时机调用 all-gather / all-reduce。这类代码与算法逻辑纠缠在一起,冗长、易错,且难以修改。

DTensor 把分布信息记录在 Tensor 自身:DeviceMesh 描述"分布在哪些设备上",Placements 描述"如何切分"。算子根据这两个属性自动完成切分与通信,分布式代码因此变得和 single-device program 一样简洁易用:Tensor 始终以完整的逻辑形态出现在代码中,"分布在哪些设备上、如何切分"只是它的属性。DTorch 则原生支持 DTensor,所有算子直接接受 DTensor 输入,并自动推导输出的 DeviceMesh 与 Placements。由此带来一连串的连锁简化:

  • 加载权重自动切分:加载一份完整的 state_dict,框架按 Placements 自动切分到各 rank,无需像 Megatron-LM 那样按 TP / PP rank 预先切分和转换权重;
  • 取值自动聚合to_torch() 时自动把各 rank 分片聚合回完整 Tensor,无需手动 all-gather;
  • 不均匀切分自动处理:维度长度不能被设备数整除时自动计算本地 shape,通信时自动 padding / unpadding,全程对用户透明。

PyTorch 中虽然也实现了 DTensor,但至今仍处于 alpha state and under development

没有 DTensor 的框架如何解决这一问题。Megatron-LM 和 vLLM 都给出了自己的答案——提供一套专用并行层,把通信写死在层的实现里:

框架 做法 代价
PyTorch 不解决,用户手动切分与通信 分布式代码冗长易错
Megatron-LM 提供 ColumnParallelLinear / RowParallelLinear 等专用并行层,通信写死在 forward 中 模型必须用 Megatron 算子重写;权重需按 TP / PP rank 手动切分,训练与推理之间需要 checkpoint conversion
vLLM 同样维护一套并行层(QKVParallelLinear 等),并为分布式推理重写了几乎所有主流模型 每个模型在 HuggingFace 实现之外再维护一份并行实现,新模型接入成本翻倍
DTorch 分布信息是 Tensor 的属性,算子原生支持 模型代码与单卡一致,仅需配置 DeviceMesh 与 Placements

2 Single-Controller 的优势

有了 DTensor,就不再需要 Multi-Controller + SPMD。 Multi-Controller + SPMD 并非一个主动选择的理想设计,而是 DTensor 缺位时期的产物:既然每个进程只能看到局部数据,那就让每个进程各自管理自己,用 rank 和集合通信相互协调。当 DTensor 把全局视角的数据描述带回来之后,用户程序天然就是"全局的",此时更自然的做法是让唯一的 Controller 直接消费这份全局信息——数据切分、任务分发、通讯协调、资源管理全部由框架完成,这正是 Single-Controller(正如设计决策中所说:Multi-Controller 方案以普通 Tensor 为设计前提,切换到 DTensor 后,Single-Controller 的易用性更好)。

落到使用体验上,Single-Controller 把整个分布式集群抽象成一个线程:用户只接触一个普通的 Python 线程——没有 torchrun、没有多进程、没有 rank,以全局视角描述计算即可。Multi-Controller 则把多进程模型直接暴露给用户:进程的启动与生命周期、每个进程各自的状态、以 rank 区分的执行路径,都需要用户以进程为单位理解和管理。这是 Single-Controller 易用性最直观的来源。

在此之上,框架接手了原先由用户承担的协调工作,优势体现在两个层面。

自动通讯与调度。Controller 持有全局计算图:redistribute() 是普通算子,框架可以选择最高效的实现、插入必要的同步节点(CUDA Event)、实现计算与通信的重叠;通信死锁可以在构图时被检测和规避。用户不需要创建和管理 ProcessGroup,不需要关心 all-reduce 应该插在哪里。

全局视野带来的系统能力——这些能力在 Multi-Controller 的 SPMD 范式下需要大量手动协调,或根本无法实现:

  • 全局设备管理与设备虚拟化:DeviceMesh 仅仅是一组 Device ID 的描述,物理设备由框架统一管理,因此可以把多个虚拟设备映射到同一物理设备上——单卡模拟分布式正是设备虚拟化的一个具体体现(环境变量 DTORCH_DTENSOR_IN_SAME_DEVICE=1,默认模拟 8 卡,在单卡上完成分布式程序的开发与调试);
  • 运行时动态增删计算节点:Controller 统一管理所有 Worker,可以在程序运行过程中向集群加入新的计算节点、或将其移除(配合上述任务迁移),实现集群的弹性伸缩——SPMD 的进程组在 torchrun 启动时便已固定,运行期扩缩容需要重组进程组并重启;
  • 故障排查:Controller 掌握全局状态,"哪个设备、哪个算子出了问题"可以被集中观测和定位,而不是散落在每个进程各自的日志里;
  • 错误设备驱逐与自动恢复:Controller 知道每个任务在哪些设备上执行,可以把故障设备上的任务迁移到健康设备继续执行,而不是整个作业失败重启;
  • Auto Parallel(自动并行策略搜索)与 JIT 编译(依靠全局信息进行编译优化)。

当前进度:DTorch 已实现基于 gRPC 双向心跳的进程故障检测与优雅关闭(秒级发现,避免资源泄漏与死锁),为上述能力打下基础;自动恢复等能力是 Single-Controller 架构上自然的扩展方向。

2.1 跨设备的计算编排:RL 的训练与推理

Single-Controller 还解决了一类 SPMD 难以处理的问题——把不同角色放在不同设备上并编排它们的执行。强化学习训练(PPO / GRPO)是最典型的场景:actor 生成(推理)、reference model 推理、reward model 推理、参数更新(训练)各自占用一组 GPU,交替执行且存在复杂的依赖关系。Multi-Controller 生态中,vLLM 和 veRL 为此引入 Ray 做中心调度,构成 Single-Controller + Multi-Controller 的混合范式,还需要专门解决训练引擎与推理引擎之间的权重同步问题。

DTorch 原生就是 Single-Controller:将训练和推理分别放置在不同的 GPU 上,直接使用不同的 DeviceMesh 即可——这只是配置层面的差异,不涉及代码逻辑的修改(划分独立的推理组与训练组,或让两者共用同一组 GPU,程序其余部分完全一致)。所有角色由同一个 Controller 编排;角色之间的数据传输(激活、权重)是普通的 DTensor 操作(redistribute),不需要跨引擎的格式转换与同步协议:

# 示意:RL 训练中不同角色的编排(训练能力在路线图中)
mesh = init_device_mesh("cuda", (4, 2), mesh_dim_names=["dp", "tp"])

policy  = Policy(config, device_mesh=mesh)    # actor / rollout
ref     = Policy(config, device_mesh=mesh)    # reference model(冻结)
trainer = Trainer(policy, device_mesh=mesh)   # 参数更新

for batch in data:
    responses = policy.generate(batch.input_ids)   # 推理
    rewards   = reward_model(responses)            # 推理
    trainer.update(policy, responses, rewards)     # 训练:权重原地更新,
                                                   # 无跨引擎同步

3 接入新的算法更方便

分布式生态中,新算法的落地成本极高——每个新算法都要在多个框架中各适配一遍。DTorch 把接入成本压回到"接近单卡程序"的水平:

新算法类型 传统生态的做法 DTorch 中的做法
新模型结构 分别适配 HF 单卡版、Megatron 训练版、vLLM 推理版 按单卡写法实现一次(或从 HF 迁移),配置 DeviceMesh / Placements 即可分布式运行
新并行策略 在 SPMD 代码中重写通信逻辑 表达为新的 Placements 组合或并行 Module,实现一次、所有模型受益
新优化技术(量化、缓存、融合) 各框架分别实现 在算子 / 框架层接入,通过 ExecuteConfig 开关,模型代码不用改

以 DTorch 已落地的扩散模型推理为例:

  • diffusers==0.34.0 迁移 SD3 / FLUX 时,仅做必要修改(替换模块基类与 import、包装 tokenizer 输出等),目录结构与原库一致,模型代码与单卡版本几乎相同;
  • Sage Attention 量化在 SDPA 算子内接入,通过 QuantizeConfig 一行开关,所有模型受益;
  • First Block Cache 残差缓存、RoPE / SiLU-Linear / LayerNorm 算子融合,均通过 ExecuteConfig 开启,不侵入模型代码;

更根本的变化是协作方式的改变。当前实践中,分布式代码通常由经验丰富的工程人员魔改,算法人员使用魔改后的结果进行训练和推理——这堵墙已经严重限制了算法创新。DTorch 把分布式代码的形态拉回到与单卡一致,算法人员可以直接阅读、修改和调试分布式程序,新想法从提出到在集群上验证的路径大幅缩短。对算法迭代频繁的研究型团队而言,这意味着更高的实验吞吐。

4 调试与可观测性

调试是分布式开发中被低估的大头。SPMD 生态中,多进程作业的调试体验很差:日志按 rank 混排,堆栈分散在每个进程里;断点需要 attach 到特定进程,而集合通信会因一个进程被断下、其余进程等待而整体 hang;NCCL 通信死锁时,所有进程停在等待中,难以拿到有效的堆栈。定位"哪个 rank、哪次通信、哪个 shape 不匹配"的问题,往往要靠反复二分和重跑。

从用户视角看,DTorch 的分布式程序就是一个单线程、全局视角的普通 Python 程序,调试体验回到单卡时代:

  • 直接调试print 直接输出 DTensor 的全局 shape 与 placements,断点和单步在 Client 线程上即可跟踪整个分布式程序的执行;
  • 单卡复现:结合单卡模拟分布式(见第 2 节),多卡环境的问题可以在单卡上稳定复现、反复调试;
  • 集中观测:Controller 掌握全局状态,日志、心跳与错误信息集中在一处,而不是散落在每个进程;

由此,分布式程序的调试成本从"经验丰富的工程师专职排查"回到"算法人员自己上手"的水平——这也是第 3 节协作方式改变的另一面。

5 实现单机和分布式计算统一的框架

当前 PyTorch 生态中,"单机代码"和"分布式代码"是两套独立维护的实现:单卡训练和推理用 HuggingFace transformers / diffusers;分布式训练用 Megatron-LM / DeepSpeed 的改写版;分布式推理用 vLLM / SGLang 的重写版。同一个模型至少存在两到三份实现,新模型发布后,社区往往需要数周时间才能完成各框架的分布式适配。

DTorch 中,单机与分布式是同一份代码在不同 DeviceMesh 上的执行

  • DeviceMesh 只有一个设备时,DTorch 程序就是普通单卡程序,性能与 PyTorch 基本一致(SD3 单卡推理端到端耗时 -2.08%,见介绍博客);
  • 同一份 Llama 实现支持 DP、TP、PP、CP 的任意组合,仅通过 DeviceMesh 与 Placements 配置切换(见 Llama 并行示例);
  • 从单机多卡扩展到多机,也只是 DeviceMesh 的不同配置;
  • 结合第 2 节的单卡模拟分布式能力,可以在一块 GPU 上开发、调试任意并行组合的分布式程序(如模拟 16 个虚拟 GPU 跑 DP × TP × PP × CP),确认无误后再部署到真实集群。
从单卡扩展到多卡 PyTorch 生态 DTorch
代码 换一套框架实现(Megatron / vLLM …) 同一份代码
改动点 重写模型、切分权重、管理通信 修改 DeviceMesh 与 Placements
调试 在多进程环境调试 可先单卡调通,再扩展;单卡上可模拟分布式
部署 提前确定并行方案 按可用资源选择并行度,运行时切换

单机与分布式统一的意义:开发与调试成本(先单卡调通算法逻辑,再切换 DeviceMesh 扩展,无需重写)、部署灵活性(同一份代码按资源规模选择 DP+TP 或叠加 PP+CP)、维护成本(一份实现,修复与新特性同时惠及单机和分布式)。

6 实现训练和推理统一的框架

训练和推理使用相同的算子与模型结构,却在工程上分裂为两个生态:训练用 Megatron-LM / DeepSpeed,推理用 vLLM / SGLang。这一分裂带来三个代价:

  1. 多份模型实现:同一模型在训练框架和推理框架中各有一份并行实现,功能与修复需要双向同步;
  2. 权重格式转换:训练产出与推理引擎的权重格式(切分方式、参数命名)不同,需要 checkpoint conversion 工具链,维护成本高且容易出错;
  3. 训推数值不一致:两侧 kernel 实现与并行切分方式不同,同一份权重在两个引擎中的输出存在数值差异。

训推不一致在强化学习中代价最大。RLHF / GRPO 训练中,rollout 由推理引擎执行,策略更新由训练引擎执行,两侧的数值差异(train-inference mismatch)会引入额外的算法噪声、影响训练稳定性,社区为此投入了大量工程做训推一致性对齐;生成与训练之间的权重同步(如 veRL 的 3D-HybridEngine)也需要专门的通信与重切分机制。

DTorch 在架构上不存在训练与推理的鸿沟:两者是同一执行引擎的两种使用方式——同一套算子(LibTorch 后端)、同一份模型代码、同一套 DeviceMesh 与 Placements 并行描述。RL 流程中 rollout 与 update 在同一个框架内编排(见第 2 节):权重是普通的 DTensor,在角色间的传输就是 redistribute 操作,没有格式转换,也没有跨引擎同步协议;训推一致性由构造保证——两侧调用的是同一个算子的同一份实现,而非两套分别对齐过的实现。

目前 DTorch 已基于此架构完成扩散模型(SD3 / FLUX)分布式推理的原型,训练能力在路线图中。训练与推理由同一框架承载,是 DTorch 相对"训练框架 + 推理引擎"组合的长期结构性优势。

7 基于 LibTorch,开发成本可控

一套覆盖主流模型的算子库需要多年的工程积累。DTorch 采用 PyTorch 的 C++ 算子库 LibTorch 作为计算后端,大多数 Kernel 直接调用 LibTorch API 完成计算,新增算子只需实现数据结构与调度的一层"外壳":

  • 算子无需重写:绝大多数 Kernel 直接调用 LibTorch API 完成计算,只需几人的小团队,即可完成功能的开发和迭代;
  • 精度天然对齐:DTorch 与 PyTorch 调用同一算子库、运行同样的 CUDA kernel,在输入和计算逻辑一致时,输出可以做到逐 bit 一致,省去了漫长的精度对齐工作;
  • 生态直接复用:算子覆盖面随 LibTorch 版本升级自动扩展,HuggingFace 模型代码只需少量修改即可迁移;
  • 社区算子库直接复用:社区有大量基于 PyTorch 自定义算子实现的高性能算子库(如 SageAttention、FlashAttention)。DTorch 的 Python Kernel 支持在 C++ Kernel 中调用 Python 代码,这些库无需移植 kernel 即可直接复用。

8 按使用量付费的 GPU 云服务

当前 GPU 云服务的商业模式有两种:按 token(请求)付费按硬件使用时长付费

  • 按 token (请求)付费:平台托管模型并提供推理服务(如各家的大模型 API),用户按请求或生成的 token 数计费。该模式下只能运行固定且被广泛使用的模型;使用方式虽然固定,但平台可以围绕它做精心优化,GPU 利用率极高。
  • 按硬件使用时长付费:平台提供 GPU 机器,用户在 docker 容器中直接运行 PyTorch 计算,按时长计费。用户可以任意使用,但 GPU 利用率取决于用户自己的使用方式,通常最低。
按 token 付费 按硬件时长付费
可运行的模型 平台托管的固定主流模型 任意模型、任意计算
优化责任 平台精心优化 用户自行负责
GPU 利用率 极高 取决于用户,通常最低

两种模式各占一端:前者牺牲灵活性换取利用率,后者牺牲利用率换取灵活性。当用户想运行自己的模型和算法、同时又希望获得高利用率时,两个选项都无法满足。

DTorch 的架构为更细粒度的计费模式提供了可能。Client、Controller 与 Worker 之间通过异步消息解耦,三者的职责恰好构成用户与云平台之间的分工界面:Client 负责构建计算节点,与业务逻辑紧密相关,运行在用户侧;Worker 只执行计算,不持有业务逻辑,是无状态的执行单元;Controller 持有全局计算图,知道计算需要什么资源、何时可以释放

Worker 支持动态增删(见第 2 节"运行时动态增删计算节点"):需要 CPU、GPU 计算资源时,动态向云平台请求;完成计算后即释放。在这种模式下,用户负责业务逻辑,云平台负责调度存储、计算和网络,将硬件资源按需、按时间片动态地卖给用户。同时,由于调度权集中在平台手中,云平台可以对存储、计算和网络进行深度优化,持续提高资源利用率——用户获得接近"按硬件租用"的自由度,平台达到接近"按 token 托管"的利用率。

这也是 Pathways 所设想的"集群操作系统"方向。DTorch 当前的资源管理仍是按作业静态分配的原型,这条路径是 Single-Controller 架构打开的长期想象空间。

9 行业的机会

目前流行的深度学习分布式训练和推理框架基本都基于 PyTorch 开发,看似整个生态已经收敛到 PyTorch,实则还存在一些机会:

  1. PyTorch 在分布式上的易用性不足,对 DTensor 的支持不够完善。PyTorch 的 DTensor 至今仍处于 alpha state and under development,也未被主流的大模型训练和推理框架广泛使用。原生支持 DTensor、易用的分布式 API 层存在空位。
  2. PyTorch 的定位是框架而非解决方案。vLLM、SGLang、veRL 等解决方案依托 PyTorch 生长并各自取得了成功,说明价值捕获发生在"解决方案"层。依托于 PyTorch,可以衍生解决各个领域问题的解决方案——DTorch 可以占据的正是"分布式解决方案的基座"这一层。
  3. PyTorch 对国产 GPU 的支持较差。新硬件需要完整的分布式软件栈,而逐个适配 PyTorch 及其上层框架的周期很长。DTorch 的规模小、演进快,设备层与通信层(CUDA、NCCL)是独立的适配模块,可以为新硬件快速提供分布式能力,成为国产 GPU 分布式软件栈的选项之一。

为了和 PyTorch 形成错位竞争,DTorch 的定位是针对具体领域的分布式计算解决方案,如 Diffusion Transformer 模型的推理(SD3 / FLUX 已落地)、量化训练及推理(Sage Attention)、强化学习训练等,而非做完善的基础框架。实现解决方案并将代码开源,可以凭此证明并推广 DTorch 分布式 API,吸引其他开发者参与到 DTorch 的开发中。DTorch 把重心放在具体问题的解决方案上,可以小步快跑,逐步增加人力投入,并产生相应的业务价值。

10 劣势

任何技术路线都有代价。DTorch 的劣势中,前两项源于架构本身,后两项源于项目阶段:

  • 中心化调度的单点与吞吐上限(结构性):整个集群依赖同一个 Controller,它一旦崩溃,作业随之失败;计算节点的构建与分发也都经过同一个 Python 线程,在数千卡的规模上,中心化调度的消息吞吐可能先于 GPU 饱和——这也是 SPMD 至今仍主导超大规模训练的原因之一。DTorch 已实现 Worker 进程的故障检测与优雅关闭;Controller 容错与控制器分片(Pathways 的方向)是长期课题。
  • 数据依赖控制流的往返延迟(结构性):if loss < threshold: 这类按数据分支必须把值从 Worker 传回 Client,构成一次跨进程往返与同步点;TensorFuture 异步取值可以与计算重叠,但无法根除。控制流越密集,代价越明显。
  • 算子覆盖与生态成熟度(阶段性):算子需逐个接入,覆盖面远小于 PyTorch;CUDA Graph、torch.compile 等优化需框架层逐一适配;复杂度集中在 C++ 引擎中,贡献门槛高。
  • 训练能力尚未完成(阶段性):已落地的只有扩散模型(SD3 / FLUX)推理,autograd、优化器等训练能力仍在路线图中——第 6 节的训推统一目前是架构潜力而非现实能力。

这些劣势划出了 DTorch 当前的边界:短期不以数千卡超大集群为目标,立足中等规模、具体领域的解决方案(如扩散模型推理),再随生态成熟逐步扩展。

11 当前进展和实现路径

当前,DTorch 已跑通单机多卡下扩散模型分布式推理的完整链路,多机集群的通讯、调度相关的设施也接近完成:

  • DTensor 与算子体系:DeviceMesh、Placements(Shard / Replicate / Partial)与输出自动推导已就绪,权重自动切分、取值自动聚合、不均匀切分对用户透明;
  • 并行策略:DP、TP、PP、Ulysses / Ring CP 已落地,同一份 Llama 实现支持任意组合(见 Llama 并行示例);
  • 应用与优化:SD3 / FLUX(diffusers 0.34.0)分布式推理已验证,Sage Attention 量化、First Block Cache、算子融合经 ExecuteConfig 开启,单卡端到端性能与 PyTorch 基本一致(SD3 耗时 -2.08%);
  • 系统与质量:单机多线程 / 集群多进程(ZMQ + gRPC 心跳)两种执行模式,单卡模拟分布式与 TensorFuture 异步取值;三层测试体系以 PyTorch 为基准,输出可做到逐 bit 一致。

后续按"先推理、后训练,先能力、后规模"推进。短期内 DTorch 被定位为扩散模型的分布式推理框架;中期补齐训练能力后,可以成为 Transformers 模型预训练和后训练的训练框架;远期可以往"集群操作系统"的方向演进。

阶段 目标 关键内容
近期 完善推理 扩大算子与模型覆盖,集成社区优秀的融合算子实现,推理性能对齐框里框架的水平
中期 补齐训练 autograd 与优化器、量化训练,落地训推统一,进而编排 RL 训练(PPO / GRPO)
远期 扩展规模 Controller 容错与控制器分片、故障设备自动恢复、动态扩缩容;探索 Auto Parallel、JIT 编译与按使用量付费的云服务

12 展望

易用性对于深度学习框架至关重要。回顾深度学习框架的发展历程,与其他同期框架相比,PyTorch 在性能上往往处于劣势地位,但其凭借易用性的优势吸引了大量开发者,渐渐成为事实上的行业标准。单卡时代的故事如此;分布式时代,易用性的高地尚无人占据。DTorch 基于 Single-Controller + DTensor 的技术路线,在分布式场景上的易用性优于 PyTorch——以"分布式程序的写法与单卡程序一致"为目标,把算子实现建立在 LibTorch 之上,用可控的开发成本,以此为突破点重构当前 Transformer 模型训练和推理流程,并支持未来出现的新的训练范式。

13 总结

本文从技术优势、商业模式、行业机会、劣势与进展路径五个角度,讨论了 DTorch 的差异化定位:

# 优势 核心要点
1 DTensor 分布信息是 Tensor 的属性,算子原生支持,分布式代码与单卡一致
2 Single-Controller 全局视角带来自动通讯调度、故障处理、动态伸缩与 RL 编排能力
3 新算法接入 接入成本回到"接近单卡程序",优化技术在算子层落地、全模型受益
4 调试与可观测性 分布式程序可像单卡程序一样调试、复现与观测
5 单机与分布式统一 同一份代码在不同 DeviceMesh 上执行
6 训练与推理统一 同一执行引擎的两种用法,训推一致由构造保证
7 开发成本可控 复用 LibTorch 与社区生态,几人小团队即可开发迭代
8 按使用量付费的云服务 Client / Controller / Worker 构成用户与平台的分工界面,自由度与利用率兼得
9 行业机会 分布式易用性空位、解决方案层定位、国产 GPU 适配
11 进展与路径 扩散模型推理链路已落地,沿"先推理后训练、先能力后规模"推进
12 展望 负载复杂化与 SPMD 的错位持续放大,易用性高地尚无人占据

这些优势同源于一个架构选择:Single-Controller + DTensor。它让分布式程序回到单卡程序的形态(1–6),把开发成本压缩到小团队可承载的范围(7),自然划分出用户与云平台的分工界面(8)——在 PyTorch 生态收敛的表象之下,这些正是 DTorch 的机会所在(9、12)。与此同时,第 10 节讨论的劣势同样是这一架构选择的代价:中心化调度的吞吐上限需要随规模演进正面应对,生态与训练能力的差距则决定了 DTorch 短期内以具体领域的解决方案为立足点,逐步向外扩展——第 11 节的进展与路径,正是这一扩展的具体安排。

DTorch 的代码和文档均在 GitHub 开源,欢迎关注与参与。

延伸阅读