为什么 VLA 训练选 FSDP 而不是 Megatron:四个约束同时成立

本文是 从 LLM Infra 到具身智能 Infra:仿真集群、VLA 训练栈与数据管线全景解析 的姊妹篇——那篇给出了"VLA 训练用 FSDP 而非 Megatron"的结论,这篇把背后的逻辑彻底拆透。

核心原因一句话:VLA 模型的参数规模(1B-34B)刚好落在 FSDP 的"甜点区",加上它基于 HuggingFace 生态的异构多模态结构让 Megatron 类框架难以适配,而科研社区主导的开源生态天然偏向 PyTorch 原生方案。下面把逻辑拆开讲。

一、先看 VLA 训练的真实规模和场景

VLA 训练的参数量和你熟悉的 LLM 预训练完全不是一个量级:

场景 参数规模 典型硬件 训练方式
VLA 预训练 7B(OpenVLA)/ 3B(π0) 64 张 A100 × 15 天 全量
VLA 微调(主流场景) 7B 单节点 8 张 A100 LoRA 或全量
VLA-RL 后训练 7B 单节点到多节点 FSDP 管理
LLM 基座预训练 70B-万亿级 数千卡、跨多机 Megatron 3D 并行

这是最关键的判断依据:VLA 的主战场是"7B 量级 + 单节点到少量多机 + 高频微调迭代",不是"千亿级 + 数千卡"。这直接决定了框架选型。

二、为什么不是 Megatron:三个硬性障碍

障碍 1:VLA 是异构多模态模型,Megatron 的统一并行拓扑难以适配

Megatron-LM 官方已明确说明:它目前不原生支持"同一多模态模型内不同子模块使用不同并行配置"的场景——例如"视觉编码器 TP=1、LLM TP=8"这种分配在 Megatron 里开箱即用是做不到的。原因是 Megatron 的 TP/PP 通信组、流水线调度、参数同步、ckpt 加载都建立在全模型统一的全局并行拓扑之上,而 VLA 恰恰是"DINOv2/SigLIP 视觉编码器 + LLM 主干 + Action Expert"三段异构结构,各段的计算特征、显存压力、并行诉求都不同。

FSDP 对此天然宽容——它只是按 module 边界做参数分片,每个子模块可以用不同的 auto_wrap_policy 灵活处理。

障碍 2:VLA 用不上流水线并行

Megatron 的核心优势在 TP+PP+CP 组合下的千亿级扩展,流水线并行(PP)需要模型深、参数大才有意义。VLA 的 7B 模型在 8 卡单机上全量微调或 LoRA 都能跑下,用 PP 是杀鸡用牛刀——切两个 stage 带来的 bubble 反而拖慢吞吐。而 FSDP 在单节点 8 卡场景的开销几乎可以忽略,纯数据并行+参数分片已经够用。

障碍 3:Megatron 的代码侵入性太强

Megatron 要求模型按它的 GPTModel/MegatronModule 规范重写,层的切分逻辑(ColumnParallelLinear/RowParallelLinear)要显式嵌入模型代码。OpenVLA 是在 Prismatic VLMs 这个 HuggingFace 生态项目上改出来的,π0 也是标准的 PyTorch 模型结构——把这样的模型迁移到 Megatron 意味着重写主干,成本高且每次跟进上游开源模型更新都会冲突。

三、为什么是 FSDP:四个匹配点

匹配 1:参数规模刚好卡在 FSDP 的效率甜点区

FSDP 通过参数/梯度/优化器状态分片(FULL_SHARD,等价于 ZeRO-3)解决显存问题,代价是引入 all-gather 通信。在 7B-34B 这个规模、单节点到少量多机(InfiniBand/NVLink 内),通信开销可以被计算很好地掩盖,吞吐接近线性。OpenVLA 官方明确表述:

“Powered by PyTorch FSDP and Flash-Attention, we can quickly and efficiently train models from 1B - 34B parameters”

这条参数区间就是 FSDP 的舒适区。模型再往上(百亿到千亿),Megatron 的 TP+PP 才开始显著占优。

匹配 2:HuggingFace 生态的无缝集成

VLA 训练栈强依赖 HF 全家桶:transformers 加载 Llama/PaliGemma 底座、timm 加载视觉 encoder、peft 做 LoRA、datasets 处理数据。PyTorch FSDP 是 HF Accelerate 原生支持的分布式后端,配置 YAML 就能开启;peft 也有专门的 FSDP + LoRA 集成文档。这意味着你在 HF 上能跑通的单卡模型,几乎不改代码就能扩展到 8 卡 FSDP。Megatron 则需要脱离这套生态,自己处理 tokenizer、模型加载、数据格式——对快速迭代的科研代码是巨大负担。

匹配 3:微调场景的灵活性和易用性完胜

VLA 的高频操作是"拿 base ckpt 在自己机器人数据上微调",需要频繁切换 LoRA / 部分冻结 / 全量微调。FSDP 在这种场景下:

  • LoRA 微调:OpenVLA 用 LoRA 只训 1.4% 参数就能达到全量微调效果,FSDP+LoRA 在单卡到 8 卡上都能跑
  • 全量微调:OpenVLA 官方路径就是"native PyTorch FSDP using the Prismatic training script",推荐硬件是一个节点 8 张 A100
  • 代码侵入:训练脚本就是 torchrun 启动 + FSDP wrap,几十行搞定

相比之下 DeepSpeed 也能做 ZeRO-3,但 FSDP 的错误栈更可读、PyTorch 原生集成更深、不需要额外装包。HF 官方也指出 FSDP 在 bf16 精度下的行为比 DeepSpeed 更可预期。

匹配 4:科研开源社区的主流选择形成正循环

看看 VLA 领域代表性工作的训练框架选择:

项目 出处 训练框架
OpenVLA 伯克利 PyTorch FSDP + Flash-Attention
VLA-RL Online RL for VLA PyTorch FSDP 管理分布式训练
openpi Physical Intelligence PyTorch 生态,多 GPU 训练通过 FSDP 配置开启

这是典型的路径依赖:第一波开源 VLA 都在 HF 生态上构建,FSDP 是这个生态里最自然的分布式方案,后续工作为了复用数据加载、模型结构、LoRA 配方,继续沿用 FSDP。对要在这种生态里做 infra 的你来说,这是好消息——轮子都已经造好,你要做的是在上面加具身特有的能力(数据管线、评测、仿真调度)。

四、把这张选型逻辑说透

决策因素 Megatron DeepSpeed ZeRO-3 FSDP 对 VLA 的意义
参数规模适配 70B+ 优势明显 7B-100B 1B-34B 最优 ✅ VLA 主流在此区间
异构多模态支持 ❌ 统一拓扑 中 ✅ 按 module 灵活分片 ✅ VLA 三段式结构
HF 生态集成 ❌ 需脱离 中 ✅ Accelerate 原生 ✅ VLA 依赖 HF
单节点 8 卡微调 过度设计 可以 ✅ 最轻量 ✅ VLA 高频操作
LoRA 集成 复杂 可以 ✅ peft 原生 ✅ 微调主流
代码侵入度 高 中 低 ✅ 科研代码友好
错误可读性/调试 复杂 中 ✅ 更友好 ✅ 快速迭代

一句底层逻辑总结:FSDP 赢在"7B-34B 参数规模 + HuggingFace 生态 + 单节点微调为主 + 异构多模态结构"这四个 VLA 的关键约束同时成立——而 Megatron 的优势区间(千亿、数千卡、纯文本 LLM 预训练)在 VLA 场景里基本用不上。

1
2
3
4
参数规模轴
1B ────────── 7B(VLA主流) ── 34B ─────── 70B+ ──── 千亿
     ◄──────── FSDP 甜点区 ────────►  ◄── Megatron 占优 ──►
           (单节点~少多机,HF生态)         (数千卡,TP+PP+CP)

五、对 Megatron 背景的人意味着什么

需要主动迁移的心智:

  • 通信分析:Megatron 的 TP all-reduce / PP p2p 分析框架仍然适用,FSDP 的 all-gather + reduce-scatter 模式可以类比理解
  • ckpt 管理:FSDP 用 SHARDED_STATE_DICT,和你做过的全周期 ckpt 管理思路一致
  • 显存预算:FSDP 的参数分片量、激活 recomputation、通信-计算 overlap 这些调优手段你都熟

需要重建的认知:

  • 不再写 ColumnParallelLinear——FSDP 是在现有模型代码外做 wrap,模型本身保持纯 PyTorch 写法
  • 不再配 3D 并行——VLA 场景下通常就是 FULL_SHARD + 可选的 CPU offload,配置面窄得多
  • 不再关心 micro-batch 流水线调度——单节点数据并行,batch size 直接除以卡数

入职具身团队后:大概率你不需要从零搭训练框架(openpi/openvla 已经能用),你的价值在于把 FSDP 训练栈接入公司的数据管线、评测链路、ckpt 管理体系——这正是你 Megatron 经验最能发挥的地方。如果未来模型规模上到 30B+ 或者引入大规模 RL 后训练(VLA-RL 路线),可以再评估是否需要引入更重的并行策略,但那是后话。

六、附:模型格式正交——权重存储 vs 架构设计

聊完训练框架选型,顺带厘清一个常被混在一起的概念:“格式”。这个词其实混了两个维度——一个是权重怎么存的,一个是模型怎么设计的。Megatron 格式、HF 格式、具身模型格式这三个概念并不在同一层:前两者是权重存储/ckpt 组织方式的区别,而"具身模型"是架构设计理念的区别。厘清这一点,三者关系就一目了然。

6.1 三个概念到底在比什么

维度 Megatron 格式 HF 格式 具身模型
本质 分布式训练的 ckpt 存储规范 通用权重存储/加载规范 模型架构范式(VLM+Action)
组织单位 按 rank 切片的分布式 shard 单文件或分片的 safetensors 多模块组合(视觉+语言+动作)
核心诉求 大规模并行训练的高效保存/恢复 通用性、可移植、易加载 支持多模态输入+连续动作输出
典型使用者 Megatron-LM / NeMo / Megatron Bridge Transformers 生态全家桶 OpenVLA / openpi / LeRobot
典型文件形态 model_optim_rng.pt / __0_0.distcp 等 config.json + model.safetensors + tokenizer.json HF 外壳 + 多模块权重 + policy 配置
能否直接推理 ❌ 需转换或恢复并行环境 ✅ 直接 from_pretrained ✅ 但需配套前/后处理器

理解这层区分后,一个关键结论浮出来:同一个 VLA 模型,它可以以 HF 格式保存,也可以以 Megatron 格式保存——格式和架构是正交的。OpenVLA 的官方仓库同时提供 HF 格式 ckpt(用于推理/LoRA 微调)和 Megatron 格式 ckpt(用于大规模预训练/全量微调),就是这两个维度正交的直接体现。

6.2 Megatron 格式:为"分布式并行训练"而生

Megatron 格式不是一种文件类型,而是一种按并行度切分的目录结构:

1
2
3
4
5
6
7
8
9
10
checkpoint/
└── iter_0001000/
    ├── mp_rank_00_000/          # TP rank 0, PP rank 0
    │   ├── model_optim_rng.pt   # 模型+优化器+RNG 状态
    │   └── distrib_optim.pt     # 分片优化器状态
    ├── mp_rank_01_000/          # TP rank 1, PP rank 0
    ├── mp_rank_00_001/          # TP rank 0, PP rank 1
    ├── mp_rank_01_001/
    ├── ...
    └── latest_checkpointed_iteration.txt

关键特征:

  • 按 rank 切片存储:一个 8B 模型在 TP=2、PP=4 的配置下训练,会产生 8 份 mp_rank_* 文件,每份只包含该 rank 负责的参数切片。保存时不做聚合——每个 rank 直接写自己那份,从而避免了单卡 OOM。
  • 优化器状态也要存:训练恢复需要 optimizer momentum/variance,且这些也按 rank 切分。
  • 新版支持异步保存:Megatron v0.7.0 之后引入异步 ckpt,文件形态变成 __0_0.distcp(基于 PyTorch DistributedCheckpoint),保存时和训练计算重叠,几乎不阻塞。
  • 并行度信息嵌入文件名:换并行度(比如从 TP=2 改成 TP=4)需要专门的转换工具。

这套格式的价值在于"训练态高效"——每个 rank 独立读写、支持异步、支持换并行度恢复。代价是"推理态不方便"——单卡想加载它,必须先转换。社区已经把 HF↔Megatron 的双向转换工具链打通了:Megatron Bridge 提供 import_ckpt / export_ckpt 两个一键 API,支持在转换过程中保持并行度信息;NeMo-RL 也提供 convert_megatron_to_hf.py 脚本。但转换永远会存在精度/元信息风险(比如 embedding tied/untied 处理、RMSNorm 权重命名差异),每次转换都要做一遍推理对齐验证。

6.3 HF 格式:为"通用性"而生

HF Transformers 的模型目录是整个开源生态的事实标准,结构简单到近乎简陋:

1
2
3
4
5
6
7
8
9
openvla-7b/
├── config.json                     # 模型结构超参
├── tokenizer.json                  # 分词器
├── tokenizer_config.json
├── processor_config.json           # 视觉处理器(VLA 会加)
├── model.safetensors               # 完整权重(或分片)
├── model-00001-of-00002.safetensors
├── model-00002-of-00002.safetensors
└── generation_config.json

关键特征:

  • 单文件或少量分片:模型权重整体存成 safetensors(或分片),加载时一次性读进单进程,from_pretrained 一个调用搞定。
  • safetensors 格式:相对 pytorch_model.bin 的优势是零反序列化风险(不做 pickle)、加载更快、跨框架可读。
  • config.json 是语义描述:记录层数、hidden size、attention 头数、初始化方式,AutoModel 根据它实例化正确的类。
  • Tokenizer/Processor 独立保存:文本 tokenizer、视觉 processor 与权重解耦,VLA 里会额外放 processor_config。

HF 格式的价值定位:它是"交付格式"而不是"训练格式"。下载、加载、推理、LoRA 微调、发布到 Hub——全部以这个格式为通用货币。但如果你要做千亿模型的 Megatron 3D 并行训练,HF 格式反而是负担——单文件存不下分布式 shard,必须先转换到 Megatron 格式再启动训练。

6.4 具身模型格式:架构层面的"多模态 + Action"

到了具身模型,"格式"的含义才真正变化——它不仅是存储方式变了,模型本身的构成也和纯文本 LLM 不同。一个 VLA 模型至少包含三个语义不同的模块:

模块 典型选择 参数量级 训练时是否常冻结
视觉骨干 DINOv2 + SigLIP 双流(OpenVLA) ~500M 常 LoRA / 冻结
语言主干 Llama-2 / PaliGemma / Qwen 3B-7B 常做部分冻结
动作专家 离散 token 头 / flow matching 头 ~100M-1B 高频训练

以 π0 为例:它用一个大型 VLM backbone + 一个小得多的 action expert,action expert 专门处理机器人状态和动作,通过共享 attention 和 VLM 交换信息。这个结构叫 Mixture-of-Transformers(MoT),是当前 VLA 和世界动作模型(WAM)的默认架构。核心好处是:模态专属参数分离表征,但共享 attention 让视频和动作信息流通——既避免了"同一套权重同时拟合稠密视觉 token 和稀疏动作目标"的双重优化冲突,又保持了模态间的耦合能力。

这带来一个重要的格式后果:VLA 的 ckpt 天然是"多模块组合",不是单一 state_dict。具体体现在:

  • 视觉/语言/动作模块参数各自独立命名空间:加载时要分别处理,LoRA 可能只挂在语言主干和 action expert 上。
  • 必须配套前后处理器:视觉输入要归一化,动作输出要 un-normalize 到具体机器人的动作空间——OpenVLA 的 model card 明确说,“模型输出归一化的 7-DoF end-effector 增量,执行时需按 per-robot/per-dataset 的统计量 un-normalize”。这些统计量本身也是 ckpt 的一部分。
  • 多本体适配层:同一个模型要服务多种机器人本体,需要 per-embodiment 的动作空间映射。

LeRobot 的 policy 目录结构:具身模型格式的代表

LeRobot 作为 HF 出具的具身全家桶,为"policy"(即具身模型的部署形态)定义了一套更完整的目录规范,比标准 HF 目录多了策略层的东西:

1
2
3
4
5
6
7
8
pi0_checkpoint/
├── config.json                         # policy 配置(含机器人类型、action_dim)
├── model.safetensors                   # 策略权重
├── policy_preprocessor_config.json     # 输入归一化配置
├── policy_postprocessor_config.json    # 动作 un-normalize 配置
├── normalizer.safetensors              # 归一化统计量
├── unnormalizer.safetensors            # 反归一化统计量
└── train_config.json                   # 训练配置(可选)

这就是"具身模型格式"和"纯文本 LLM 格式"最本质的区别:除了权重和结构,还必须携带一整套与具体机器人本体绑定的语义元信息。一个 Llama ckpt 下载下来就能推理,但一个 π0 ckpt 如果不知道目标机器人的动作空间定义、归一化统计量、相机配置,就完全无法使用。

一个现实痛点:具身模型格式尚未统一

这个领域的格式还在演化中,不同框架之间的转换有真实的坑。一个典型案例:openpi(JAX 训练)导出的 π0 PyTorch ckpt,目前没有官方路径能转换成 LeRobot 兼容的 policy 目录——不是权重本身读不了,而是缺少 policy_preprocessor/postprocessor 配置、normalizer.safetensors、以及 LeRobot 特定的 config.json 结构。这背后反映的是:具身模型格式的"语义元信息"部分还处在各家自定义的状态,没有像 LLM 那样形成 HF 一统天下的格局。

6.5 三者的转换关系与选型逻辑

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
                  ┌─────────────────────────────────────┐
                  │   同一个 VLA 模型权重,三种形态       │
                  └─────────────────────────────────────┘
                                   │
       ┌───────────────────────────┼───────────────────────────┐
       ↓                           ↓                           ↓
┌─────────────┐          ┌──────────────────┐        ┌──────────────────┐
│ HF 格式      │←──转换──→│ Megatron 格式     │        │ LeRobot Policy   │
│ (交付/推理/  │          │ (大规模训练态)     │        │ (具身部署态)      │
│  LoRA微调)   │          │                   │        │                  │
└─────────────┘          └──────────────────┘        └──────────────────┘
       ↑                                                    ↑
       │                                                    │
  [Megatron Bridge /                                     [LeRobot 自带]
   NeMo-RL converters]

选型逻辑一句话:

  • 你要做千亿模型的 Megatron 3D 并行预训练 → 用 Megatron 格式,这是唯一高效方案
  • 你要发布开源模型、跑推理、做 LoRA → 用 HF 格式,通用货币
  • 你要部署到机器人本体、做真机评测、构建数据飞轮 → 用 LeRobot Policy 格式,因为它必须携带归一化统计量和前后处理器配置

对具身 Infra 工程师来说,最核心的心智转变是:具身模型不是"一个权重文件",而是"权重 + 本体配置 + 归一化统计 + 前后处理器"的完整包。这决定了你做过的全周期 ckpt 管理工作在具身场景下要多管一层:ckpt 不光要能恢复训练,还要能精确还原到某台机器人上能推理出正确动作。这个多出来的语义层,就是具身 infra 和 LLM infra 在"模型格式"这个话题上的分水岭。

参考