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()) – wraps mixle.utils.parallel.dcp_checkpoint.save_sharded()’s underlying torch.distributed.checkpoint call so a checkpoint does not block training: the (bounded, D2H-copy) cost of cloning the state dict to a frozen CPU snapshot happens synchronously on the caller’s thread, and the (unbounded, I/O-latency-bound) cost of actually writing it to disk happens on a background thread. Loader state rides along in the same checkpoint directory (a sibling loader_state.json), so a resume restores model + optimizer + data position together, atomically from the caller’s point of view.

  • 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 and resume from the last checkpoint’s loader state (not from scratch, not re-running the whole job).

  • 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)[source]

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:
advanced(n=1)[source]

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

Parameters:

n (int)

Return type:

LoaderState

resume_batches(corpus, state)[source]

Resume a loader shaped like mixle.data.streaming_corpus.StreamingCorpus (anything exposing epoch_batches(epoch) -> Iterator[(x, y)] 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)[source]

Bases: object

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

Parameters:
save_checkpoint_async(module, optimizer, path, loader_state, *, extra=None)[source]

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

Refines mixle.utils.parallel.dcp_checkpoint.save_sharded() for the async case: that function calls dcp.save directly on the live state dict, which blocks the caller for the full write – fine for a synchronous checkpoint, unsafe to background (the live tensors keep changing under the writer). Here, 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 actual dcp.save call – and the sibling loader_state.json write – 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 (plus any caller-supplied extra, e.g. every rank’s LoaderState in a multi-rank job) is written alongside the DCP checkpoint directory as JSON – resuming needs both the model/optimizer AND the data position, and this keeps them physically bundled under one path.

Parameters:

  • module (Any)

  • optimizer (Any)

  • path (str)

  • loader_state (LoaderState)

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

Return type:

AsyncCheckpointHandle

load_checkpoint(module, optimizer, path)[source]

Load a checkpoint written by save_checkpoint_async() (or plain save_sharded, if a sibling loader_state.json was written by hand) into module/optimizer in place; returns the captured LoaderState so the caller’s data loader can resume from the exact same position.

Parameters:
Return type:

LoaderState

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

Bases: object

Parameters:
class SimulatedRank(rank_id, model_factory, batch_fn)[source]

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()[source]

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

Return type:

None

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

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 and resume from the last checkpoint’s model/optimizer/loader state instead of the whole job restarting from scratch.

Holds one canonical (model, optimizer) (what gets checkpointed and what a respawned rank loads); 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.

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)

  • health_monitor (TrainingHealthMonitor | None)

run_step(step, kill_ranks=frozenset())[source]

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:
Return type:

dict[str, Any]

checkpoint(path=None)[source]

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)[source]

Elastic restart: bring rank_id back from the last checkpoint – model, optimizer, and every rank’s loader state – instead of restarting the whole job from scratch. Mirrors resilient_em’s _respawn_worker: same rank id, resumed data position, job otherwise 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