ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

VLA训练栈深度揭秘:OpenDM的Muon+AdamW优化器、FSDP分布式训练与混合精度策略

VLA训练栈深度揭秘:OpenDM的Muon+AdamW优化器、FSDP分布式训练与混合精度策略 VLA训练栈深度揭秘OpenDM的MuonAdamW优化器、FSDP分布式训练与混合精度策略【免费下载链接】opendmAn Open-World Foundation Model for General-Purpose Embodied Intelligence.项目地址: https://gitcode.com/gh_mirrors/op/opendmOpenDM 是面向开放世界具身智能的视觉-语言-动作VLA基础模型 DM0.5 的开源训练与推理框架。本文将带你完整拆解它的 VLA 训练栈核心自研的 MuonAdamW 混合优化器、基于 FSDP 的多卡分布式训练以及 bf16/fp32 双套混合精度策略帮助你快速理解大模型微调背后的工程细节。 先搞清楚OpenDM 训练栈由什么组成OpenDM 的训练入口集中在opendm/exp/dm05_exp.py整个训练流程由四块配置驱动模型配置、优化器配置、Trainer 配置和数据配置。其中与训练效率直接相关的三大件是优化器MuonAdamW 混合优化器代码在 opendm/optimizer/muon_adamw.py分布式训练基于 HuggingFace Transformers Trainer 扩展的DMTrainer代码在 opendm/trainer/trainer.py混合精度策略精度常量与策略定义在 opendm/constants/precision.py多卡训练通过 script/dm05_launcher.sh 一键拉起它内部调用torchrun启动 8 卡进程并设置 NCCL 通信参数上手体验可参考 docs/zh/dm05_finetuning.md。 核心揭秘一Muon AdamW 混合优化器这是 OpenDM 训练栈最有意思的部分。为什么不用一个优化器打天下因为 VLA 模型里有两类性格完全不同的参数VLM 主干视觉编码器 语言模型参数量大、梯度稀疏敏感适合稳定的 AdamWAction Expert动作专家模块的隐层权重矩阵全是 2D 矩阵负责把视觉语言表征转成机器人动作用 Muon 正交化更新收敛更快参数如何分流在 FSDP 包裹模型之前mark_muon_parameters会扫描所有可训练参数按以下规则选中 Muon 组见 opendm/optimizer/muon_adamw.py 中的is_default_muon_parameter条件说明必须 2 维Muon 只能作用于矩阵1D 参数如 bias、LayerNorm一律排除名称匹配参数路径包含action_expert且位于layers下、以.weight结尾排除项嵌入层、lm_head之外的 embedding、所有 norm 层未被选中的参数自动进入 AdamW 组。整个分流逻辑由 opendm/exp/dm05_exp.py 中DM05OptimizerConfig的optim字段控制设为muon_adamw即启用混合优化器adamw则退化为纯 AdamW。Muon 的工作原理牛顿-舒尔茨正交化MuonMomentUm的核心思想是不直接用梯度方向更新权重而是先把梯度正交化到最优的更新方向上。OpenDM 用5 步牛顿-舒尔茨Newton-Schulz迭代近似矩阵的零次幂见zeropower_via_newtonschulz5相比 SVD 分解只需几次矩阵乘法就能逼近算力友好得多。关键超参默认值DM05OptimizerConfig超参默认值含义muon_momentum0.95动量系数muon_nesterovTrue启用 Nesterov 动量muon_ns_steps5牛顿-舒尔茨迭代步数muon_moonlight_coefficient0.2学习率缩放系数base_lr2.5e-5基础学习率Muon 组的有效学习率还有一个巧妙的修正lr × muon_lr_scale × 0.2 × sqrt(max(shape))用矩阵最大维度开根号来平衡不同形状层的更新幅度避免大矩阵更新过小、小矩阵更新过大。为 FSDP 分片量身定制的细节分布式训练下每个 GPU 只持有矩阵的一部分切片而正交化必须基于完整矩阵计算。OpenDM 为此设计了MuonShardPlan分片计划先通过all_gather探测各 rank 持有的分片大小拼出完整矩阵 → 执行牛顿-舒尔茨正交化 → 再切回本 rank 的切片做参数更新。这种分片计算、整矩阵正交的设计是 Muon 能无缝跑在 FSDP 上的关键也解释了为什么必须在 FSDP 包裹前先记录参数形状。 核心揭秘二FSDP 分布式训练策略DMTrainer继承自 HuggingFace 的Trainer在_link_exp_config中为多卡场景配置了 FSDP见 opendm/trainer/trainer.pyfsdp shard_grad_op参数和梯度、优化器状态都跨卡分片显存占用最低适合 DM0.5 这种大模型use_orig_params True保留原始参数视图这正是 Muon 需要读取/写回完整矩阵形状的前提sync_module_states True多机场景下从 rank 0 同步初始权重保证各节点一致backward_prefetch BACKWARD_PRE提前预取下一层参数让通信与计算重叠seed rank每张卡使用不同随机种子让数据增强等随机过程真正并行化在自动包裹auto-wrap策略上OpenDM 做了自定义不采用 Transformers 默认的 PEFT 自动包裹而是让DM05ForConditionalGeneration.fsdp_wrap_modules()显式返回分片边界——整个 Action Expert VLM 的最后若干层作为一个单元既保证分片粒度合理又避免了逐层包裹带来的通信开销。此外训练还默认开启VLM 与 Action Expert 双路梯度检查点vlm_gradient_checkpointing/ae_gradient_checkpointing并用 Liger Kernel 融合 RMSNorm 等操作进一步压低显存、提升吞吐VLM 的 token 嵌入层及其共享权重的lm_head默认冻结。⚖️ 核心揭秘三bf16 与 fp32 两套混合精度策略OpenDM 把精度策略抽象为一个precision_policy字段定义在 opendm/constants/precision.py提供两档bf16_mixed默认策略模型以 bfloat16 存储走标准 AMP 自动混合精度路径。bf16 与 fp32 具有相同位宽和更大的动态范围训练大模型不易溢出是吞吐与稳定性的最佳平衡点。fp32_mixed高精度策略模型权重保持 fp32MODEL_DTYPE torch.float32仅前向计算通过torch.autocast以 bf16 执行见 opendm/model/dm05/dm05_utils.py 的dm05_autocast同时显式关闭 TF32TF32_ENABLED False保证矩阵乘精度可控FSDP 侧通过set_mixed_precision(fp32)覆盖默认配置权重分片以 fp32 保存敏感线性层如动作投影通过linear_fp32强制在 fp32 下完成投影避免低精度污染动作输出这个策略适合对数值精度敏感的场景代价是显存和速度。两种策略只需一个配置项切换--model-config.precision-policy bf16_mixed # 默认追求吞吐 --model-config.precision-policy fp32_mixed # 高精度追求数值稳定 训练栈默认参数速查表项目默认值配置位置优化器adamw可切 muon_adamwDM05OptimizerConfig.optim基础学习率2.5e-5DM05OptimizerConfig.base_lr学习率调度cosine_with_min_lrmin_lr_rate0.1DM05TrainerConfig梯度裁剪max_grad_norm 1.0DMTrainer._link_exp_config每卡 batch size4DM05TrainerConfig.per_device_train_batch_size精度策略bf16_mixedDM05ModelConfig.precision_policyFSDP 模式shard_grad_op use_orig_paramsDMTrainer多卡时自动启用梯度检查点VLM Action Expert 均开启DM05ModelConfig 如何动手从 Demo 训练开始理解了上面的机制上手只需三步启动多卡训练运行script/dm05_launcher.sh --exp playground/dm05_sft_demo.py --nproc_per_node 8 --task train ...脚本会自动设置 NCCL 环境变量并拉起 8 卡详见 script/dm05_launcher.sh切换优化器在实验配置中将--optimizer-config.optim muon_adamw传入即可体验 MuonAdamW 混合优化调整精度策略用--model-config.precision-policy fp32_mixed开启高精度模式更多数据集注册、归一化统计与训练流程细节建议配合官方文档 docs/zh/dm05_finetuning.md 和 docs/zh/dm05_libero_lora_training.md 一起阅读。OpenDM 这套Muon 正交化 FSDP 分片感知 双档混合精度的组合展示了 VLA 大模型训练栈如何在不牺牲收敛质量的前提下把显存和算力用到极致。掌握这些底层机制后你可以更有底气地为自己的机器人模型调参。【免费下载链接】opendmAn Open-World Foundation Model for General-Purpose Embodied Intelligence.项目地址: https://gitcode.com/gh_mirrors/op/opendm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表