tensor_parallel

Why we included this project

Running a PyTorch model that outgrows one GPU usually means rewriting your training loop around a distributed framework. This library sidesteps that: wrap the model with tensor parallelism and it splits each layer's weights across your GPUs, then reassembles the outputs, working for both inference and backpropagation. That keeps fine-tuning a large transformer as straightforward as before, and it shines on a single multi-GPU machine or in a notebook where DeepSpeed or Megatron feel like overkill. It also includes helpers for memory-efficient dispatch and for collapsing a sharded model back into a normal single-GPU checkpoint, which saves you trouble when moving a fine-tuned model elsewhere. If you already work with Hugging Face transformers and just need to fit a 13B-class model on two or more GPUs without adopting a new framework, this is a low-friction option.