batching¶
Batch several invoke inputs into one forward and scope interventions to rows.
with model.trace() as tracer: may contain several with tracer.invoke(x):
blocks. Their inputs are combined into a single batched forward, and each block's
interventions see only its rows of every activation.
A Batcher (one per trace) collects each invoke's input and assigns it a
batch_group — a [start, size] row range in the combined batch. At run time
Batcher.narrow slices a full batched activation down to a block's rows when
it reads, and Batcher.widen splices an edit back into the full tensor. The row math
is dim-0 only; the model's _batch equalizes everything else (e.g. sequence
length) when it builds the combined input.
Two consequences a block's author meets. Equalizing the sequence length pads every invoke out to the batch's longest input, so a position index counted from the left names a different token depending on what else is in the batch, while one counted from the right does not.
Batching only actually narrows when there are two or more non-empty invokes — a lone invoke is the whole batch, so it sees every row untouched, and neither the row scoping nor that row check applies to it.
Skip replacements collected across the invokes of one batched forward.
.skip() bypasses a module's body and substitutes a value for its output.
In a batched forward there is no body output to splice into — the body didn't
run — so the combined output is built from the invokes' replacements alone
(see Batcher.gather_skip / Batcher.assemble_skip). Each entry
is a (group, replacement) pair.
Collects invoke inputs for one trace and builds the combined forward input.
Each add records an invoke's input and returns its batch_group.
assemble hands the collected invokes to the model's _batch to
produce the actual (args, kwargs) for the run. narrow/widen
scope a batched activation to a group's rows and splice an edit back; a model
whose batch layout isn't a plain dim-0 stack overrides the per-tensor
_narrow_tensor/_widen_tensor (e.g. diffusion's
classifier-free-guidance doubling) and picks its subclass via
_batcher_class.
property
¶
Whether narrowing applies — true once two or more invokes contribute rows.
narrow(value: Any, group: BatchGroup) -> Any
Slice every batched tensor in value down to group's rows.
A tensor is batched only when its leading dim equals total (the
combined batch size), so non-batched tensors pass through untouched. Returns
the whole value when not actually batching or for a groupless (empty) invoke.
widen(full: Any, group: BatchGroup, edited: Any) -> Any
Splice edited (a block's rows) back into full (the whole batch).
Walks full and edited in parallel; for each batched tensor in
full writes the corresponding edited tensor into the group's rows via
_widen_tensor. Returns edited unchanged when not batching or for
a groupless invoke.
gather_skip(running: Any, group: BatchGroup, replacement: Any) -> Any
Collect one invoke's skip replacement for its group's rows.
A lone invoke is the whole batch, so its replacement is the output
outright. With two or more, there's no body output to splice into — the
skip fires before the body runs — so accumulate the replacements and let
assemble_skip build the combined output once every invoke's is in.
Concatenate collected skip replacements into the full-batch output.
A no-op unless running is the SkipParts a batched skip built.
The replacements must tile the whole batch: every invoke has to skip the
module, since a shared forward can't run for only the rows that didn't.
add(*inputs: Any, **kwargs: Any) -> BatchGroup
Record one added input set; return its [start, size] row group.
A set that contributes no rows — params only (e.g. max_new_tokens=) or an
empty invoke() — returns None (a groupless, whole-batch worker) and
folds its kwargs into extra_kwargs, so assemble lays them onto
the combined call.
Build the combined (args, kwargs) for fn from the collected input
sets; params-only kwargs (extra_kwargs) are laid on top and win.
Concatenate a list of like-shaped structures along dim 0, leaf by leaf.
Every structure is one invoke's skip replacement, so they share a shape:
tensors concatenate, containers recurse in parallel, and a non-tensor leaf
(a None in a layer's output tuple, say) is taken from the first — they
agree by construction.