职位描述
-负责 PyTorch 基础设施在昆仑芯 AI 芯片上的适配与优化,涵盖算子、分布式训练、编译及运行时等基础能力。
-负责 Megatron、DeepSpeed、Ray、verl、slime、Swift+TRL 等训练及强化学习框架、产品,以及 vLLM、SGLang 等推理框架的适配、调试与落地。
-面向大模型训练和推理场景,开展算子、显存、并行策略、通信与计算融合及端到端性能优化。
-负责模型精度、性能和稳定性分析,支持千卡、万卡规模集群调试及客户需求落地。任职要求
-计算机、软件工程、电子信息等相关专业,本科及以上学历。
-熟练掌握 Python、C/C++,熟悉 Linux 开发环境,具备良好的编程和问题分析能力。
-熟练掌握 PyTorch 框架本身,能够独立完成模型接入、调试、扩展和性能优化。
-掌握深度学习基础理论,了解 Transformer、大语言模型的训练与推理流程。
加分项
-了解 TensorFlow,熟悉 torch.compile、TF XLA、Triton 等 AI 编译技术。
-熟悉 Mega Kernel、Flux 等通信与计算融合技术。
-熟悉 CUDA 算子编写及性能优化。
-熟悉数据并行、张量并行、流水线并行、序列并行、专家并行等并行策略。
-具备精度调试、性能 Profiling 或千卡、万卡规模集群调试经验,熟悉 PyTorch Profiler、Nsight Systems、Nsight Compute 等工具。