10 · Tensor Parallel：把一层模型拆到多张 GPU

为什么不是简单地复制模型

数据并行会在每张 GPU 放完整模型，适合训练或并行处理独立批次；但当单个模型放不进一张 GPU，或者希望单次前向由多 GPU 协作时，需要 Tensor Parallel，把层内矩阵切分。

两种基本线性层

Column Parallel

按输出维度切权重矩阵。每张 GPU 计算一部分输出列。QKV 投影和 MLP 的上投影常用这种方式，因为后续计算也可以继续使用分片输出。

Row Parallel

按输入维度切权重矩阵。每张 GPU 计算局部和，随后通常通过 all_reduce 汇总完整输出。Attention 输出投影和 MLP 下投影常见这种方式。

nano-vLLM 的实现位置

• layers/linear.py：并行线性层、权重加载器和通信。
• layers/embed_head.py：词嵌入与 LM Head 的词表并行。
• models/qwen3.py：根据 world size 计算每个 rank 的 head 数和维度。
• engine/model_runner.py：每个 rank 初始化 NCCL 进程组与对应 CUDA device。

多进程控制

rank 0 负责主控制流，其他 rank 作为 worker。源码使用共享内存传递序列化调用参数，并用事件 / barrier 协调。每张 GPU 都执行同一个逻辑操作，但持有不同权重分片。

必须满足可整除

总 attention head 数、KV head 数、相关隐藏维度必须能被 tensor parallel size 整除。源码中的 assert 不是保守限制，而是保证每个 rank 得到一致分片。

通信可能抵消收益

更多 GPU 不一定更快。小模型或小 batch 时，NCCL 通信、进程同步和 kernel 启动可能比计算本身更显著。Tensor Parallel 首先解决容量问题，其次才是性能问题。

HTML INTERACTIVE LAB10 · Tensor Parallel：把一层模型拆到多张 GPU

单独打开 ↗

课后习题等待完成

Row Parallel 线性层为什么通常需要 all-reduce？

动手任务

选择一个 8×8 矩阵和 2 张 GPU，分别画出 Column Parallel 与 Row Parallel 的切分方向，并标出哪一种需要汇总输出。

下一节：CUDA Graph