mixle.utils.parallel.training_health module

Training-health receipts: MFU, loss/grad-norm anomaly detection, precision checks, restart continuity.

The frontier-scale trainer (the vendored TorchTitan/Megatron loop atop mixle.models.transformer, sharded via mixle.utils.parallel.torchrun / mixle.utils.parallel.dcp_checkpoint) runs for weeks unattended – the only way to know it is healthy is receipts computed from the loop itself, not a human staring at a loss curve. TrainingHealthMonitor is that receipt machine: call observe_step(...) once per optimizer step (mirroring mixle.telemetry.core.Telemetry.record()) and call report() once at the end (mirroring mixle.evolve.ledger.EvolutionLedger / mixle.evolve.population.OperatorBandit.report()) for a structured, JSON-serializable summary.

Four things are tracked:

  • MFU – ModelFlopConfig computes the theoretical FLOPs/step for a transformer config (the standard 6N + attention accounting, same formula nanoGPT’s estimate_mfu uses); achieved FLOPs/sec comes from a caller-supplied wall-clock step time (real timing, whatever hardware the loop runs on); MFU is the ratio against a caller-supplied hardware peak.

  • Loss-spike / changepoint detection – a robust (median/MAD) rolling z-score per step.

  • Grad-norm / precision anomalies – the same rolling z-score on grad-norm, plus NaN/Inf checks on both streams (these always fire, independent of the rolling window’s warmup).

  • Per-restart continuity – restart=True on the first step after a checkpoint resume marks the rolling baseline boundary; the next step’s z-score is evaluated against the pre-restart baseline (it has not been updated with any post-restart value yet), so a resume that silently drops optimizer/RNG state and produces a real loss jump is caught as restart_discontinuity – a well-behaved resume is not.

  • Dead-rank liveness – observe_rank_step(rank, step) is a per-rank heartbeat: a data-parallel loop calls it once per step per rank (mirroring how ElasticTrainingJob already tracks dead_ranks for gradient-averaging purposes, but that bookkeeping never surfaced as a health receipt – this does). check_rank_liveness(current_step) flags any rank that has gone silent for more than rank_heartbeat_threshold steps as dead_rank, once per outage.

No cluster is required to exercise any of this: the FLOPs accounting and anomaly math are exact regardless of scale, and the tests drive a real (tiny) mixle.models.transformer.build_causal_lm() for a handful of steps. The absolute MFU number a laptop/CI runner produces is not comparable to a real cluster’s – that comparison is deferred until the real distributed trainer (roadmap F1) exists.

class Anomaly(step, kind, value, baseline, z_score, detected_at_step)[source]

Bases: object

One flagged anomaly: what kind, at which step, and the baseline it was judged against.

Parameters:
latency(injected_at_step)[source]

Steps between the true anomaly step and detection – 0 means caught on the same step.

Parameters:

injected_at_step (int)

Return type:

int

class ModelFlopConfig(n_params, n_layer, n_head, d_model, seq_len)[source]

Bases: object

The shape a transformer’s FLOPs/step depends on – enough to compute theoretical FLOPs exactly.

n_params should exclude position-embedding params (they are a lookup, not a matmul); use flop_config_from_causal_lm() to derive this correctly from a real mixle.models.transformer.CausalLM.

Parameters:
class MFUSample(step, step_flops, step_time_s, peak_flops_per_sec)[source]

Bases: object

One step’s MFU measurement: real wall-clock timing against the theoretical FLOPs for that step.

Parameters:
class RollingBaseline(window=20, min_periods=5)[source]

Bases: object

A robust rolling baseline (median + scaled MAD) over a trailing window of values.

Median/MAD rather than mean/std so a previous spike does not itself blow up the spread used to judge the next one. z_score(value) is evaluated against the window as it stood before value was seen, so calling update only after scoring makes the check causal (no leakage of the current point into its own baseline) – this is also what makes the restart-continuity check work: the window carries only pre-restart history until the caller updates it.

Parameters:

  • window (int)

  • min_periods (int)

baseline()[source]

(median, scaled_mad) of the current window, or None during warmup.

Return type:

tuple[float, float] | None

z_score(value)[source]

Robust z-score of value against the window before value is added, or None in warmup.

Parameters:

value (float)

Return type:

float | None

state()[source]

Serializable snapshot – carry this across a checkpoint/restart to preserve continuity.

Return type:

dict[str, Any]

class StepRecord(step, loss, grad_norm=None, step_time_s=None, restart=False)[source]

Bases: object

One observed training step.

Parameters:
class TrainingHealthMonitor(*, flop_config=None, peak_flops_per_sec=None, loss_window=20, loss_min_periods=5, loss_z_thresh=6.0, grad_window=20, grad_min_periods=5, grad_z_thresh=6.0, rank_heartbeat_threshold=50)[source]

Bases: object

The run object: observe_step(...) per optimizer step, report() once at the end.

Follows the shape of mixle.telemetry.core.Telemetry (record/buffer) and mixle.evolve.ledger.EvolutionLedger (.report()-style terminal summary): no I/O, pure in-process accounting, JSON-serializable output.

Parameters:

  • flop_config (ModelFlopConfig | None)

  • peak_flops_per_sec (float | None)

  • loss_window (int)

  • loss_min_periods (int)

  • loss_z_thresh (float)

  • grad_window (int)

  • grad_min_periods (int)

  • grad_z_thresh (float)

  • rank_heartbeat_threshold (int)

observe_step(step, loss, *, grad_norm=None, step_time_s=None, batch_size=None, restart=False)[source]

Record one step; returns any anomalies raised at this step (also appended to self.anomalies).

Parameters:
Return type:

list[Anomaly]

observe_rank_step(rank, step)[source]

Record a per-rank heartbeat: rank reported liveness (e.g. completed a local forward+backward) at step. Call once per step per rank; combine with check_rank_liveness() to detect a rank that has stopped reporting. A rank reporting again after an outage clears its dead_rank flag so a later re-death can be flagged again.

Parameters:
Return type:

None

check_rank_liveness(current_step)[source]

Flag any known rank that has not reported a heartbeat (observe_rank_step()) for more than rank_heartbeat_threshold steps as dead_rank. Raised once per outage (not once per step) – the flag is cleared the next time that rank reports in, so a respawned/recovered rank can be caught going dead again later.

Parameters:

current_step (int)

Return type:

list[Anomaly]

continuity_ok()[source]

True unless any restart_discontinuity was ever flagged.

Return type:

bool

report()[source]

A complete, JSON-serializable summary: step count, MFU, anomalies by kind, continuity verdict.

Return type:

dict[str, Any]

theoretical_flops_per_iter(*, n_params, n_layer, n_head, d_model, seq_len, batch_size)[source]

Theoretical forward+backward FLOPs for one training iteration of a decoder-only transformer.

The standard two-term accounting (see nanoGPT’s estimate_mfu / the PaLM paper appendix): 6*N per token for the dense matmuls (forward+backward is ~3x forward, and each matmul is a multiply-add = 2 FLOPs, giving the well-known 6N per-token constant), plus 12*L*H*Q*T per token for the attention matmuls (QK^T and attn@V), which scale with context length T and are not captured by the 6N param-count term. Multiplying by T (all tokens in the sequence) and batch_size gives the FLOPs for one full iteration.

Parameters:
Return type:

float

flop_config_from_causal_lm(model, seq_len)[source]

Derive a ModelFlopConfig from a real mixle.models.transformer.CausalLM.

Position-embedding params are excluded from the count (a lookup table, not a matmul) – the same convention nanoGPT’s get_num_params(non_embedding=True) uses.

Parameters:
Return type:

ModelFlopConfig