研究型实习生-Training Infra工程师

AI 摘要智能生成

负责大模型分布式训练基础设施建设

职位描述

负责大模型分布式训练基础设施建设，支撑新模型结构快速接入及大规模训练任务稳定、高效运行。 1. 负责新模型结构在 Megatron 等并行训练框架中的适配，包括模型实现、并行切分、权重转换和训练链路接入； 2. 建设和优化多机多卡训练系统，合理组合 TP、PP、DP、CP、EP、SP 等并行策略； 3. 完善分布式 Checkpoint、自动 Resume、故障检测及容错恢复机制，提升长周期训练任务的稳定性； 4. 支持 Attention、稀疏 Attention、线性 Attention、MoE 等模型结构的高效训练； 5. 定位和解决长上下文及超大规模训练中的显存、通信、负载均衡、数据流水线和算子瓶颈； 6. 建设多卡数值对齐、梯度正确性、Checkpoint 一致性、显存和吞吐回归测试体系； 7. 优化混合精度训练、通信计算重叠、显存管理及关键 CUDA/Triton 算子。

任职要求

• 熟练掌握 Python、PyTorch，熟悉 Megatron-Core、Transformer Engine 等大模型训练框架；
• 有新模型结构接入分布式训练框架的实践经验，能够独立完成模型适配、权重转换和训练验证；
• 深入理解 TP、PP、DP、CP、EP、SP 等并行策略的原理、约束及组合方式；
• 理解 Attention、稀疏 Attention、线性 Attention、MoE 等模型结构及其并行训练特点；
• 熟悉 PyTorch Autograd、混合精度、梯度累积、分布式通信和分布式 Checkpoint；
• 能够系统定位 Loss 异常、数值不稳定、训练 Hang、OOM、通信超时和性能退化等问题。
• 具备多机多卡训练的数值正确性与性能测试能力。
• 具备良好的系统设计、工程实现和复杂问题排查能力。
• 具备 CUDA、Triton 或其他高性能算子开发和优化经验；
• 熟悉 NCCL、RDMA、通信拓扑以及通信计算重叠优化；
• 有 MoE、超长上下文或千卡级大模型训练经验；
• 熟悉 FP8/BF16 训练、Loss Scaling、累加误差及数值稳定性优化；
• 有训练平台、自动容错、可观测性或大规模集群调度系统建设经验。