什么是 张量并行 (intra-layer)
1️⃣ 考察意图
面试官想考察你对大规模模型分布式训练中“模型并行”的深入理解，特别是张量并行（Tensor Parallelism, TP） 作为 intra-layer 切分的核心机制。这属于系统设计 + 工程取舍类型题目。刁钻点在于：不仅要背定义，还要能清晰区分 TP 与流水线并行（PP）、数据并行（DP）的通信模式和适用场景，并解释为什么 TP 在 Transformer 中更“重计算、轻通信”。答好了能展示你对 Megatron-LM、DeepSpeed 等框架底层原理的掌握，以及处理超大规模模型（如 175B+）时的显存与通信权衡能力。
2️⃣ 标准答
定义与核心思想张量并行（intra-layer parallelism）是将神经网络单层内的矩阵运算（如 Transformer 的注意力 QKV 投影、FFN 的线性变换）切分到多个 GPU 上并行执行。每个 GPU 只持有权重矩阵的一部分，计算完成后通过通信聚合部分结果。这与流水线并行（inter-layer）将不同层分配到不同设备完全不同。
实现方式：Megatron-LM 的列并行与行并行以 Transformer 的 FFN 层（Y = GeLU(X * A) * B）为例，Megatron-LM 采用两种切分：
- 列并行（Column Parallel）：将权重矩阵
A按列切分到N个 GPU，每个 GPU 计算X * A_i（A_i是A的列切片）。前向时无需通信，反向时需对A_i的梯度做 AllReduce 求和。 - 行并行（Row Parallel）：将权重矩阵
B按行切分到N个 GPU，每个 GPU 持有B_i（B的行切片）。前向时，每个 GPU 先计算GeLU(X * A_i)，然后通过 AllReduce 聚合所有 GPU 的部分结果得到最终输出Y。反向时，梯度也需 AllReduce 分发。
通信模式与 trade-off
- 通信量：每个 Transformer 层需要 2 次 AllReduce（前向一次、反向一次），通信量为
O(batch_size * seq_len * hidden_dim)。对于 175B 模型，单次 AllReduce 数据量可达数百 MB，对互联带宽要求极高（必须 NVLink/NVSwitch，否则成为瓶颈）。 - 计算 vs 通信：TP 是计算密集型切分——每个 GPU 计算量减少为原来的
1/N，但通信量不随 N 线性减少（因为 AllReduce 需要聚合所有部分）。因此，TP 适合单节点内（8 卡）使用，跨节点时通信延迟会急剧恶化。 - 显存节省：每个 GPU 只存储
1/N的权重和优化器状态（如 Adam 的 momentum 和 variance），显存占用从16 * param_size降到16 * param_size / N（假设混合精度训练）。这是 TP 的核心价值——让单卡装不下的大模型跑起来。
实际落地的坑 + 解法
- 坑 1：负载不均衡。如果切分维度（如 hidden_dim）不能被 GPU 数整除，部分 GPU 会多算或少算。解法：在模型设计时确保 hidden_dim 是 GPU 数的倍数（如 8 卡时 hidden_dim 设为 8192 而非 8000）。
- 坑 2：通信与计算重叠。AllReduce 是同步操作，会阻塞计算。解法：使用 Tensor Fusion（Megatron-LM 的
fuse_communication参数）将多个小 AllReduce 合并为一次大通信，减少通信次数；或利用 CUDA 流（stream）将 AllReduce 与后续计算重叠。 - 坑 3：与数据并行（DP）共存。TP 通常与 DP 结合（即 3D 并行：DP + TP + PP）。此时，TP 组内通信（AllReduce）和 DP 组间通信（AllReduce 梯度）会竞争带宽。解法：设置
tensor_model_parallel_size和data_parallel_size，确保 TP 组内使用 NVLink，DP 组间使用 InfiniBand，避免带宽争抢。
与流水线并行（PP）的区别
- TP：切分单层，通信量大，适合单节点内。
- PP：切分层间，通信量小（仅传递激活值），适合跨节点。
- 典型组合：8 卡节点内用 TP（
tp=8），跨 4 个节点用 PP（pp=4），每个节点内再叠加 DP（dp=2），总卡数 = 842=64。
3️⃣ 答题模板（30 秒电梯版）
“这个问题我从定义、实现方式、通信模式三个层面回答。定义上，张量并行是将单层内的矩阵运算切分到多个 GPU，每个 GPU 只持有一部分权重。实现上，以 Megatron-LM 为例，FFN 层采用列并行和行并行，前向和反向各需一次 AllReduce 通信。通信模式上，TP 是计算密集型切分，通信量随 hidden_dim 线性增长，适合单节点内使用，必须搭配 NVLink。总结一句：张量并行是训练超大规模模型（如 175B+）时，解决单卡显存瓶颈的关键技术，但通信开销高，需与流水线并行、数据并行组合成 3D 并行。”
4️⃣ 高频追问 & 应对
追问 1：张量并行和流水线并行在训练 175B 模型时，哪个通信开销更大？为什么？
应对策略：张量并行通信开销更大。因为 TP 每层需要 2 次 AllReduce，通信数据量为
2 * batch_size * seq_len * hidden_dim（约 2 * 4 * 2048 * 12288 ≈ 200MB/层）。而 PP 每层只需传递激活值（batch_size * seq_len * hidden_dim），且跨节点时通信频率低（每 micro-batch 一次）。所以 TP 适合单节点内（NVLink 带宽 600GB/s），PP 适合跨节点（InfiniBand 带宽 200GB/s）。实际中，TP 组大小通常 ≤ 8，PP 组大小可到 32+。
追问 2：如果我用 4 张 A100 做张量并行，但 AllReduce 通信占用了 40% 的时间，怎么优化？
应对策略：三个方向。1）通信与计算重叠：使用 Megatron-LM 的
async_tensor_model_parallel_allreduce参数，让 AllReduce 与后续的 GeLU 激活计算在 CUDA 流上并行执行，可减少 20-30% 的通信开销。2）减少通信量：检查是否开启了fuse_communication，将多个小 AllReduce 合并为一次大通信。3）调整切分策略：如果 hidden_dim 不是 4 的倍数，补零到 4 的倍数，避免负载不均衡导致部分 GPU 等待。如果仍不行，考虑将 TP 组大小从 4 降到 2，改用 PP 或 DP 补足。
追问 3：张量并行中，为什么 FFN 层用列并行 + 行并行，而注意力层用另一种切分？
应对策略：因为注意力层有 QKV 投影和输出投影。Megatron-LM 对注意力层采用：QKV 投影用列并行（每个 GPU 计算
X * [Q_i, K_i, V_i]），输出投影用行并行。这样前向时，QKV 投影无需通信，输出投影需要一次 AllReduce。而 FFN 层是GeLU(X * A) * B，列并行切分A，行并行切分B，前向需要一次 AllReduce。两种切分都保证了每个 GPU 的计算量相等，且通信次数最少（每层 2 次 AllReduce）。
5️⃣ 避坑 · 常见错误答法
- ❌ “张量并行就是把模型参数分到不同 GPU 上，每个 GPU 算一部分，最后合并。” → ✅ 必须明确是单层内的切分，且区分列并行和行并行，说明 AllReduce 通信时机（前向和反向各一次）。
- ❌ “张量并行通信开销小，因为每个 GPU 只算一部分。” → ✅ 通信开销不小，因为 AllReduce 需要聚合所有部分结果，通信量随 hidden_dim 线性增长。实际中 TP 组大小通常 ≤ 8，否则通信成为瓶颈。
- ❌ “张量并行和流水线并行一样，都是模型并行。” → ✅ 必须强调区别：TP 是 intra-layer（切分单层），PP 是 inter-layer（切分层间），通信模式和适用场景完全不同。
6️⃣ 简历呼应
- 如果你有大规模分布式训练项目：从“在 64 卡集群上训练 13B 模型时，对比了 TP=8 和 TP=4 的吞吐量差异，发现 TP=8 时 AllReduce 通信占比从 15% 升到 30%，最终选择 TP=4 + PP=2 + DP=8 的组合”切入，展示工程调优能力。
- 如果你只做过单卡小模型：用“类比”迁移：单卡训练时矩阵乘法是
[batch, hidden] * [hidden, 4*hidden]，TP 相当于把[hidden, 4*hidden]按列切分成 2 份，每张卡算[batch, hidden] * [hidden, 2*hidden]，最后 AllReduce 合并。强调“计算量减半，通信量不变”的 trade-off。 - 如果你是校招无项目：聚焦 Megatron-LM 论文（Shoeybi et al., 2019）的复现 demo：在 2 张 GPU 上用 PyTorch 实现 FFN 层的列并行和行并行，测量显存从 4GB 降到 2.2GB，但通信时间增加 0.3ms。展示对论文原理的动手理解。
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)
- DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters (Rasley et al., 2020)
- 博客：NVIDIA 官方 “Tensor Parallelism in Megatron-LM” 技术白皮书
- 论文：Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022) —— 讨论 TP 下的激活内存优化