tensor_parallel

收录理由

当PyTorch模型超出单卡显存时，通常得把训练循环改写成分布式框架的写法，这套库则绕开了这一步：用张量并行把模型包一层，各层权重自动切到多张GPU上，前向和反向传播都能正常跑，输出也会自动拼回来。这样一来，微调大号Transformer的代码几乎不用动，尤其适合单机多卡或笔记本环境，DeepSpeed、Megatron这类重框架反而显得杀鸡用牛刀。它还带了显存友好的调度辅助，以及把切分后的模型合并回普通单卡检查点的功能，换机器部署微调结果时省去不少麻烦。如果你本来就在用Hugging Face Transformers，只是想把13B级别的模型塞进两张或更多卡里，又不想引入新框架，这个工具上手成本很低。