小米 校园招聘
顶尖应届-大模型训练框架优化工程师-自动驾驶
北京 ·经验不限·学历不限
北京
薪资面议
前往 小米 官网投递
职位描述
1、负责基于现有大模型训练框架分析自动驾驶与机器人模型训练链路,包括数据加载、分布式并行、显存占用、通信开销和算子执行等环节,定位性能瓶颈;
2、围绕训练吞吐、GPU 利用率和模型迭代效率开展优化,推动训练周期缩短和训练资源成本下降;
3、结合 PyTorch、Megatron-LM、DeepSpeed 等技术栈,开展分布式训练策略、混合精度、算子融合、通信优化和 CUDA 高性能算子优化;
4、推进训练框架与 AI 平台、调度系统和监控体系协同,建立训练效率、集群算力利用率和任务稳定性的评估机制;
5、探索 AI Agent 在自动化 profiling、瓶颈定位和代码优化建议中的应用,沉淀可复用的训练优化工具链。
【任职要求】
1、本科及以上学历,扎实的编程基础(Python/C++)与良好的工程习惯;
2、精通深度学习框架底层原理(PyTorch 核心机制如 torch.compile、Autograd、TorchScript、XLA);
3、熟悉分布式训练框架(Horovod、Megatron-LM、DeepSpeed 等)的实现逻辑;
4、熟悉 CUDA 编程及 GPU 架构,能独立开发高性能自定义算子(CUDA、CUTLASS、TileLang、Triton);
5、理解大模型训练框架工程化实践,具备在现有框架基础上进行性能分析、模块扩展或业务适配的能力;
6、有深度学习编译技术栈(TVM、MLIR 等)、AI 开源社区贡献、训练性能优化或大规模模型训练经验者优先。
官网发布:2026-06-03 · 最后确认在招:2026-09-03 23:17:24 · 来源平台:xiaomi