mixle.utils.parallel.torch_training module

Executable PyTorch distributed-gradient backend.

This is the native path for arbitrary Mixle nn.Module objects. It uses public PyTorch APIs: DeviceMesh/DTensor tensor parallelism, FSDP2 (including HSDP), DDP, distributed checkpoint-compatible optimizer state, and the experimental context-parallel SDPA context manager. Pipeline and expert parallel transformer training belong to the Megatron backend, whose model partitioner owns those semantics.

class TorchDistributedBackend

Bases: object

Native backend for Mixle modules that remain ordinary PyTorch modules.

prepare(module, *, plan, **kwargs)

Materialize mesh, transforms, optimizer, and resumable session.

Parameters:

• module (Any)

• plan (ParallelPlan)

• kwargs (Any)

Return type:

TorchDistributedSession

class TorchDistributedSession(module, *, plan, capabilities, mesh, device, optimizer, optimizer_receipt, parameter_layouts, precision, max_grad_norm, scheduler=None, manual_data_parallel=False)

Bases: object

A transformed module plus optimizer and explicit process-mesh semantics.

Parameters:

• module (Any)

• plan (ParallelPlan)

• capabilities (BackendCapabilities)

• mesh (Any)

• device (Any)

• optimizer (Any)

• optimizer_receipt (dict[str, Any])

• parameter_layouts (tuple[ParameterLayout, ...])

• precision (str)

• max_grad_norm (float | None)

• scheduler (Any)

• manual_data_parallel (bool)

train_batch(inputs, targets)

Accumulate one local batch and commit when the configured clock fires.

Parameters:

• inputs (Any)

• targets (Any)

Return type:

StepReceipt

discard_accumulation()

Explicitly discard pending gradients without changing parameters or optimizer state.

Return type:

StepReceipt | None

save_checkpoint(path, *, loader_state=None, typed_scheduler_state=None, extra=None, asynchronous=False)

Save all state needed to resume this session.

Parameters:

• path (str)

• loader_state (Any)

• typed_scheduler_state (Any)

• extra (dict[str, Any] | None)

• asynchronous (bool)

Return type:

Any

load_checkpoint(path, *, allow_world_size_change=False)

Restore this transformed model/optimizer and resume its step clock.

Parameters:

• path (str)

• allow_world_size_change (bool)

Return type:

dict[str, Any]

reduce_sums(*values)

Sum scalar metrics over the two-dimensional data-parallel domain.

Parameters:

values (float)

Return type:

tuple[float, …]

class TorchSessionState(step, pending_accumulation, optimizer_receipt)

Bases: object

Serializable non-tensor clocks owned by a native training session.

Parameters:

• step (int)

• pending_accumulation (int)

• optimizer_receipt (dict[str, Any])

describe_parameter_layouts(module, plan)

Capture stable logical identities before distributed wrappers rewrite parameters.

Parameters:

• module (Any)

• plan (ParallelPlan)

Return type:

tuple[ParameterLayout, …]