mixle.utils.parallel.model_decomposition module

Structural model-decomposition planner (component C2 of the model-parallel design).

Turns any mixle model into a model-parallel placement by walking the opt-in decomposition contract (mixle.stats.compute.decomposition) instead of a blind reflective __dict__ walk:

  • size_model_tree() – a structural byte sizing of the model tree (own params per node, recursing into children via the declared axis), so shared subtrees are not double-counted and a node’s own parameter footprint (e.g. an HMM’s dense S*S transition block) is visible rather than silently zeroed. Replaces the reflective estimate_model_nbytes for sizing decisions.

  • decompose_model() – choose the cut (data vs model vs none) for the root shardable node and bin-pack the component / factor units across the device budget, carrying the per-cut reduction. This is the general form of, and a real consumer-shaped output for, planner.model_sharding_plan (which only ever understood mixture components and had no consumer).

It works for every family: nodes that do not opt into the contract report Decomposition.atomic() and are simply replicated (data-parallel, already optimal). See ~/codex/notes/model-parallel-design.md.

class NodeSize(path, axis, num_units, own_param_bytes, subtree_param_bytes, subtree_work=1.0, children=())[source]

Bases: object

Structural size of one model-tree node.

Parameters:

  • path (str)

  • axis (DecompAxis)

  • num_units (int)

  • own_param_bytes (int)

  • subtree_param_bytes (int)

  • subtree_work (float)

  • children (tuple[NodeSize, ...])

path: str
axis: DecompAxis
num_units: int
own_param_bytes: int
subtree_param_bytes: int
subtree_work: float = 1.0
children: tuple[NodeSize, ...] = ()
size_model_tree(model, _path='', _seen=None)[source]

Recursively size the model tree via the decomposition contract (shared subtrees counted once).

Parameters:
Return type:

NodeSize

shard_children(node, dc=None)[source]

Return the actual child distributions a node splits along its declared axis (else ()).

Parameters:

  • node (Any)

  • dc (Decomposition | None)

Return type:

tuple[Any, …]

cost_children(model)[source]

ALL child distributions of a node (for COST), not just the shardable ones.

shard_children returns only the axis a node can be split along; for costing we need every nested distribution’s compute counted – e.g. an HMM is atomic (not shardable here) yet its S emission distributions and len_dist are real work, and a heavy leaf buried under a non-shardable wrapper still costs FLOPs. Discovery is reflective (any SequenceEncodableProbabilityDistribution held directly or inside a list/tuple/dict), so it works for any model without per-family wiring.

Parameters:

model (Any)

Return type:

tuple[Any, …]

subtree_work(model, _seen=None)[source]

Total compute weight of a model subtree – own emission cost plus ALL descendants (counted once).

Recurses over cost_children() (every nested distribution), so a unit’s cost includes heavy subtrees the executor can’t split (a nested HMM, a GP leaf), making the balance honest about where the FLOPs actually are – not just where the model happens to be shardable.

Parameters:
Return type:

float

compute_cost(model, _seen=None)[source]

(flops_per_observation_proxy, bytes) for the whole model – compute load and memory footprint, the two resources the balancer trades off (compute is the load, memory is the constraint).

Parameters:
Return type:

tuple[float, int]

class AxisCandidate(path, axis, reduction, num_units, unit_works)[source]

Bases: object

One shardable axis somewhere in the model tree, with the compute weight of each of its units.

Parameters:

  • path (str)

  • axis (DecompAxis)

  • reduction (ReductionOp)

  • num_units (int)

  • unit_works (tuple[float, ...])

path: str
axis: DecompAxis
reduction: ReductionOp
num_units: int
unit_works: tuple[float, ...]
property total_work: float
tree_axes(model)[source]

Enumerate every shardable axis anywhere in the tree (not just the root), each with per-unit work.

Parameters:

model (Any)

Return type:

list[AxisCandidate]

best_parallel_axis(model, max_workers=None)[source]

Pick the axis whose parallelization removes the most serial wall-time (heaviest, parallelizable).

Benefit of cutting an axis with P available workers is total_work * (1 - 1/min(P, num_units)) – favouring the axis that carries the most work AND has enough units to keep the workers busy. This looks at the WHOLE tree, so a heavy mixture nested inside a thin composite is found (the root-only planner missed it). Returns None when nothing in the tree is worth splitting.

Parameters:

  • model (Any)

  • max_workers (int | None)

Return type:

AxisCandidate | None

class ModelCut(device, start, stop, reduction)[source]

Bases: object

One model-parallel cut: a contiguous unit range placed on one device, with its reduction.

Parameters:

  • device (DeviceSpec)

  • start (int)

  • stop (int)

  • reduction (ReductionOp)

device: DeviceSpec
start: int
stop: int
reduction: ReductionOp
class ModelDecomposition(axis, reduction, num_units, cuts, rationale, extra=<factory>)[source]

Bases: object

The chosen decomposition of a model onto a device budget.

Parameters:

  • axis (DecompAxis)

  • reduction (ReductionOp)

  • num_units (int)

  • cuts (tuple[ModelCut, ...])

  • rationale (str)

  • extra (dict[str, Any])

axis: DecompAxis
reduction: ReductionOp
num_units: int
cuts: tuple[ModelCut, ...]
rationale: str
extra: dict[str, Any]
property is_model_parallel: bool
decompose_model(model, resources, *, n_data=None, min_components_per_shard=1, prefer_data_parallel=True)[source]

Decide how to place model across resources for model parallelism.

The WHOLE tree is searched for the heaviest worth-splitting axis (best_parallel_axis()), so a big mixture nested inside a thin composite is found – not just the root. A model that does not opt into the contract, or whose subtree fits replicated and is best served by data parallelism, yields an axis="none" plan (replicate the model, shard the data – already optimal for large N). Otherwise that axis’s units are bin-packed across devices BALANCED BY COMPUTE COST, carrying the per-cut reduction.

Parameters:

  • model (Any)

  • resources (Resources)

  • n_data (int | None)

  • min_components_per_shard (int)

  • prefer_data_parallel (bool)

Return type:

ModelDecomposition