Pastens

veScale 深度技术分析

11分钟阅读

veScale 深度技术分析

一、概览

属性内容
项目veScale — ByteDance Seed 的开源 PyTorch 分布式训练库
GitHubhttps://github.com/volcengine/veScale ⭐ ~0.8K
论文 1veScale: Consistent and Efficient Tensor Programming with Eager-Mode SPMD (arXiv Sep 2025)
论文 2veScale-FSDP: Flexible and High-Performance FSDP at Scale (arXiv Feb 2026)
机构ByteDance Seed (字节跳动 Seed 团队)
LicenseApache 2.0
核心创新RaggedShard 非对称分片格式 + Eager-mode SPMD + 零拷贝 FSDP 通信

项目演化

veScale 经历了两次重大迭代：

  1. Legacy veScale (2023–2025) — 基于 PyTorch DTensor 的 N-Dim 并行训练框架，支持 TP/SP/DP/ZeRO/PP 全栈并行，已归档到 legacy/ 目录
  2. New veScale (2025–2026) — 精简为 Pure DTensor Extension，核心贡献是 RaggedShard Placement，代码仅 ~2K 行 Python

二、论文解读

Paper 1: veScale — Eager-Mode SPMD (2509.07003)

核心问题

PyTorch DTensor 存在两个关键缺陷：

  • 语义不一致: 分布式执行无法保证与单设备执行 bit-wise 一致 (weight init、dropout 随机数在不同并行策略下产生不同结果)
  • 性能低下: DTensor dispatch 引入高 CPU 开销 (每个算子 0.5ms+)，通信粒度过细

核心贡献

  1. 一致性分布式 RNG 算法 — 首次实现单设备语义等价 (single-device-semantic) 的分布式随机数生成。通过 Virtual Thread + Virtual Offset 机制，确保无论分片方式如何，合并后的全局张量与单设备生成结果完全一致。

  2. 零模型代码变更的 Plan API — 通过 VescalePlan + parallelize() API，将并行策略与模型定义完全解耦。支持正则表达式匹配多参数、分阶段 sharding（如 ZeRO-3 在 INIT/RUN 阶段不同 placement）。

  3. DTensor 性能优化 — 三大优化：

    • Static Eager 执行模式：预编译 dispatch 逻辑，消除运行时开销
    • 通信融合：跨多个并行维度和多个 DTensor 融合为单次 collective
    • 算子级优化：减少 _cvt_dtensor 路径上的不必要转换

评估结果

  • 比 TorchTitan 快 2.2×
  • 代码复杂度降低 78.4%（LoC 对比 Megatron-LM）
  • 在 LLaMA-3-70B、Mixtral 8×7B、DiT 上验证

Paper 2: veScale-FSDP — Flexible FSDP (2602.22437)

核心问题

现有 FSDP 系统（DeepSpeed ZeRO、FSDP1、FSDP2、Megatron-FSDP）都存在共同局限：

  • 固定分片粒度：element-wise 或 row-wise，无法支持 block-wise 量化（DeepSeek-V3 风格）和矩阵优化器（Shampoo、Muon）
  • 通信/内存开销：interleaved copy (FSDP2)、padding inflation (Megatron)、fragmented collectives (DeepSpeed)

核心贡献: RaggedShard

提出 RaggedShard — 一种通用的非对称张量分片格式，作为 PyTorch DTensor 的第四种 Placement：

class RaggedShard(Placement):
    dims: tuple[int, ...]        # 分片的维度前缀
    local_units: tuple[int, ...] # 各设备上的相对分配比例

表达能力：通过调整 dims 和 local_units，RaggedShard 可以表达：

  • Element-wise Shard (极端粒度)
  • Symmetric Row-wise Shard (FSDP2 的 Shard(0))
  • Asymmetric Row-wise Shard (不同设备持有不同行数)
  • Block-wise Shard (自定义 block size，完美对齐量化 block)
  • 任意非均匀分布

组成性：与标准 DTensor placements (Shard, Replicate, Partial) 完全可组合，通过 StridedRaggedShard 处理 2D 并行 (FSDP×EP/TP) 中 placement 顺序与逻辑顺序的差异。

性能优化: Planning Algorithm + DBuffer

  1. NP-hard 规划问题：将 RaggedShard 张量布局重排形式化为 NP-hard 优化问题，使用多项式时间启发式算法求解
  2. Distributed Buffer (DBuffer)：支持 RaggedShard 张量的全局缓冲区，实现：
    • 零拷贝通信 — AllGather/ReduceScatter 直接在 DBuffer 上操作，消除 interleaved copy
    • 批量内存分配 — 减少 fragmentation
    • 确定性内存管理 — 替代 PyTorch record_stream 非确定性释放

评估结果

指标veScale-FSDPvs DeepSpeedvs FSDP2vs Megatron-FSDP
吞吐量提升5–66%✅✅✅
内存降低16–30%✅✅✅
10K GPU 弱扩展Near-linear———
2.4T 参数训练1K GPU 可行OOMOOMOOM

两篇论文的关系

论文 1 (2509.07003) ──→ veScale 基础框架 (Plan API, RNG, DTensor 优化)
         │
         └──→ 论文 2 (2602.22437) ──→ veScale-FSDP (RaggedShard, Planning, DBuffer)
                     │
                     └──→ 开源代码 (RaggedShard Placement 实现)

论文 1 建立了 eager-mode SPMD 的基础设施，论文 2 在此基础上专攻 FSDP 场景，通过 RaggedShard 解决灵活性和性能问题。

三、代码库分析

架构总览

当前新 veScale 非常精简，核心在 vescale/dtensor/ 目录下，整体结构：

veScale/
├── vescale/
│   ├── dtensor/                    ← 核心代码
│   │   ├── _api.py                 # DTensor 类 (继承 torch DTensor)
│   │   ├── placement_types.py      # RaggedShard, _StridedRaggedShard 实现
│   │   ├── _redistribute.py        # 重分片逻辑 (RaggedShard 感知)
│   │   ├── _dispatch.py            # OpDispatcher (RaggedShard 特殊 handler)
│   │   ├── _sharding_prop.py       # sharding 传播 (继承 PyTorch)
│   │   ├── _collective_utils.py    # mesh_scatter_ragged (send/recv 实现)
│   │   ├── _dtensor_spec.py        # DTensorSpec + is_ragged_shard
│   │   ├── _ops/                   # 算子规则 (pointwise, matrix, tensor 等)
│   │   └── vescale_utils/          # 工具函数 (ragged_shard_utils, checkpoint)
│   └── utils/
│       └── monkey_patch.py         # 运行时补丁 (patch_method)
├── legacy/                         ← 旧版 veScale (已归档)
└── test/                           # 测试用例

核心组件详解

1. RaggedShard Placement (placement_types.py, ~270 行)

RaggedShard 是一个 frozen dataclass，继承自 PyTorch 的 Placement：

  • _split_tensor(): 基于 local_units 比例将 flatten 后的张量切分给各设备
  • _ragged_shard_tensor(): 分发张量到 mesh，使用 mesh_scatter_ragged (send/recv 实现)
  • _to_replicate_tensor(): 通过 all_gather 收集成全量张量
  • _to_new_ragged_shard(): 从一个 RaggedShard 重分片到另一个（不同 local_units），使用 all_to_all
  • reconstruct_tensor_from_flat(): 从 flat buffer 重建张量

关键设计决策：RaggedShard 操作的是 flattened 存储，即所有非分片维度被合并为一个 “tail”。dims 参数指定哪些前缀维度被分片。

2. Redistribute (_redistribute.py, ~200 行)

重写了 PyTorch 的 redistribute_local_tensor，处理 4 种场景的组合：

Current \ Target有 RaggedShard无 RaggedShard
有 RaggedShard_to_new_ragged_shard (all-to-all)_to_replicate_tensor (all-gather) → 标准 redistribute
无 RaggedShard标准 redistribute → _split_tensor调用 PyTorch 原版的 redistribute

3. OpDispatcher (_dispatch.py, ~380 行)

继承 PyTorch OpDispatcher，添加 RaggedShard 特殊处理：

  • ragged_norm_op_handler(): 处理 linalg_vector_norm，在分片边界不匹配时使用 ragged_norm_kernel（@torch.compile 编译的 kernel）
  • fused_adamw_sgd_op_handler(): 将 fused optimizer 的参数直接从 DTensor 拆包为 local tensor
  • found_inf_reduce_handler(): 处理 _amp_foreach_non_finite_check_and_unscale_

4. 集体通信 (_collective_utils.py)

mesh_scatter_ragged() 使用 P2P send/recv 替代标准 mesh_scatter，支持不同大小的分片。代码注释标注了性能问题：“serialized send/recv, can we do async launch?”。

5. Monkey Patch (utils/monkey_patch.py)

veScale 通过 patch_method 在运行时给 PyTorch 的类打补丁（如给 Placement 添加 is_ragged_shard() 方法）。这种方式比 fork PyTorch 更轻量，但可能引入兼容性问题。

数据流追踪

用户代码: DTensor.from_local(tensor, mesh, [RaggedShard(...)])
    │
    ├─ _FromTorchTensor.forward()
    │   └─ compute_global_tensor_info() 计算 global shape/stride
    │
    ├─ 算子执行 (如 matmul)
    │   ├─ OpDispatcher.dispatch()
    │   │   ├─ _cvt_dtensor() 将普通 Tensor 转为 Replicate DTensor
    │   │   ├─ sharding_propagator.propagate() 推算输出 placement
    │   │   ├─ redistribute_local_args() 重分片输入到一致 placement
    │   │   └─ 本地算子执行 → OpDispatcher.wrap()
    │   └─ 特殊 handler (ragged_norm / fused_adamw)
    │
    ├─ DTensor.redistribute() 重新分片
    │   └─ Redistribute.forward()
    │       └─ redistribute_local_tensor()
    │           ├─ RaggedShard→RaggedShard: all_to_all
    │           ├─ RaggedShard→Replicate: all_gather
    │           └─ Replicate→RaggedShard: redistribute → _split_tensor
    │
    └─ DTensor.to_local() / full_tensor() 获取局部/全局张量

Legacy 代码

legacy/ 目录包含完整的旧版 veScale（~30K+ 行代码），包括：

  • vescale/dmodule/ — Distributed Module (类似 FSDP)
  • vescale/dmp/ — Distributed Model Parallelism (auto TP/SP plan)
  • vescale/pipe/ — Pipeline Parallelism
  • vescale/checkpoint/ — 分布式 checkpoint (含 gRPC server)
  • vescale/devicemesh_api/ — N-Dim device mesh API
  • legacy/examples/ — Mixtral、LLaMA2、nanoGPT 的 4D 训练示例

四、亮点与局限

亮点

  1. 极简设计: 新 veScale 核心代码仅 ~2000 行 Python，通过继承/补丁 PyTorch DTensor 实现，可维护性极强
  2. RaggedShard 通用性: 统一了 element-wise、row-wise、block-wise 三种分片格式的表达，是 FSDP 领域的重大抽象创新
  3. 生产验证: 已在 ByteDance Seed 部署到 10K+ GPU 集群，训练 2.4T 参数模型
  4. 零代码变更: 用户无需修改模型代码即可获得 FSDP + block-wise 量化 + Muon 优化器支持
  5. 论文-代码一致: 开源代码实现了论文中 RaggedShard 的核心思想，且代码质量高、注释清晰

局限

  1. 仅支持单一 RaggedShard: 当前限制每个 DTensor 只能有一个 RaggedShard placement（见 get_ragged_shard() 中的断言 “only 1 ragged shard is allowed for now”）
  2. P2P 通信性能: mesh_scatter_ragged 使用串行 send/recv 而非 broadcast，作者已标注性能问题
  3. 缺少完整 FSDP 实现: 开源代码仅包含 RaggedShard placement 和 redistribute 逻辑，论文中提到的 Planning Algorithm 和 DBuffer 未开源（是 ByteDance 内部实现）
  4. 3D/4D 并行示例缺失: 新代码没有像 legacy 那样的端到端训练示例
  5. 旧代码被归档: legacy/ 代码不再维护，但其中包含许多有价值的并行实现（pipeline、auto-plan 等）

五、个人评价

veScale 是字节跳动 Seed 团队在分布式训练领域的一次重要技术输出。RaggedShard 的设计优雅地解决了 FSDP 多年来的灵活性瓶颈——将 element-wise 和 block-wise 分片统一在同一抽象下。这与笔者之前分析的 msOpProf 的 LD_PRELOAD 思路类似：都是通过一个统一的抽象层来简化原本复杂的系统问题。

不过，目前开源的只是 RaggedShard Placement 本身（约 2000 行），论文中关键的 Planning Algorithm 和 DBuffer 尚未开源。这使得开源版本更多是一个 概念验证（PoC） 而非生产级 FSDP 系统。期待 ByteDance 后续开源更多内部组件。

从代码质量看，veScale 使用了 patch_method 补丁机制来扩展现有 PyTorch 类，避免了 fork，这是值得学习的轻量扩展模式。

六、参考资料

关系图谱

反向链接