在之前的单卡阶段结束和进行阶段性的总结之后，经过几天的休息（实际上是作者被导师叫去忙活论文投稿和下一篇论文启动的事情了），我们终于可以进入Pico-vLLM开发的下一个阶段：分布式的开发。而分布式的开发，根据计划，我们将要实现如下几个重要特性：张量并行TP，流水线并行PP和PD分离。

首先做的是张量并行，Tensor Parallel。其实我在一开始规划的时候，曾经想过要不要先做流水线并行。然而经过进一步的思考和分析，我发现了一个令人惊讶的事实：张量并行这个看上去切分粒度更细、通信要求和频率更高的方案，在实现难度上居然是最低的，而pipeline并行这种方案，需要的改动反而更大。下面会详细解释实现所涉及到的改动，以及实现的具体技术手段和注意事项、细节。

张量并行涉及到的改动

张量并行涉及到的改动意外的少，这个少是相较于其他并行方案而言。这其实是因为张量并行并不改变任何一张单个GPU上所需要加载的模型的结构：所有层数没有变，block没有变，变化的只有Q head和KV head的头数。

模型

模型的结构并不需要修改，仅仅将num_heads相关的参数和hidden_dim相关的参数修改成相应的local版本即可，层数等等结构相关的都不需要改变。

KV cache的存储

和先前的模式一模一样，每个GPU对应的程序仍然要单独管理自己的KV cache，有自己的Scheduler。区别在于，每个KV cache存储的slot大小变化了，因为每个GPU对应的head数量变化了（被分给了不同的GPU），其中缩小的倍率和TP的size相同。除此之外，几乎没有任何需要修改的地方，包括Triton Kernel。

引擎Engine

几乎没有需要修改的地方，只需要在初始化的时候给cfg传入相应的额外参数即可。

另一个需要注意的事情是，如果采用的是有随机性的采样，需要在每一张卡的程序实例初始化的时候，都进行相同的随机数指定。否则，不同的卡将可能会在Logits生成之后产生不同的采样结果，而整个sequence的生成Token是每个GPU上各自解码、各自保存的，每次Decode步之间并没有广播同步（为了性能），因此这将引发不同GPU间的序列不一致性，导致严重的问题。考虑到整体的设计模式，最好的指定是在初始化Engine实例的时候作为一个参数传入。因为用户不一定想到做这件事而且其具体影响不是很大，最好将其设置为一个可选的默认参数：比如42。

张量并行的实现手段

意外的简单，整个改造实现完全在pytorch内部进行。nccl相关的操作被torch.distribute完全封装，而且同样可以被CUDA Graph捕获，正常执行。

需要注意的技术细节

需要格外说明的是，CUDA Graph和torch.distribute+nccl backend会产生一些奇妙的化学反应。具体来说，CUDA Graph在capture的过程中，如果使用了nccl的后端作为torch.distribute的调用，那么CUDA Graph会直接固定的占用nccl的资源。这会导致在程序结束后，dist.destroy_process_group()无法正确完成，整个集群hang住，无法结束。为了解决这个问题，需要显式的删除CUDA Graph的对象：即执行del g, static_output类似的语句。这是一个说大不大但是很烦人的bug，花了一会儿才定位出来，一开始还以为是哪里同步错误，导致计算流本身有问题。

张量并行在Qwen2.5-1.5B上的效果

效果其实并没有预想中好，主要的原因是——虽然难以启齿——没钱。在2x5090，PCIe连接的的autodl平台上，总体来说加速比只有1.3x，效果并不是非常理想。通过Profiling排查原因可以看到，在有Profiling的情况下，通信（All-Reduce）占据了~48%的时间，几乎就是GPU time的一半了。这主要是因为PCIe连接的Latency问题：每次传输的内容其实只有几KB，数据量本身并不构成时间来源。但是，nccl+pcie的allreduce本身的overhead启动延迟就非常长，每次启动几乎需要26μs的时间，这在28层*2次调用的前向传播来说占据的时间长度几乎和真正的计算和内存存取一样多，从而大大拉低了潜力。这个问题没有什么好办法解决，是完全硬件上的。要解决倒也简单：NVLink，它的通信延迟比PCIe低很多。

画饼

想办法预定到了8卡B200的集群，真是太好了。现在是在autodl上做的，因为是PCIe互联，本身就不是非常适合TP这个并行模式，所以加速效果也有限，更何况只有2卡的5090。过几天马上就上去试试集群，看看Profiling的结果会不会有什么变化？能够预计，它能够达到的有效加速比应该比PCIe的情况高。