Tensor parallelism¶

A model too big for one GPU can be split across several. transformers shards each linear layer's weights; nnsight gathers the sharded activations so your block reads the whole tensor, and re-splits whatever you write before the model carries on.

Everything below was run on four A100s against Llama-3.3-70B-Instruct, transformers 5.16.1 and torch 2.9.1, and the outputs are the ones those runs produced. They are shown here as cells for readability, but you cannot run them in a notebook kernel. Tensor parallelism needs the calling process to be a rank, so the script goes into a file and the file is launched with torchrun.

In [ ]:

import torch
from transformers.distributed import DistributedConfig
from nnsight.modeling.transformers import TransformersModel

model = TransformersModel(
    "meta-llama/Llama-3.3-70B-Instruct",
    task="text-generation",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4),
    dispatch=True,
)

with model.trace("The Eiffel Tower is in the city of"):
    mlp = model.model.layers[40].mlp.output.save()
    logits = model.lm_head.output.save()

print(mlp.shape, logits.shape)
print(model.tokenizer.decode(logits[0, -1].argmax()))

import torch from transformers.distributed import DistributedConfig from nnsight.modeling.transformers import TransformersModel model = TransformersModel( "meta-llama/Llama-3.3-70B-Instruct", task="text-generation", dtype=torch.bfloat16, distributed_config=DistributedConfig(tp_size=4), dispatch=True, ) with model.trace("The Eiffel Tower is in the city of"): mlp = model.model.layers[40].mlp.output.save() logits = model.lm_head.output.save() print(mlp.shape, logits.shape) print(model.tokenizer.decode(logits[0, -1].argmax()))

torch.Size([1, 11, 8192]) torch.Size([1, 11, 128256])
 Paris

Run it with torchrun, not python — every rank runs the whole script, including your block:

torchrun --nproc_per_node=4 your_script.py

--nproc_per_node and tp_size have to be the same number. mlp and logits come back at their full width on every rank, exactly as they would on one GPU. Requires transformers >= 5.16.

A bare tp_plan="auto" works too and gets the same pre-load check — it names no degree of its own, so nnsight reads it off the world size, and a checkpoint that cannot be split is refused before any weights are fetched. Prefer distributed_config=DistributedConfig(tp_size=N): it states the degree in the code, and it is the only form that can ask for expert parallelism.

What is gathered for you¶

Every module's .input and .output. That covers most of what people read:

	Gathered?
Column-parallel output — q_proj, gate_proj	yes
Row-parallel input — o_proj.input, down_proj.input	yes
Row-parallel output, whole modules, the LM head	already whole
Parameters — q_proj.weight	no, see below
Values between two sharded modules	no, see below
embed_tokens.output	whole — unless the plan shards it, see below

The gather fires only when your block is actually waiting on that location, so reading a handful of places does not pay for the hundreds you ignore.

In [ ]:

layer = model.model.layers[40]

with model.trace("The Eiffel Tower is in the city of"):
    q_out      = layer.self_attn.q_proj.output.save()
    o_in       = layer.self_attn.o_proj.input.save()
    o_out      = layer.self_attn.o_proj.output.save()
    gate_out   = layer.mlp.gate_proj.output.save()
    down_in    = layer.mlp.down_proj.input.save()
    layer_out  = layer.output.save()
    head_out   = model.lm_head.output.save()

for name, value in [("q_proj.output", q_out), ("o_proj.input", o_in),
                    ("o_proj.output", o_out), ("gate_proj.output", gate_out),
                    ("down_proj.input", down_in), ("layers[40].output", layer_out),
                    ("lm_head.output", head_out)]:
    print(f"{name:20s} {tuple(value.shape)}")

layer = model.model.layers[40] with model.trace("The Eiffel Tower is in the city of"): q_out = layer.self_attn.q_proj.output.save() o_in = layer.self_attn.o_proj.input.save() o_out = layer.self_attn.o_proj.output.save() gate_out = layer.mlp.gate_proj.output.save() down_in = layer.mlp.down_proj.input.save() layer_out = layer.output.save() head_out = model.lm_head.output.save() for name, value in [("q_proj.output", q_out), ("o_proj.input", o_in), ("o_proj.output", o_out), ("gate_proj.output", gate_out), ("down_proj.input", down_in), ("layers[40].output", layer_out), ("lm_head.output", head_out)]: print(f"{name:20s} {tuple(value.shape)}")

q_proj.output        (1, 11, 8192)
o_proj.input         (1, 11, 8192)
o_proj.output        (1, 11, 8192)
gate_proj.output     (1, 11, 28672)
down_proj.input      (1, 11, 28672)
layers[40].output    (1, 11, 8192)
lm_head.output       (1, 11, 128256)

hidden_size is 8192 and intermediate_size is 28672 on this checkpoint, so every one of those is the full width — not the 2048 and 7168 each rank actually computed.

Three obligations on your block¶

Every rank runs it, which is what keeps the collectives lined up.

1. No rank-dependent control flow. If the ranks take different paths they stop agreeing on when to gather, and the run deadlocks. There is no exception, and no watchdog fires for minutes. Killing torchrun does not take the rank processes with it either; find their pids and kill those, or the cards stay occupied by the run you thought you ended.
2. Seed before you sample. torch.initial_seed() differs per rank under torchrun. If sampling diverges the ranks generate different tokens, and the model's own all-reduces then sum activations from different sequences, so the output is wrong on every rank, not merely different. Use greedy decoding, or torch.manual_seed(0) identically everywhere. Many checkpoints ship do_sample: true in generation_config.json, so this bites without you asking for sampling.
3. Clone before you edit. A gathered value is the output of a collective, and torch refuses an in-place write into one:

In [ ]:

with model.trace("The Eiffel Tower is in the city of"):
    layer.mlp.gate_proj.output[..., :1024] = 0

with model.trace("The Eiffel Tower is in the city of"): layer.mlp.gate_proj.output[..., :1024] = 0

RuntimeError: Output 0 of SliceBackward0 is a view and is being modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is forbidden. You can fix this by cloning the output of the custom Function.

Clone, edit, assign back:

In [ ]:

with model.trace("The Eiffel Tower is in the city of"):
    edited = layer.mlp.gate_proj.output.clone()
    edited[..., :1024] = 0
    layer.mlp.gate_proj.output = edited
    ablated = model.lm_head.output.save()

print(model.tokenizer.decode(ablated[0, -1].argmax()))

with model.trace("The Eiffel Tower is in the city of"): edited = layer.mlp.gate_proj.output.clone() edited[..., :1024] = 0 layer.mlp.gate_proj.output = edited ablated = model.lm_head.output.save() print(model.tokenizer.decode(ablated[0, -1].argmax()))

 Paris

Replacing the whole value (... .output = torch.zeros_like(...)) needs no clone. nnsight does not clone for you: on a model this size that is a copy of the activation on every gather, for the many traces that only read.

Reading a value between two sharded modules¶

This is the one place a trace does not read as it would on one GPU.

The shard is created by a column-parallel module and consumed by the row-parallel one after it, so everything in between is one rank's slice:

gate_proj [colwise] ─▶  (1,11,64)  →  (1,11,32)   ┐
act_fn                  (1,11,64)  →  (1,11,32)   │  every value here
up_proj   [colwise]     (1,11,64)  →  (1,11,32)   │  is this rank's slice
down_proj [rowwise] ─▶  (1,11,16)  →  (1,11,16)   ┘  ← whole again

nnsight does not gather these for you, because nothing on the value says how it is split. transformers tracks the layout only while a value is inside the module that produced it; on the way out it unwraps to an ordinary tensor, and a 32-wide slice of a 64-wide activation is then indistinguishable from a genuine 32-wide one.

And the axis moves. Inside attention, view/transpose puts the shard on the head dimension:

q_proj [colwise]      (1,11,16)  →  (1,11,8)     sharded on the last dim
  .view(...)          (1,11,4,4) →  (1,11,2,4)   sharded on dim 2
  .transpose(1,2)     (1,4,11,4) →  (1,2,11,4)   sharded on dim 1  ← heads
o_proj [rowwise]      (1,11,16)  →  (1,11,16)    whole again

So there is not even a fixed axis to assume. You know what the forward did, so you say which axis:

In [ ]:

from nnsight.modeling.tp import gather, shard

attn = model.model.layers[1].self_attn

with model.trace("The Eiffel Tower is in the city of"):
    q = attn.source.query_states_0.output       # (1, heads/tp_size, seq, head_dim)
    local = torch.tensor(q.shape).save()
    heads = gather(model, q, dim=1).save()      # (1, heads, seq, head_dim) — every head

print("this rank's slice:", tuple(local.tolist()))
print("gathered:         ", tuple(heads.shape))

from nnsight.modeling.tp import gather, shard attn = model.model.layers[1].self_attn with model.trace("The Eiffel Tower is in the city of"): q = attn.source.query_states_0.output # (1, heads/tp_size, seq, head_dim) local = torch.tensor(q.shape).save() heads = gather(model, q, dim=1).save() # (1, heads, seq, head_dim) — every head print("this rank's slice:", tuple(local.tolist())) print("gathered: ", tuple(heads.shape))

this rank's slice: (1, 16, 11, 128)
gathered:          (1, 64, 11, 128)

64 heads, 16 per rank at tp_size=4. To edit one, put this rank's piece back before the forward continues:

In [ ]:

with model.trace("The Eiffel Tower is in the city of"):
    q = attn.source.query_states_0.output
    heads = gather(model, q, dim=1).clone()
    heads[:, 3] = 0                                  # ablate head 3, whoever holds it
    attn.source.query_states_0.output = shard(model, heads, dim=1)
    logits = model.lm_head.output.save()

print(model.tokenizer.decode(logits[0, -1].argmax()))

with model.trace("The Eiffel Tower is in the city of"): q = attn.source.query_states_0.output heads = gather(model, q, dim=1).clone() heads[:, 3] = 0 # ablate head 3, whoever holds it attn.source.query_states_0.output = shard(model, heads, dim=1) logits = model.lm_head.output.save() print(model.tokenizer.decode(logits[0, -1].argmax()))

 Paris

gather and shard are collectives, so obligation 1 applies with full force: every rank must reach them, so call them unconditionally and never inside a branch that could go differently on different ranks. Both are no-ops on an unsharded model, so the same script runs on one GPU and on eight.

To find the axis, print the shape at tp_size=1 and at tp_size=2. The one that shrank by tp_size is the one to pass.

When width does not tell you¶

That comparison works on a plain tensor, which is what a value between two modules is. It does not work on a .source value taken from inside a sharded module. Those come back as DTensors, and a DTensor reports the global shape while holding this rank's data. So nothing shrinks, at any degree, and reading a number out of it gives the wrong answer with no warning:

In [ ]:

with model.trace("The Eiffel Tower is in the city of"):
    inner = layer.mlp.gate_proj.source.F_linear_0.output.save()

print("type        ", type(inner).__name__)
print("shape       ", tuple(inner.shape), "<- the global shape")
print("to_local()  ", tuple(inner.to_local().shape), "<- what this rank holds")
print("sum()       ", float(inner.sum()))
print("full_tensor()", float(inner.full_tensor().sum()))

with model.trace("The Eiffel Tower is in the city of"): inner = layer.mlp.gate_proj.source.F_linear_0.output.save() print("type ", type(inner).__name__) print("shape ", tuple(inner.shape), "<- the global shape") print("to_local() ", tuple(inner.to_local().shape), "<- what this rank holds") print("sum() ", float(inner.sum())) print("full_tensor()", float(inner.full_tensor().sum()))

type         DTensor
shape        (1, 11, 28672) <- the global shape
to_local()   (1, 11, 7168) <- what this rank holds
sum()        -4736.0                       # rank 0; rank 1 says -5024.0
full_tensor() -19200.0                     # the same on every rank

So the test for one of these is value.placements, not its width. gather handles it (the value carries its own layout, so gather uses that rather than the dim you name), and so does .full_tensor().

Parameters¶

layer.weight is a DTensor: this rank holds 1/tp_size of it, but .shape reports the whole, so the shape alone will not tell you it is split.

In [ ]:

w = model.model.layers[40].self_attn.q_proj.weight

print("type        ", type(w).__name__)
print("w.shape     ", tuple(w.shape))
print("w.placements", w.placements)
print("w.to_local()", tuple(w.to_local().shape))

w = model.model.layers[40].self_attn.q_proj.weight print("type ", type(w).__name__) print("w.shape ", tuple(w.shape)) print("w.placements", w.placements) print("w.to_local()", tuple(w.to_local().shape))

type         DTensor
w.shape      (8192, 8192)
w.placements (Shard(dim=0),)
w.to_local() (2048, 8192)

nnsight does not reassemble weights for you: they are what tensor parallelism exists to split, and gathering one would allocate the whole tensor on every rank in exactly the situation where memory was tight enough to reach for TP.

Reducing one gives you this rank's answer, and nothing says so. w.mean(), w.norm(), w.abs().max() all come back as a DTensor with a Partial placement — the reduction over this rank's slice, still waiting to be combined. float() reads that partial number, which differs between ranks and matches none of them:

In [ ]:

print(f"float(w.mean())               = {float(w.mean()):.6e}")
print(f"float(w.mean().full_tensor()) = {float(w.mean().full_tensor()):.6e}")
print(f"float(w.norm())               = {float(w.norm()):.4f}")
print(f"float(w.norm().full_tensor()) = {float(w.norm().full_tensor()):.4f}")

print(f"float(w.mean()) = {float(w.mean()):.6e}") print(f"float(w.mean().full_tensor()) = {float(w.mean().full_tensor()):.6e}") print(f"float(w.norm()) = {float(w.norm()):.4f}") print(f"float(w.norm().full_tensor()) = {float(w.norm().full_tensor()):.4f}")

# rank 0                                   # rank 1
float(w.mean())               = -8.702278e-06    5.811453e-06
float(w.mean().full_tensor()) = 3.576279e-07     3.576279e-07
float(w.norm())               = 62.7500          49.7500
float(w.norm().full_tensor()) = 117.5000         117.5000

Call .full_tensor() on the reduction, not on the weight: the result is a scalar, so it costs one small collective rather than a copy of the layer. Any weight-norm sweep or layer-magnitude plot needs this, or it plots per-rank noise. (w.cpu() is still a DTensor, so moving it off the GPU does not help.)

A sharded weight also cannot be edited in place through the DTensor. weight[:, :10] = 0 raises NotImplementedError: Operator aten.fill_.Tensor does not have a sharding strategy registered. Write through .to_local(), which is the right thing anyway: each rank edits the rows or columns it holds.

with torch.no_grad():
    layer.mlp.gate_proj.weight.to_local()[:, :10] = 0

The embedding¶

Some plans shard the embedding and some do not, and it makes a difference to what you can read. A plan that shards it names it embedding_rowwise, which in practice means the tied-embedding checkpoints, Llama-3.2-1B and -3B among them. Where the embedding is sharded, its value's layout is single-use, and parking on embed_tokens.output raises a bare AssertionError from torch's embedding op. The same takes out a tracer.cache() called with no arguments, which selects every module and so reaches the embedding.

"embed_tokens" in (model.config.base_model_tp_plan or {})
# True  — Llama-3.2-1B, Llama-3.2-3B     embed_tokens.output raises
# False — Llama-3.3-70B, Qwen3-8B        embed_tokens.output reads whole

model.model.layers[0].input is the same tensor one module later, and it is whole on either plan.

Before you allocate the GPUs¶

tp_size has to divide the model's attention heads, key/value heads and intermediate size. max_tp_size answers from the config alone: no weights, no GPUs, no torchrun.

In [ ]:

from transformers import AutoConfig
from nnsight.modeling.tp import max_tp_size

for name in ["meta-llama/Llama-3.3-70B-Instruct", "meta-llama/Llama-3.2-3B",
             "Qwen/Qwen2.5-0.5B", "openai-community/gpt2"]:
    print(f"{name:36s} {max_tp_size(AutoConfig.from_pretrained(name))}")

from transformers import AutoConfig from nnsight.modeling.tp import max_tp_size for name in ["meta-llama/Llama-3.3-70B-Instruct", "meta-llama/Llama-3.2-3B", "Qwen/Qwen2.5-0.5B", "openai-community/gpt2"]: print(f"{name:36s} {max_tp_size(AutoConfig.from_pretrained(name))}")

meta-llama/Llama-3.3-70B-Instruct    8
meta-llama/Llama-3.2-3B              8
Qwen/Qwen2.5-0.5B                    2
openai-community/gpt2                None

Every workable degree is a divisor of that number. Qwen2.5-0.5B stops at 2 because it has 2 key/value heads. None means the checkpoint cannot be split at all — gpt2 publishes no base_model_tp_plan, and neither do several recent Qwen releases, so this is not only an old-model problem. Asking anyway raises UnshardableCheckpoint, listing the degrees that would have worked. The check sits on the loading path, so with the default dispatch=False it arrives at .dispatch() rather than at the constructor.

Expert parallelism¶

A mixture-of-experts checkpoint can be split by expert instead. transformers applies the model's expert plan in place of its tensor-parallel one:

model = TransformersModel(
    "Qwen/Qwen1.5-MoE-A2.7B",
    task="text-generation",
    distributed_config=DistributedConfig(tp_size=2, enable_expert_parallel=True),
    dispatch=True,
)

Traces read the same. mlp.experts.output and mlp.gate.output come back at their single-GPU shapes and values, and edits to them carry through. The degree has to divide the expert count rather than the head counts; max_tp_size(config, expert_parallel=True) is the pre-flight for that.

Precision¶

An all-reduce sums in a different order than one big matmul, so results drift a little from a single-GPU run. How much depends entirely on the dtype. Llama-3.2-3B at tp_size=4 against one GPU, as the largest elementwise difference over the tensor's own scale:

	layer 0 gate_proj	layer 13	layer 27	logits
bfloat16	3.4e-3	9.0e-3	7.5e-3	8.2e-3
float32	2.5e-7	7.3e-7	9.7e-7	9.1e-7

Greedy token choices are identical in both. If you are using agreement with a single-GPU run as a sanity check, expect the bfloat16 numbers — a float32 mismatch above 1e-6 is a real difference, not arithmetic order.

Related¶

• Intermediate Operations — .source, used above to reach query_states
• vLLM tensor parallelism — the same idea in the vLLM runtime, where gathering is automatic everywhere