Distributed Transformer Training

Mixle has two distributed-gradient backends. They share an explicit ParallelPlan and fail before fitting when a backend cannot execute a requested axis. This is separate from the Spark, Dask, Ray, and MPI encoded-data backends, which merge sufficient statistics rather than gradients.

Backend split

torch_native

Keeps an arbitrary Mixle nn.Module in PyTorch. It supports replicated data parallelism (DDP), sharded or hybrid data parallelism (FSDP2/HSDP), MLP tensor parallelism through DTensor, and PyTorch’s CUDA context-parallel SDPA path. TP+CP composition, pipeline parallelism, and expert parallelism are capability errors on this backend.

megatron

Maps the plan to a Megatron Bridge model provider. Megatron owns full transformer/MoE TP, PP, CP, EP, ETP, distributed optimizer overlap, and its pretraining schedule. Mixle does not replace those kernels with local simulations.

EP and ETP describe overlapping Megatron process groups. They do not multiply the physical world size as independent DeviceMesh dimensions. EP must divide the data-parallel domain.

Native packed training

The compact LM surface uses packed all-position teacher forcing:

lm.fit(
    token_ids,
    dense=True,
    distributed=True,
    distributed_backend="torch_native",
    dp_shard=8,
    precision="bf16",
    microbatches=4,
    gradient_accumulation_steps=8,
)

Rows are shuffled globally and deterministically, then partitioned over the data-parallel domain. Model-parallel ranks see the same rows. Step receipts distinguish local examples, global examples, microbatches, and accumulation. The legacy streaming one-target objective accepts data parallelism only; it refuses TP/PP/CP/EP/ETP rather than running a different plan.

Megatron Bridge

For a Bridge provider, resolve and configure the backend directly:

from mixle.utils.parallel import ParallelPlan, get_training_backend

plan = ParallelPlan(
    dp_replicate=8,
    tp=8,
    pp=4,
    cp=2,
    ep=8,
    etp=2,
    microbatches=16,
)
backend = get_training_backend("megatron")
backend.configure_provider(provider, plan)
backend.run_pretrain(config, forward_step, plan=plan)

The dependency is optional and imported only when selected. A cluster image must provide Megatron Bridge, CUDA, and NCCL.

Cluster launch

RayTrainLauncher starts one worker per physical rank and passes the complete plan to the worker configuration. LightningFabricLauncher is intentionally limited to DDP/FSDP launches; model axes are capability errors. These launchers manage processes only. The selected training backend still owns parameter layout, collectives, optimizer state, and checkpoints.

Optimizer and checkpoint semantics

Optimizer routes are planned from stable logical parameter names before wrapping. Global matrix geometries such as Muon are not applied independently to parameter shards: the native backend records a momentum fallback until it has a declared global factor/gather implementation. Adam and AdamW remain explicit options, not the automatic default.

save_training_state stores DCP model and optimizer shards plus scheduler, gradient scaler, Python/NumPy/Torch RNGs, loader position, optimizer step, parallel geometry, typed-scheduler state, and caller metadata. A checkpoint is loadable only after its _SUCCESS manifest is committed. Native asynchronous saves use DCP’s future and report write failures from wait().

Evidence boundary

CPU tests establish objective parity, accumulation semantics, provider mapping, checkpoint completeness, and capability failures. They do not establish GPU throughput, MFU, overlap efficiency, or multi-node recovery. Those claims remain unverified until retained hardware runs are attached to the GPU claims ledger.