mixle.utils.parallel.training_contracts module

Contracts shared by distributed gradient-training backends.

The encoded-data backends reduce sufficient statistics. Frontier neural training has a different lifecycle: materialize a process mesh, transform a module, bind optimizer state to logical parameters, execute steps, and save a reshardable training state. Keeping that lifecycle separate prevents a Ray or Megatron integration from pretending to be an EncodedDataHandle.

class BackendCapabilities(name, axes, precisions=frozenset({'bf16', 'fp32'}), distributed_optimizer=False, reshardable_checkpoint=False, elastic_restart=False, requirements=(), incompatible_axis_sets=())

Bases: object

What a backend actually executes, rather than merely accepting as flags.

Parameters:

• name (str)

• axes (frozenset[ParallelAxis])

• precisions (frozenset[str])

• distributed_optimizer (bool)

• reshardable_checkpoint (bool)

• elastic_restart (bool)

• requirements (tuple[str, ...])

• incompatible_axis_sets (tuple[frozenset[ParallelAxis], ...])

class CollectiveKind(*values)

Bases: StrEnum

Collective or point-to-point operation required by an update.

class DistributedTrainingBackend(*args, **kwargs)

Bases: Protocol

Structural extension point implemented by Torch, Megatron, or other engines.

class DistributedTrainingSession(*args, **kwargs)

Bases: Protocol

A materialized model, optimizer, mesh, and resumable step clock.

class DistributedUpdate(node_id, payload, collective, mesh_axes, state_layout, exact, notes=(), contract_exact=False, determinism_observed=None, maximum_absolute_error=None, maximum_relative_error=None, numerics_evidence_id=None, numerics_sample_count=0)

Bases: object

Executable communication declaration with explicit numerical evidence.

Parameters:

• node_id (str)

• payload (PayloadKind)

• collective (CollectiveKind)

• mesh_axes (tuple[ParallelAxis, ...])

• state_layout (StateLayout)

• exact (bool)

• notes (tuple[str, ...])

• contract_exact (bool)

• determinism_observed (bool | None)

• maximum_absolute_error (float | None)

• maximum_relative_error (float | None)

• numerics_evidence_id (str | None)

• numerics_sample_count (int)

class ParallelAxis(*values)

Bases: StrEnum

Named dimensions understood by Mixle and external training engines.

class ParallelPlan(dp_replicate=1, dp_shard=1, tp=1, pp=1, cp=1, ep=1, etp=1, microbatches=1, gradient_accumulation_steps=1)

Bases: object

An explicit N-D process mesh and batch schedule.

Every process belongs to exactly one coordinate in the product below. Keeping replicate and shard data-parallel dimensions separate permits HSDP without conflating it with tensor or context parallelism.

Parameters:

• dp_replicate (int)

• dp_shard (int)

• tp (int)

• pp (int)

• cp (int)

• ep (int)

• etp (int)

• microbatches (int)

• gradient_accumulation_steps (int)

property world_size: int

Number of physical ranks.

EP and ETP are overlapping process groups cut from the data/model domains, not extra orthogonal DeviceMesh dimensions.

property mesh: tuple[tuple[str, ...], tuple[int, ...]]

Return stable DeviceMesh names and shape, including a 1-rank DP dimension.

class ParameterLayout(logical_id, global_shape, placements=(), shared_group=None, optimizer_state=StateLayout.REPLICATED)

Bases: object

Stable logical identity and global layout for one parameter.

Parameters:

• logical_id (str)

• global_shape (tuple[int, ...])

• placements (tuple[tuple[str, str], ...])

• shared_group (str | None)

• optimizer_state (StateLayout)

class PayloadKind(*values)

Bases: StrEnum

Value communicated when an update is distributed.

class StateLayout(*values)

Bases: StrEnum

Physical ownership of mutable training state.

class StepReceipt(step, loss, local_examples, local_tokens, microbatches, accumulation_steps, data_parallel_size, optimizer, precision, collective_bytes=0, skipped=False, extra=<factory>, global_examples_observed=None, global_tokens_observed=None)

Bases: object

Unambiguous accounting for one committed optimizer update.

Parameters:

• step (int)

• loss (float)

• local_examples (int)

• local_tokens (int)

• microbatches (int)

• accumulation_steps (int)

• data_parallel_size (int)

• optimizer (str)

• precision (str)

• collective_bytes (int)

• skipped (bool)

• extra (dict[str, Any])

• global_examples_observed (int | None)

• global_tokens_observed (int | None)

get_training_backend(name)

Resolve a registered backend, calling a zero-argument lazy factory once.

Parameters:

name (str)

Return type:

Any

register_training_backend(name, backend, *, override=False)

Register a backend instance or lazy factory without importing optional dependencies.

Parameters:

• name (str)

• backend (Any)

• override (bool)

Return type:

None