Tensor parallelism

Tensor parallelism (TP) shards a model's weights across GPUs along a tensor dimension so a single forward/backward pass uses all GPUs simultaneously. The library's TP integration was rewritten in PR #36539 (March 2025) and now lives at src/transformers/integrations/tensor_parallel.py (66K LOC).

How it works

For a linear layer y = x W, two sharding choices exist:

• Column-parallel — each rank holds a column slice of W. The output y is split along the last dimension; an all-gather (or no-op when followed by a row-parallel layer) reassembles it.
• Row-parallel — each rank holds a row slice of W. The input x must be sharded along the last dimension; the output is reduced via all-reduce.

Stacking column-parallel followed by row-parallel (with no inter-layer all-gather) is the standard pattern used in transformer attention and MLP blocks (Megatron-LM style).

tp_plan

A model declares its sharding strategy as a dict:

class LlamaModel(LlamaPreTrainedModel):
    _tp_plan = {
        "self_attn.q_proj": "colwise",
        "self_attn.k_proj": "colwise",
        "self_attn.v_proj": "colwise",
        "self_attn.o_proj": "rowwise",
        "mlp.gate_proj": "colwise",
        "mlp.up_proj": "colwise",
        "mlp.down_proj": "rowwise",
        "embed_tokens": "rowwise",
        "lm_head": "colwise",
    }

from_pretrained(..., tp_plan="auto") reads the plan, shards weights across the world group during loading, and replaces forward calls with the sharded variants.

The implementation in src/transformers/integrations/tensor_parallel.py provides:

• colwise, rowwise, gather, replicate, local_colwise, local_rowwise, local_packed_rowwise strategies.
• Distributed weight loading helpers compatible with safetensors.
• Hooks that insert all-reduce / all-gather ops in the right places.
• A registry of strategies so models can declare their own (e.g., MoE-aware sharding).

Activating TP

torchrun --standalone --nproc_per_node=8 my_script.py

import os, torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-70B",
    dtype=torch.bfloat16,
    tp_plan="auto",
)

Trainer users can also set TrainingArguments.tp_size > 1.

Combining with other parallelisms

• TP + DDP — DDP across nodes, TP within a node.
• TP + FSDP — FSDP shards the optimizer state across data-parallel ranks; TP shards within each FSDP unit.
• TP + EP (expert parallel) — for MoE models, src/transformers/integrations/moe.py shards experts across a separate group.
• TP + PP — pipeline parallel is supported via the accelerate and deepspeed paths.

The examples/3D_parallel.py script shows a 3D combination.

Models that support tp_plan

Models with an explicit _tp_plan ship sharded out of the box. As of v5 most decoder LMs (Llama family, Qwen family, Mistral family, Gemma family, Phi family, DeepSeek-V2/V3, GLM, Granite, Falcon, Mixtral, Olmo, Cohere, …) and many VLMs are TP-ready. Models without _tp_plan either need adoption (open contribution) or fall back to non-TP loading.

Test infrastructure

• tests/tensor_parallel/ contains correctness and load tests.
• tests/test_tensor_parallel_mixin.py (23K LOC) is the shared mixin that per-model tests inherit when they declare is_tensor_parallel_test.
• The flash_attn_test, is_tensor_parallel_test, and is_training_test markers in pyproject.toml allow CI to filter.

Continuous batching + TP

Continuous batching uses the same TP machinery; the paged attention kernels are TP-aware. transformers serve --continuous-batching works on multi-GPU when launched with torchrun.

Integration points

• Modeling — _tp_plan, _no_split_modules.
• from_pretrained — sharded weight loading.
• Trainer — tp_size argument.
• Continuous batching — paged + TP combo.
• Integrations — tensor_parallel.py, moe.py.

Entry points for modification

• New sharding strategy → register a strategy in src/transformers/integrations/tensor_parallel.py and document in the docstring.
• Add TP to a model → declare _tp_plan on the relevant <Arch>Model (or <Arch>PreTrainedModel) class and add a test in tests/models/<arch>/test_modeling_<arch>.py that opts into TensorParallelTesterMixin.
• Combined TP+EP for MoE → see src/transformers/integrations/moe.py and existing examples in DeepSeek-V3 / Qwen2-MoE.