AI Infra 高级研发工程师

职位描述

负责 Momenta 自动驾驶大模型、世界模型的分布式训练推理框架研发，主要包括：

1. 参与设计、研发、维护团队内部的模型分布式训练框架，擅长分析并深度优化训练各个阶段的性能瓶颈，包括计算效率、通信延迟、显存占用等；

2. 结合不断迭代的模型算法逻辑，设计并实现针对性的高效分布式并行训练策略；

3. 深入研究 CUDA、NCCL、RDMA 等编程范式和通信库，针对团队内部的 GPU 和 集群拓扑约束，开发高性能算子 并 优化分布式通信效率，达到行业 SOTA 水平；

4. 深入研究低精度混合精度训练策略，在保证模型精度满足预期的情况下，探索低精度（FP8、FP4）训练的性能极限；

5. 配合算法需求，开发 RL 训练框架、迭代 RL 训练算法逻辑，优化在线 Rollout 推理性能，深度优化 分布式 On-Policy/异步 RL 训练效率。

任职要求

1. 学历背景： 计算机相关专业硕士及以上学历，具备扎实的理论基础和强大的动手能力；

2. 编程能力：

- 熟悉 Python/C++/CUDA/CUTLASS/Triton 编程，熟悉 PyTorch 框架及其底层实现；

- 熟悉分布式系统开发与调试，熟悉分布式通信开发 NCCL/RDMA/IB/RoCE，有多进程调度与并行算法优化经验者优先；

3. 训练与推理经验：

- 熟悉大模型结构与算法（LLM/VLM/VLA/DiT、MoE 架构、各种 RL 算法 等）；

- 熟悉分布式训练/推理的常用策略（DP/TP/PP/EP/CP、recompute、offloading、PD 分离 等）；

- 有主流开源分布式训练框架（如 Megatron-LM、DeepSpeed、VeRL）或推理框架（如 SGLang、vLLM）的深度二次开发与优化经验。