使用 Megatron-Infinigence 训练工具
Megatron-Infinigence 是一个专为大规模分布式深度学习训练设计的训练工具组合，集成了多种高性能计算工具、深度学习框架和优化库。该工具可帮助用户在大规模集群上高效地完成对于 MoE 模型预训练。
目前对 Deepseek-V3(671B) 以及 Qwen3-235B-A22B 模型的训练做了性能优化。
技术特性
- DualpipeV 流水线调度: 基于 DeepSeek 优化的流水线方案，实现计算与通信的深度重叠，提升 MoE 模型训练效率
- Flash Attention 3 原生支持: 针对 Hopper 架构 GPU 的最新优化，需配合 local transformer-impl 使用
- TEgroupgemm 融合算子: 集成 Transformer Engine 的高效 groupgemm，将多个 expert 矩阵乘融合为单一算子
- 智能显存管理: MLP 层部分重算策略，精确控制 fc1/fc2 重算，在显存节省与性能间达到最优平衡
- 计算通信异步重叠: 支持前后计算与 MLP allgather 通信的完全重叠，最大化硬件利用率
- 高性能网络栈: 预配置 HPCX + RDMA + InfiniBand 完整解决方案，支持超大规模集群的低延迟通信
训练工具容器镜像
我们将 Megatron-Infinigence 训练工具组合打包为容器镜像，用户可直接拉取镜像进行训练，无需手动配置环境。
镜像名称： infini-ai/megatron-infinigence:v1-ngc25.04-20250725
训练框架代码
镜像中包含了：
- Megatron-infinigence 高性能 MoE 训练框架：
/root/megatron-infinigence/ - Loss 对齐工具：
/root/megatron-infinigence/tools/loss_align/ - 显存优化工具：
/workspace/memory_tools
其他工具
除训练框架外，infini-ai/megatron-infinigence:v1-ngc25.04-20250725 镜像中还内置了以下工具。
网络和通信优化
- NCCL Tests: NVIDIA 集合通信库测试工具（v2.15.1）
- HPCX: 高性能计算通信库
- RDMA Core: 远程直接内存访问核心库（v57.0）
- Mellanox Tools: Mellanox 网络工具集（v25.07.0）
性能测试和基准测试
- nvbandwidth: GPU 带宽测试工具（v0.7）
- FIO: 灵活的 I/O 测试工具（v3.40）
- GPU-Burn: GPU 压力测试工具
- VDBench: 存储性能测试工具
- PerfTest: InfiniBand 性能测试工具
预训练模型和数据集
镜像预置了多个主流大语言模型的配置文件和分词器：
- DeepSeek-V3: 最新的 DeepSeek 系列模型
- Qwen3-30B-A3B: 通义千问 3 系列模型
- Qwen2.5-7B: 通义千问 2.5 系列模型
- Llama-2-7b: Meta 的 Llama 2 模型
- Mixtral-8x7B: Mistral AI 的混合专家模型
- GPT-2: OpenAI 的 GPT-2 模型
数据处理和工具
- WuDao 数据集: 预处理好的开源数据集
- OSS 工具: 阿里云对象存储服务客户端
- AWS CLI: 亚马逊云服务命令行工具
- Kafka: 分布式流处理平台客户端
监控和调试工具
- Infinigence Telemetry Kit: 无问芯穹自研的遥测监控工具包
- Memory Tools: 显存管理和优化工具
- Troubleshoot: 故障排查工具集
- Infini-TBench: 性能基准测试套件