零一万物AI Infra工程师面试经历:分布式训练+GPU优化+集群管理全考察
3年AI Infra经验面试零一万物AI Infra工程师岗位,三轮技术面详细复盘,涵盖3D并行与ZeRO优化、CUDA编程与GPU优化、集群调度与容错设计等核心考点
背景介绍
我做了3年AI Infra相关的工作,之前在一家大厂的AI平台部门负责GPU集群管理和分布式训练框架的维护,日常和DeepSpeed、Megatron打交道,也做过CUDA kernel优化。零一万物是李开复博士创办的AI公司,技术氛围很好,他们的AI Infra团队在招人,我觉得这是个很好的机会,既能继续做我擅长的方向,又能参与从零构建大模型训练基础设施的过程。投简历后大概5天约了面试。
面试流程复盘
一面:分布式训练 + DeepSpeed/Megatron(约2小时)
一面的面试官是Infra团队的核心工程师,上来就问分布式训练的问题,节奏很快。
第一个问题:数据并行、张量并行、流水线并行的原理分别是什么?我从三个维度讲了:数据并行是每个GPU持有完整模型副本,不同数据切片;张量并行是把模型的权重矩阵切分到不同GPU上;流水线并行是按层切分模型,不同GPU负责不同层。面试官追问了3D并行中这三种并行方式怎么组合,我说了通常先做流水线并行(层间切分),再在每个stage内做张量并行(层内切分),最后做数据并行。面试官点了点头。
接下来是DeepSpeed的重点:ZeRO的三个级别分别优化了什么?我详细讲了ZeRO-1切分优化器状态、ZeRO-2额外切分梯度、ZeRO-3额外切分模型参数,以及每个级别节省的显存比例。面试官追问了ZeRO-3的通信开销怎么优化,我说了预取参数、通信计算重叠、连续内存分配几个方法。
然后问Megatron:Megatron-LM的张量并行是怎么做的?我讲了列并行和行并行的切分方式,以及在前向和反向传播中如何通过All-Reduce同步。面试官追问了Megatron的序列并行是什么,我说了把LayerNorm和Dropout的激活值也切分到不同GPU上,进一步节省显存,面试官说理解正确。
还问了一个比较深入的问题:混合精度训练的原理和注意事项,我讲了FP16计算+FP32主权重+Loss Scaling的流程,面试官追问了动态Loss Scaling和静态Loss Scaling的区别,我说了动态会根据梯度是否溢出自动调整scale,更稳健但稍慢。
最后问了一个实战题:如果训练过程中出现OOM,你会怎么排查?我说了先看是哪个阶段OOM(前向/反向/优化器更新),然后分析显存占用(模型参数/梯度/优化器状态/激活值),再根据情况选择优化策略(ZeRO级别/梯度检查点/激活重计算)。面试官觉得排查思路比较系统。
二面:GPU优化 + CUDA(约2小时)
二面的面试官是做GPU优化的专家,问得非常硬核。
开场就是:CUDA的线程层次结构是什么?我讲了Grid、Block、Thread三级结构,以及Warp的概念。面试官追问了Warp Divergence是什么?对性能有什么影响?我讲了同一个Warp内的线程如果执行不同的分支路径,会串行执行,导致性能下降。面试官追问了怎么避免Warp Divergence,我说了重新组织线程映射、使用分支消除技巧、调整数据布局等方法。
然后是重点:GPU的内存层次结构是什么?各有什么特点?我讲了Global Memory(大但慢)、Shared Memory(小但快,同一Block内共享)、寄存器(最快但最少)、L1/L2 Cache。面试官追问了Shared Memory的Bank Conflict是什么?怎么避免?我讲了32个Bank的并行访问机制,如果多个线程访问同一个Bank的不同地址就会冲突,可以通过Padding、调整访问模式来避免。
CUDA编程实战题:写一个高效的矩阵乘法Kernel。我现场写了使用Tiling技术的Kernel,利用Shared Memory减少Global Memory访问。面试官看了之后问了几个优化点:怎么进一步优化?我说了使用向量内存访问(float4)、寄存器分块、双缓冲预取、Warp-level矩阵乘法指令。面试官说优化思路不错。
还问了一个很实际的问题:怎么Profile CUDA程序的性能?我介绍了Nsight Systems和Nsight Compute两个工具,前者看全局的时间线和瓶颈,后者看单个Kernel的详细性能指标。面试官追问了常见的性能瓶颈有哪些,我说了内存带宽瓶颈、计算瓶颈、启动开销、同步开销。
最后问了一个综合题:大模型训练中的通信和计算怎么重叠?我讲了梯度分片后边计算边通信、通信域的拓扑优化、NCCL的通信算法选择。面试官说理解得比较全面。
三面:集群管理 + 项目深挖(约1.5小时)
三面是Infra团队的负责人,聊集群管理和项目经验。
GPU集群的调度系统怎么设计?我讲了基于Kubernetes的GPU调度、多租户隔离、优先级调度、弹性伸缩几个核心功能。面试官追问了GPU碎片区怎么处理,我说了碎片整理、任务排队等待、小任务聚合几个策略,面试官补充说还可以做GPU时间片共享。
训练任务的容错机制怎么设计?我讲了定期保存Checkpoint、故障检测(进程心跳/NCCL超时)、自动重启恢复、弹性训练几个方面。面试官追问了Checkpoint的保存策略怎么优化,我说了异步保存、增量保存、分布式保存几个方法,面试官觉得方案比较成熟。
项目深挖环节,面试官让我讲之前做的GPU集群管理项目。他问得特别细:集群规模多大?多少个GPU?调度延迟多少?故障恢复时间多少?我一一回答了,还说了我们遇到的一个难题:大规模训练时NCCL通信经常超时,我们通过调整网络拓扑和NCCL参数解决了。
最后问了一个系统设计题:设计一个支持千卡级别大模型训练的集群管理系统。我从资源管理、任务调度、故障恢复、监控告警、成本优化几个模块设计了方案,面试官说架构设计合理,但提醒我要注意网络拓扑感知的调度和跨机房训练的网络优化。
真题汇总
1. 数据并行、张量并行、流水线并行的原理和3D并行组合
2. ZeRO三个级别的优化内容和通信开销优化
3. Megatron-LM张量并行和序列并行
4. 混合精度训练原理和动态/静态Loss Scaling
5. 训练OOM排查思路
6. CUDA线程层次结构和Warp Divergence
7. GPU内存层次结构和Shared Memory Bank Conflict
8. 高效矩阵乘法Kernel优化
9. CUDA性能Profile工具和常见瓶颈
10. 大模型训练中通信和计算重叠
11. GPU集群调度系统设计
12. GPU碎片区处理
13. 训练任务容错机制设计
14. Checkpoint保存策略优化
15. 千卡级别训练集群管理系统设计
心得建议
1. 分布式训练是AI Infra的核心:3D并行、ZeRO、Megatron这些必须深入理解,不只是知道概念,还要能说清楚实现细节和优化策略。
2. CUDA编程能力是硬门槛:AI Infra岗位对CUDA要求比其他方向高很多,线程模型、内存层次、性能优化这些都要扎实,面试官会直接让你写Kernel。
3. 集群管理经验很加分:GPU调度、容错恢复、监控告警这些工程经验面试官很看重,有Kubernetes经验会加分。
4. 要有大规模系统的经验:AI Infra不只是单机优化,还要考虑千卡级别的集群管理,网络拓扑、通信优化、故障恢复这些系统级问题面试官会考察。
5. 性能调优要有方法论:不是凭感觉优化,而是先Profile找到瓶颈,再有针对性地优化,面试官很看重你的系统性思维。
FAQ
Q:面试对CUDA编程要求多高?
A:要求比较高,二面直接让写矩阵乘法Kernel,还要说优化方案。建议至少写过几个常用的Kernel,理解Tiling、Shared Memory优化等基本技术。
Q:需要了解零一万物的技术栈吗?
A:面试中没直接问,但从问题来看,用了DeepSpeed、Megatron、自研的调度系统,了解这些框架会有帮助。
Q:对网络知识有要求吗?
A:有,特别是RDMA、InfiniBand、NCCL这些,大规模训练中网络通信是瓶颈,面试官会问相关的问题。
Q:面试有手写代码吗?
A:一面有写伪代码,二面有写CUDA Kernel,三面主要是系统设计,代码量比其他方向多。
Q:面试周期多长?
A:从投递到三面结束大概三周,每面间隔4-7天,节奏偏慢但可以充分准备。