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, in one of two ways, because it has
used both. Up to 5.15 apply_tensor_parallelism stamped every module it
sharded with _hf_tp_plan (the style name) and _hf_device_mesh. 5.16
rebuilt tensor parallelism on DTensor and stamps nothing: the plan stays on the
model as _tp_plan — glob patterns over module paths — with the mesh at
_device_mesh, and each module is resolved against it by transformers' own
matcher. Either way 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.
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 is arranged differently per
backend. 5.15 registered them as torch pre/post hooks, which bracket a forward on
their own. 5.16 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. On the DTensor backend 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'}}
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:
with model.trace(prompt):
q = layer.self_attn.source.query_states_0.output # (1, heads/N, seq, dim)
whole = tp.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.
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.
with model.trace(prompt):
q = layer.self_attn.source.query_states_0.output
whole = tp.gather(model, q, dim=1)
whole[:, 3] = 0 # ablate head 3
layer.self_attn.source.query_states_0.output = tp.shard(model, whole, dim=1)
Same rule: a collective, so every rank must reach it.