tp¶
Tracing a model sharded with transformers tensor parallelism.
Nothing here needs installing or enabling. A
HuggingFaceModel is always
built with a TPFragments, which
stays inert unless it finds the model actually sharded. See
fragments for which values are pieces and how
they are reassembled, and
nnsight.intervention.fragments for when.
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'}
module-attribute
¶
__all__ = ['MINIMUM_TRANSFORMERS', 'device_mesh', 'gather', 'shard', 'SIDES', 'UNSUPPORTED', 'TPFragments', 'UnsupportedParallelStyle', 'UnsupportedTransformersVersion', 'UnshardableCheckpoint', 'check_tp_request', 'max_tp_size', 'requested_expert_parallel', 'requested_tp_size']
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.
Bases: Exception
The model shards something interventions can't be shown whole.
Bases: RuntimeError
transformers is too old to shard a model correctly.
Bases: ValueError
A tensor-parallel degree was asked for that this checkpoint cannot serve.
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.
Raise unless tp_size is a degree config's model can really be split into.
transformers does not check this, and its two failure shapes are both worth
refusing. Asked to shard a checkpoint with no plan it shards nothing:
verify_tp_plan returns early on a None plan and
apply_tensor_parallelism installs no hooks, so every rank quietly loads a
complete copy of the weights — nothing errors, nothing warns, and the only
symptom is n times the memory for one model's worth of work. Asked for a
degree the plan cannot divide (SmolLM2's 9 heads at 2 ranks), it loads the
checkpoint sharded anyway — DTensor splits the 576 q_proj columns evenly,
4.5 heads per rank — and the first forward dies on
RuntimeError: shape '[1, 9, -1, 64]' is invalid for input of size 2592,
a reshape of the local tensor by the global head count, naming nothing about
tensor parallelism. Silent waste or a late opaque crash; the refusal here is
early and says what to do.
That is worth refusing rather than reporting, because there is no reading of "shard this over 4 GPUs" that is served by putting the whole thing on each of them. Raising here also puts the failure before the weights are fetched, where the message can still say what to do about it.
| RAISES | DESCRIPTION |
|---|---|
UnshardableCheckpoint
|
if the model cannot be split at all, or not into
exactly |
The largest tensor-parallel degree config's model supports.
None when it cannot be split at all: no plan to shard by, or a plan
containing a style TPFragments.instrument
will refuse. Refusing here keeps a model that would fail at load from being
placed as though it could be split.
The two must refuse the same set, which is why this asks
SIDES rather than
only UNSUPPORTED. A style in
neither — one transformers added, or one it never registered in
ALL_PARALLEL_STYLES — used to pass here and raise there, so a server
would allocate the cards, load the weights across them, and only then find
out. Llama4Config does exactly this: its plan is colwise_rep, which
is in no list and no registry.
Whether this distributed_config asks for expert parallelism.
Accepts the dataclass or a plain dict, because transformers does.
The degree a distributed_config asks for, or None if it asks for none.
Accepts the dataclass or a plain dict, because transformers does.