岗位职责
负责大模型分布式训练基础设施建设,支撑新模型结构快速接入及大规模训练任务稳定、高效运行。
- 负责新模型结构在 Megatron 等并行训练框架中的适配,包括模型实现、并行切分、权重转换和训练链路接入;
- 建设和优化多机多卡训练系统,合理组合 TP、PP、DP、CP、EP、SP 等并行策略;
- 完善分布式 Checkpoint、自动 Resume、故障检测及容错恢复机制,提升长周期训练任务的稳定性;
- 支持 Attention、稀疏 Attention、线性 Attention、MoE 等模型结构的高效训练;
- 定位和解决长上下文及超大规模训练中的显存、通信、负载均衡、数据流水线和算子瓶颈;
- 建设多卡数值对齐、梯度正确性、Checkpoint 一致性、显存和吞吐回归测试体系;
- 优化混合精度训练、通信计算重叠、显存管理及关键 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、累加误差及数值稳定性优化;
- 有训练平台、自动容错、可观测性或大规模集群调度系统建设经验。