plan¶
How many ways a checkpoint's weights can be split, read from its config.
Answered before anything is loaded, because a server placing a model has to decide how many cards to give it first. The whole question is divisibility: transformers shards attention by head and the MLP by its intermediate dimension, and its all-gather assumes every rank holds an equal piece, so a degree that doesn't divide those evenly is not a slower option — it fails.
The answer is a single number, the largest degree that works. Every degree that
works is a divisor of it, so a caller wanting n ranks takes the smallest
divisor >= n; if there is none the model has to be spread another way.
Bases: ValueError
A tensor-parallel degree was asked for that this checkpoint cannot serve.
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.
The degree a distributed_config asks for, or None if it asks for none.
Accepts the dataclass or a plain dict, because transformers does.
Whether this distributed_config asks for expert parallelism.
Accepts the dataclass or a plain dict, because transformers does.
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 |