envoys¶
Envoys for modules transformers split across ranks.
The ad-hoc-call half of what
TPFragments does for activations.
Interleaving already shows a worker whole activations — gathered on the way in,
re-split on the way out, once per visit — but a trace that calls a module
directly, away from its place in the forward pass, is outside that bracket. A
logit lens running lm_head on an intermediate hidden state is holding, and
wants back, whole tensors, while the sharded module deals in slices.
So the bracket is the same one vLLM's envoy uses
(nnsight.modeling.vllm.envoys): re-split the
caller's whole input, run the module, reassemble the output, all off
TPFragments' own rules.
What is not needed here is any handling of the collectives themselves.
Envoy.__call__ runs the module the
ordinary way — standing the interleaver down rather than dodging __call__ —
so the style's own transforms fire around the call, which is what makes a
row-parallel layer's all-reduce and colwise_gather_output's all-gather happen
without this module restating them.
It does mean the style's input transform runs too, which is why re-splitting
the input is conditional: see
SPLITS_ITS_OWN_INPUT.
Parameters are left alone. layer.weight is the DTensor transformers made
of it — this rank holds a slice, while .shape reports the whole — as it is
anywhere else under transformers tensor parallelism.
module-attribute
¶
module-attribute
¶
TPEnvoy(module: Module, path: str = 'model', interleaver: Interleaver | None = None, rename: dict[str, str | list[str]] | None = None, envoys: dict | None = None)
Bases: Envoy
An envoy over a module transformers may have sharded.
Inert on an unsharded model and on a one-rank mesh — TPFragments records
a rule only for a module actually split — so this is a safe envoy class for
every Linear and Embedding in a tree that was loaded across ranks.
Run this module ad hoc, on whole tensors either side.
Every rank runs the block, so every rank reaches the same collectives in the same order — as long as the call is not under rank-dependent control flow, the condition every collective in a block carries.
hook does not change this. It says whether the trace watches the
call, which is a separate question from whether the caller is holding
whole tensors; and the handoff it turns on brackets the value and puts it
back, so it leaves this caller's return value exactly as it found it. When
hook short-circuited the bracket, an observer at .output saw the
whole tensor while the caller of the same call got this rank's slice.
The envoys map pairing shardable module types with TPEnvoy.
Keyed by type rather than by style because a style belongs to the instance
— the same nn.Linear class is colwise in one place and rowwise in
another — while the map is consulted per module as the tree is built.
TPEnvoy asks TPFragments which style its own path got.
Whether this construction is going to produce a sharded model.
Asked from a constructor, before the model exists — so for a repo id it reads
the request (distributed_config) rather than the result. A ready module
is already there and can be asked directly.
Erring towards True is the safe direction: TPEnvoy is inert on a module
with no style stamped on it, so a false positive costs a different envoy
class and nothing else, while a false negative silently leaves ad-hoc calls
handing back slices. Hence tp_plan counts even with no tp_size.
Note this still asks a wider question than
HuggingFaceModel._refuse_impossible_tp: a custom tp_plan dict counts
here (the model may come back sharded) but is not degree-checked there (it
overrides the published plan the check would read).