fragments¶
Show intervention code whole tensors on a tensor-parallel model.
Under transformers tensor parallelism a module's activation can be one rank's slice rather than the real thing: a column-parallel linear splits its output across ranks, and a row-parallel linear takes its input already split. A user asked for the layer, not a quarter of it, so those values are gathered before a worker sees them and re-split before the model's own forward carries on.
Two facts make this cheap to arrange.
transformers labels the shards for us. The plan lives on the model — glob
patterns over module paths in tp_plan/ep_plan, with the mesh at
_device_mesh — and each module is resolved against it by transformers' own
matcher, so which side of which module carries a shard is read off the model,
not guessed. These rules are handed every envoy through instrument as the
tree is built, so they record themselves — and learn whether there is anything
to do at all — right there. (A transformers old enough to stamp each module
with _hf_tp_plan instead, or newer than the floor's shapes in other ways,
is refused — see MINIMUM_TRANSFORMERS below.)
nnsight sees the pre-collective value. The interleaver's handoff runs inside
a module's forward — after the style's input transform, before its output one —
so a row-parallel output is still this rank's partial sum there, and a
colwise_gather_output head still a shard.
SIDES says which, per style, and
TPFragments.whole
all-gathers a shard or all-reduces a partial. What goes back is chosen so the
module's own output transform completes the picture: a shard's slice, or — for a
partial — the whole on rank 0 and zeros elsewhere, which its reduce turns into
exactly the (possibly edited) whole on every rank.
Keeping the handoff inside those transforms takes one arrangement:
transformers applies a style by replacing module.forward with a wrapper —
the same slot nnsight's controller uses — so _keep_tp_forward installs the
controller first and re-wraps it, restoring that order. Without it the module
keeps its sharded weights but loses the code that makes its inputs match them,
and the first matmul dies with got mixed torch.Tensor and DTensor.
The rules describe module boundaries. A .source value between two ops
inside a forward can be a shard split on an axis that moves through the forward,
so a boundary rule cannot say what it is. Such a value
often carries its own placement — full_tensor then reassembles it correctly
whatever axis holds the shard, which is why the gather believes a value's own
placement over the rule whenever it has one. When it does not (transformers takes
a local fast path for nn.Linear with grad disabled, and for quantized
modules) the value is handed over as-is, as it always was: compare against a
single-GPU run if it matters, and never branch on one — the ranks would diverge.
Every rank runs the same intervention block, so every rank reaches handle at
the same location with the same parked workers and makes the same decision to
gather — which is what keeps the collectives matched, why a block whose control
flow diverges across ranks deadlocks, and why sampling must be seeded identically
on every rank.
Everything about when to gather — once per visit, only when something is
waiting, re-split on the way out — belongs to
Interleaver.handle and is
shared with every other distributed runtime. See
nnsight.intervention.fragments for that half,
including why it is the interleaver's job and not the batcher's.
module-attribute
¶
SIDES: Dict[str, Dict[str, str]] = {'colwise': {'output': 'shard'}, 'packed_colwise': {'output': 'shard'}, 'colwise_gather_output': {'output': 'shard'}, 'colwise_rep': {'output': 'shard'}, 'rowwise': {'input': 'shard', 'output': 'partial'}, 'rowwise_split_input': {'input': 'shard', 'output': 'partial'}, 'rowwise_rep': {'input': 'shard', 'output': 'partial'}, 'packed_rowwise': {'input': 'shard', 'output': 'partial'}, 'embedding_rowwise': {'output': 'partial'}, 'sequence_parallel': {'output': 'partial'}, 'all_reduce': {'output': 'partial'}, 'replicated_with_grad_allreduce': {}, 'ep_router': {}, 'grouped_gemm': {}, 'moe_tp_experts': {'output': 'partial'}, 'ep_dispatch_experts': {}}
module-attribute
¶
UNSUPPORTED: Dict[str, str] = {'megamoe_router': 'expert-parallel (MoE)', 'megamoe_experts': 'expert-parallel (MoE)', 'moe_identity_expert': 'expert-parallel (MoE)', 'mla_kv_a_proj': 'MLA split kv projection'}
Bases: RuntimeError
transformers is too old to shard a model correctly.
Bases: Exception
The model shards something interventions can't be shown whole.
Bases: Fragments
Which values a transformers-sharded model splits, and how to reassemble them.
Built for every HuggingFace model and inert (enabled=False) until
instrument finds a module actually split across ranks.
| ATTRIBUTE | DESCRIPTION |
|---|---|
enabled |
Whether anything in this tree is sharded.
|
tp_rules |
Location ->
TYPE:
|
Record what each side of this envoy's module is at the handoff.
Called as the tree is built and again on dispatch (Envoy._update), which
is when a module first carries the marks a sharded model is recognized by.
| RAISES | DESCRIPTION |
|---|---|
UnsupportedParallelStyle
|
for a style there is no rule for — refused up front rather than silently handing users a fragment. |
The parallel style and mesh recorded for the module at path.
(None, None) for a module this tree did not find sharded. Asked by
nnsight.modeling.tp.envoys.TPEnvoy, which cannot read the style off the
module: transformers keeps the plan on the model, not on each module it
shards.
Whether this location's value is one rank's piece.
A dict lookup: the rules were recorded at instrument time, so nothing is inspected here and nothing branches on rank.
Only a module's own two sides have rules. A value inside a forward — a
.source location, or any module between a column-parallel output and
the row-parallel input that consumes it — has none, and is handed over as
it comes. Nothing records which axis holds its shard once it has left the
module that made it, and the axis moves: attention's view/transpose
puts it on the head dimension. Reassembling one is the trace's job, with
gather and
shard, which take the axis from
the caller because only the caller knows it.
The real tensor, and how to put back what assembling it took.
The value's own placement wins when it has one — it knows which axis holds
the shard, which a rule cannot once a view or transpose has moved
it — and the location's rule decides otherwise. Only a module's two sides
reach here; see fragmented for what is left raw and why.
The way back is returned as a closure over what was decided here, so it cannot be confused with another location's, or consumed by an ad-hoc call made while this visit is still open.
This rank's piece of a value that was never gathered.
The rule alone, because there is nothing else: a .skip replacement and
the argument of an ad-hoc call are both the caller's own whole tensor, and
neither ever carried a placement to read.
The mesh model was sharded over, or None if it was not sharded.
Takes the model wrapper, the envoy, or the bare module — whichever is to hand inside a trace.
Every rank's piece of value, concatenated along dim.
For a value nnsight hands over as-is — anything between a column-parallel module's output and the row-parallel module that consumes it, where nothing records which axis holds the shard. You know, because you know what the forward did, so you say:
from nnsight.modeling.tp import gather
with model.trace(prompt):
q = layer.self_attn.source.query_states_0.output # (1, heads/N, seq, dim)
whole = gather(model, q, dim=1) # (1, heads, seq, dim)
A collective, so every rank must reach it: call it unconditionally, never
inside a branch that could go differently on different ranks. Returns the
value unchanged on an unsharded model, so the same block runs either way —
which also means a one-GPU run never checks dim: a wrong axis only
shows itself at tp_size > 1.
This rank's piece of value along dim — the inverse of gather.
Needed when you write an edited value back into an intermediate location: the model's forward carries on expecting this rank's piece, so a whole tensor left there is as wrong as a piece read out.
from nnsight.modeling.tp import gather, shard
with model.trace(prompt):
q = layer.self_attn.source.query_states_0.output
whole = gather(model, q, dim=1).clone() # clone before editing:
whole[:, 3] = 0 # a gathered value is a
layer.self_attn.source.query_states_0.output = ( # collective's output,
shard(model, whole, dim=1) # in-place writes raise
)
Same rule: a collective, so every rank must reach it.