SuperAPI AI 工具集

nemo-automodel-distributed-training

通用与其它

查看项目

更新于 10-16 02:33

技能介绍

用途:指导在 NeMo AutoModel 中选择与配置分布式训练策略。所有并行由单一 MeshContext 对象编排,保存设备网格、策略配置与轴名称。

主要能力:三种策略通过 distributed.strategy 选择。fsdp2 为推荐默认,支持 TP、PP、CP、EP、HSDP;megatron_fsdp 为 Megatron 风格 FSDP,不支持 PP、EP、sequence_parallel;ddp 仅支持简单数据并行,不支持 TP、PP、CP、EP。dp_size 始终由 world_size / (tp_size * pp_size * cp_size) 推断。

使用方式:在 recipe YAML 的 distributed 段配置,例如 strategy: fsdp2、tp_size、pp_size、cp_size、ep_size,以及 pipeline 与 moe 子配置。也可用 Python API:DistributedSetup.build、FSDP2Config、ParallelismSizes、initialize_distributed,并传入 from_pretrained。

适用场景:单 GPU 无需分布式配置;多 GPU 单节点用 fsdp2;多节点用 fsdp2 配合 TP/PP;MoE 专家并行用 fsdp2 且 ep_size > 1,会创建独立 moe_mesh;70B+ 大模型用 PP + TP;8K+ 长序列加 CP。

限制:TP 应保持在单节点 NVLink 域内,跨节点会严重降低吞吐。PP 要求模型类定义 _pp_plan。megatron_fsdp 不支持 EP、PP、sequence_parallel。EP 要求 ep_size 整除 dp_size * cp_size。CP 不兼容 SDPBackend.MATH。HSDP 要求 dp_replicate_size < dp_size。具体硬件与版本要求资料未提供。