mixle.utils.parallel.fault_tolerant_training module

Fault-tolerant gradient training: async DCP snapshots, loader-state capture, elastic restart, resume-with-receipts (roadmap F2).

This is the gradient-training-side analogue of mixle.utils.parallel.resilient_em (K4): that module makes the accumulator-combining EM path tolerant of a dying worker (retry, blacklist, elastic re-partition, deterministic rendezvous-based chaos test); this module carries the same PATTERN – “detect a failure, don’t restart the whole job, resume from a checkpoint” – to the gradient-training path, where the parallelism is data-parallel gradient averaging rather than additive sufficient-statistic folding, so the mechanics differ even though the shape of the fault-tolerance story does not:

• Async DCP snapshots (save_checkpoint_async()) – freezes model, optimizer, loader, and caller metadata at one instant, then writes them on a background thread through the same immutable-generation, integrity-manifest, and atomic-publication protocol as synchronous training checkpoints. A resume therefore restores model + optimizer + data position together without trusting a detached sidecar.

• Loader-state capture (LoaderState) – mirrors the resumability contract mixle.data.streaming_corpus.StreamingCorpus (F3, PR #139) already guarantees: epoch_batches(epoch) is a pure, deterministic function of (seed, epoch, rank, world_size), so the ONLY thing that changes as an epoch progresses is how many batches have been consumed – capturing (seed, epoch, rank, world_size, batch_idx) is sufficient to reconstruct the identical remaining stream. resume_batches() resumes any loader that exposes that same epoch_batches contract, F3’s or a synthetic stand-in.

• Elastic restart (SimulatedRank, ElasticTrainingJob) – mirrors ResilientMPEncodedData’s deterministic kill rendezvous (each rank signals “started” – here, once its forward+backward for a step is done – then blocks for an explicit “go” from the driver before the optimizer step commits, so a chaos test’s kill lands at a known point, not a timing race) but adapted to data-parallel gradient averaging: a dead rank’s gradient is simply excluded from the step’s average (the job continues with fewer ranks, degrading gracefully) rather than failing the whole step, and a dead rank can be elastically respawned at its last checkpointed loader position while joining the current canonical model generation (not from scratch and without rewinding surviving ranks).

• Resume-with-receipts – ElasticTrainingJob.respawn_rank() marks the NEXT observed step as restart=True when it feeds mixle.utils.parallel.training_health.TrainingHealthMonitor (F4, PR #147), so every restart automatically gets F4’s per-restart continuity verdict for free – this module does not reimplement that check, it wires into it.

Scope note (mirrors the dcp_checkpoint / resilient_em modules’ own scoping): “10k A100s” framing aside, checkpointing, async snapshotting, elastic restart bookkeeping, and loss-continuity verification are all exact regardless of scale – what genuinely does not exist on a laptop is FSDP2 sharding a model too big for one device and a real multi-node NCCL all-reduce. Those two are simulated here: world_size ranks are real concurrent OS threads (not a real distributed job), and gradient “all-reduce” is a plain mean over surviving ranks’ locally computed grads. Every other piece – DCP save/load, the CPU-clone async mechanism, loader-state round-tripping, the kill rendezvous, and the continuity check – is the same code that would run at 10k GPUs, exercised at small scale.

class LoaderState(seed, epoch, rank, world_size, batch_idx=0)

Bases: object

Resumability state of one rank’s data loader: enough to reproduce its exact next batch after a restart, without serializing RNG internals or buffered batches.

Mirrors mixle.data.streaming_corpus.StreamingCorpus’s contract: epoch_batches(epoch) is a pure function of (seed, epoch, rank, world_size) ( global_document_order() reseeds from (seed, epoch) via SeedSequence, then shard_documents_for_rank() deterministically slices per rank) – the only thing that varies as an epoch progresses is how many batches of that deterministic stream have already been consumed, i.e. batch_idx.

Parameters:

• seed (int)

• epoch (int)

• rank (int)

• world_size (int)

• batch_idx (int)

advanced(n=1)

The state after n more batches have been consumed this epoch.

Parameters:

n (int)

Return type:

LoaderState

resume_batches(corpus, state)

Resume a loader shaped like mixle.data.streaming_corpus.StreamingCorpus (anything exposing epoch_batches(epoch) -> Iterator[(context, targets, loss_mask)] with that contract) exactly at state.batch_idx.

Determinism is what makes this correct rather than approximate: re-materializing the whole epoch and discarding the already-consumed prefix reproduces bitwise-identical remaining batches to what an uninterrupted run would have yielded from that point on – the same trick mixle.utils.parallel.resilient_em.checkpointed_fold() relies on for exact accumulator recovery.

Parameters:

• corpus (Any)

• state (LoaderState)

class AsyncCheckpointHandle(thread, path, prepare_time_s, errors=<factory>)

Bases: object

A checkpoint write in flight (or finished) on a background thread.

Parameters:

• thread (Thread)

• path (str)

• prepare_time_s (float)

• errors (list[BaseException])

save_checkpoint_async(module, optimizer, path, loader_state, *, extra=None)

Snapshot (model, optimizer, loader_state) to path without blocking the training loop.

The ONLY synchronous work is get_state_dict + a detached CPU clone (a bounded D2H-copy cost that does not scale with disk/network write latency). The transactional generation write and publication happen on a background thread, so this function returns as soon as the clone is done and the training loop’s next step can start immediately.

loader_state and caller extra use the same integrity-bound typed sidecar and immutable-generation publication protocol as mixle.utils.parallel.dcp_checkpoint.

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

• loader_state (LoaderState)

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

Return type:

AsyncCheckpointHandle

load_checkpoint(module, optimizer, path)

Load a verified transactional checkpoint and return its loader position.

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

Return type:

LoaderState

class StepResult(rank: 'int', step: 'int', loss: 'float', grad_norm: 'float', grads: 'list[torch.Tensor]')

Bases: object

Parameters:

• rank (int)

• step (int)

• loss (float)

• grad_norm (float)

• grads (list[torch.Tensor])

class SimulatedRank(rank_id, model_factory, batch_fn)

Bases: object

One data-parallel rank’s local training-step worker, run on a real background thread.

Mirrors mixle.utils.parallel.resilient_em.ResilientMPEncodedData’s rendezvous: after computing a full forward+backward pass for a step – the point at which a real GPU worker would ordinarily all-reduce gradients and step the optimizer – the thread signals started and BLOCKS waiting for an explicit go from the driver. This pins “mid-step” to a known point (strictly after gradient computation, strictly before the step is applied), so a chaos test’s kill is deterministic, not a timing race: a kill issued at this rendezvous is guaranteed to land before any weight update happens.

kill() needs no OS-level teardown (real process kill, as resilient_em does, is not available for an in-process thread): simply never releasing the rendezvous IS the simulated crash – the thread times out and exits with no result, exactly as a real dead worker would leave the driver’s recv() hanging until it gives up.

Parameters:

• rank_id (int)

• model_factory (Callable[[], Any])

• batch_fn (Callable[[], tuple[Any, Any]])

release()

Wave this rank through the rendezvous – it survives this step.

Return type:

None

cancel()

Cancel this step and wake a worker already waiting at the rendezvous.

Return type:

None

class ElasticTrainingJob(model_factory, world_size, batch_fn_for_rank, checkpoint_dir, *, seed=0, lr=1e-2, step_timeout_s=5.0, health_monitor=None)

Bases: object

A data-parallel training loop, chaos-tolerant: a rank dying mid-step degrades gracefully – the job continues averaging over fewer surviving ranks that step, rather than hard-failing – and a dead rank can be elastically respawned at its last checkpointed loader position and rejoin the current canonical model generation instead of restarting or rewinding the whole job.

Holds one canonical (model, optimizer) (what gets checkpointed and what every live replica receives at the start of its next step); each SimulatedRank computes its OWN local forward+backward against a fresh copy of the canonical weights (the data-parallel replica), and the driver applies the mean of surviving ranks’ gradients to the canonical model once per step – the plain-mean “all-reduce” this module’s docstring flags as the one piece that is genuinely simulated rather than exercised for real.

Every restart is wired into F4’s continuity check for free: respawn_rank() marks the NEXT run_step call as restart=True when it feeds health, so health.report()["restarts"] always carries a real per-restart continuity verdict, not something the caller has to remember to ask for.

Single-driver only. run_step(), checkpoint(), and respawn_rank() share mutable state (canonical_model, canonical_optimizer, loader_states) with no internal locking, and are meant to be called from the one thread that constructed this job – calling any of them from a second thread while another is in flight is a misuse this class detects and rejects (RuntimeError) rather than silently letting it race: run_step advances the canonical weights before it advances loader_states, so a concurrent checkpoint() can capture a model state newer than the loader state written alongside it. Periodic/background checkpointing must be driven from the same thread as run_step (e.g. called between steps), not a separate one.

Parameters:

• model_factory (Callable[[], Any])

• world_size (int)

• batch_fn_for_rank (Callable[[int, LoaderState], tuple[Any, Any]])

• checkpoint_dir (str)

• seed (int)

• lr (float)

• step_timeout_s (float)

• health_monitor (TrainingHealthMonitor | None)

run_step(step, kill_ranks=frozenset())

Run one data-parallel step. kill_ranks simulates a mid-step death for those ranks: they reach the post-backward rendezvous (so their compute genuinely happened) but are never released, so their gradient is excluded from this step’s average – the job continues with fewer ranks rather than raising.

Parameters:

• step (int)

• kill_ranks (frozenset[int])

Return type:

dict[str, Any]

checkpoint(path=None)

Async-snapshot the canonical model/optimizer plus every rank’s loader state.

Parameters:

path (str | None)

Return type:

AsyncCheckpointHandle

respawn_rank(rank_id, checkpoint_path=None)

Bring rank_id back at its checkpointed data position without rolling back the live job.

The respawned replica receives the current canonical model at the start of its next step, exactly like every other replica. Only its own loader position comes from the checkpoint; committed model updates and surviving ranks’ positions remain untouched. Marks the next run_step as a restart so F4’s continuity check evaluates it.

Parameters:

• rank_id (int)

• checkpoint_path (str | None)

Return type:

LoaderState