工作职责: 1. 分布式训练底座 (Training Infra) 架构设计: 负责维护和优化基于 Megatron-LM, DeepSpeed, Ray 或 FSDP 的大规模分布式训练框架
通信优化: 深入优化 H/NCCL 通信库,解决 RDMA/RoCE 网络下的通信瓶颈,提升多机多卡并行效率(DP/PP/TP/CP/EP)
稳定性保障: 构建自动容错与快速恢复系统(Checkpoint 优化、故障自动检测与接续),确保千卡集群在数月跨度的训练中保持极高可用性
2. 推理加速与工程化 (Inference Infra) 高性能引擎: 深度定制或调优 vLLM, TensorRT-LLM, Triton Inference Server 等推理框架
算子开发: 针对 Transformer 结构编写高性能 CUDA/Triton kernels(如 FlashAttention, PagedAttention 的底层实现或改进)
资源调度: 优化 K8s 集群下的 GPU 资源调度,实现动态扩缩容、请求批处理(Continuous Batching)以及 Prefix Caching
3. 存储与算力管理 (Storage & Compute) I/O 优化: 优化超大规模数据集的加载速度,解决训练过程中的存储带宽瓶颈(如利用 GPFS, Lustre 或 JuiceFS)
监控与观测: 建立精细化的 GPU 利用率、显存压力、网络带宽监控体系,定位并消除“长尾效应”和计算气泡