SuperAPI AI 工具集

nemo-automodel-distributed-training

智能体与模型

查看项目

更新于 10-16 03:49

技能介绍

用途:本技能用于在 NeMo AutoModel 中选择和配置分布式训练策略。NeMo AutoModel 采用 PyTorch 原生分布式训练,所有并行由单一 MeshContext 对象编排,持有设备网格、策略配置和轴名称。

主要能力:通过 YAML 的分布式策略键选择三种策略:fsdp2(通用默认,支持 TP、PP、CP、EP、HSDP)、megatron_fsdp(不支持 PP、EP、sequence_parallel)、ddp(仅简单数据并行,不支持 TP、PP、CP、EP)。可配置 tp_size、pp_size、cp_size、ep_size,dp_size 由 world_size / (tp_size * pp_size * cp_size) 推断。MoE 模型设 ep_size > 1 会创建独立的 moe_mesh,且 dp_size * cp_size 必须能被 ep_size 整除。长序列(8K+)可加 CP。支持激活检查点、延迟梯度同步、混合精度、序列打包。

使用方式:在 recipe YAML 的 distributed 段设置策略和字段,例如: distributed: strategy: fsdp2 tp_size: 8 pp_size: 4 cp_size: 1 ep_size: 1 pipeline: pp_schedule: interleaved1f1b pp_microbatch_size: 1 TP 应保持在同一 NVLink 域(通常单节点),跨节点用 PP 或 DP 扩展。

适用场景:单 GPU(无需分布式配置)、单节点多 GPU、多节点训练、70B+ 大模型(fsdp2 加 PP+TP)、MoE 专家并行、长序列训练。

限制:megatron_fsdp 不支持 PP、EP、sequence_parallel;DDP 仅数据并行;PP 需要模型类定义 _pp_plan;CP 需要 SDPA 或 Transformer Engine 注意力,不兼容 SDPBackend.MATH;dp_replicate_size 仅 FSDP2 可用。详情以官方资料为准。