mixle.utils.parallel.megatron_training module

Megatron Bridge adapter for full transformer and MoE parallelism.

Megatron owns the transformer partitioner and training loop. Mixle owns the statistical/parallel plan and maps it into a Bridge model provider without importing Megatron unless this backend is selected.

class MegatronBridgeBackend

Bases: object

Map a ParallelPlan onto NVIDIA Megatron Bridge providers.

static configure_provider(provider, plan)

Apply the Mixle plan to a mutable Bridge model provider.

Parameters:

• provider (Any)

• plan (ParallelPlan)

Return type:

Any

static configure_distributed_optimizer(config, plan)

Enable sharded optimizer/parameter state and overlap on a Bridge config.

Parameters:

• config (Any)

• plan (ParallelPlan)

Return type:

Any

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

Finalize a provider and construct its distributed model.

Parameters:

• module (Any)

• plan (ParallelPlan)

• kwargs (Any)

Return type:

MegatronBridgeSession

run_pretrain(config, forward_step, *, plan=None)

Delegate the full training loop to Bridge’s public pretrain entry point.

Parameters:

• config (Any)

• forward_step (Any)

• plan (ParallelPlan | None)

Return type:

Any

class MegatronBridgeSession(model, *, plan, capabilities, train_batch_fn=None, optimizer=None, precision='bf16')

Bases: object

A materialized Bridge model for callers that supply a custom batch step.

Normal frontier pretraining should call MegatronBridgeBackend.run_pretrain(); Bridge then owns schedules, inter-stage transfers, distributed optimizer overlap, and checkpoint cadence. train_batch exists for integration tests and specialized loops and never silently substitutes a local step.

Parameters:

• model (Any)

• plan (ParallelPlan)

• capabilities (BackendCapabilities)

• train_batch_fn (Any)

• optimizer (Any)

• precision (str)