feat(adapter): add Megatron Bridge and NeMo RL checkpoint adapters

[g-husam]

AI generated description

This change adds two new adapters so that jobs running on Megatron Bridge — and on
NeMo RL, which is built on top of it — can write their frequent crash-recovery
checkpoints into ML Flashpoint's node-local memory instead of durable storage,
while their durable checkpoints keep going exactly where they went before.

It also adds the tooling to prove the difference: a harness that pulls checkpoint
timings out of training logs and diffs two runs.

Why

Megatron Bridge already distinguishes two kinds of checkpoint. This change maps
ML Flashpoint onto the cheap one and leaves the expensive one alone:

Checkpoint	Cadence	Written by	Durability
Persistent	save_interval	Megatron Bridge, unchanged	Durable — GCS, network FS
Non-persistent	non_persistent_save_interval	ML Flashpoint (new)	Node-local memory, replicated to a peer

The two cadences can't collide, because Bridge's own loop treats them as mutually
exclusive:

# megatron/bridge/training/train.py::checkpoint_and_decide_exit
if save and save_interval and step % save_interval == 0:
    save_checkpoint_and_time(..., non_persistent_ckpt=False)
elif save and non_persistent_save_interval and step % non_persistent_save_interval == 0:
    save_checkpoint_and_time(..., non_persistent_ckpt=True)

That elif is why this adapter needs no skip logic, unlike the NeMo 2.0 adapter and
its skip_every_n_steps.

Megatron Bridge: use the supported hook

Bridge lets a run swap in its own checkpoint manager via
CheckpointConfig.custom_manager_class. MLFlashpointBridgeCheckpointManager
implements that protocol. The whole routing decision is four lines:

# src/ml_flashpoint/adapter/megatron_bridge/checkpoint_manager.py
def save(self, ctx, callback_manager=None) -> None:
    step = ctx.state.train_state.step
    mlf_logging.update_training_step(step)

    if not ctx.non_persistent_ckpt or self._ensure_runtime() is None:
        self._delegate_save(ctx, callback_manager)   # Bridge's own save_checkpoint
        return

    try:
        self._save_ml_flashpoint(ctx, step)
    except Exception:
        _LOGGER.exception(
            "ML Flashpoint save failed at step %d. Skipping this non-persistent checkpoint and continuing.",
            step,
        )

A durable checkpoint is untouched. A non-persistent one goes to ML Flashpoint, and
if that fails, training keeps running — losing a crash-recovery checkpoint should
never take down a job.

Turning it on

from megatron.bridge.training.config import CheckpointConfig
import ml_flashpoint.adapter.megatron_bridge as mlf_bridge

checkpoint = CheckpointConfig(
    save="/gcs/my-run/checkpoints",
    save_interval=500,
    ckpt_format="torch_dist",
)

mlf_bridge.enable(checkpoint, non_persistent_save_interval=20)

enable() sets custom_manager_class, non_persistent_ckpt_type="local" and
non_persistent_save_interval — and registers ml_flashpoint with Bridge's import
allowlist, which is easy to miss:

# Without this, Bridge raises InstantiationException on a non-allowlisted prefix.
# Defaults cover megatron.*, torch.*, transformers.*, nvidia.*, numpy.*, nemo.* -- not ml_flashpoint.
register_allowed_target_prefix("ml_flashpoint")

YAML works too, as long as that prefix is registered before setup runs:

checkpoint:
  save: /gcs/my-run/checkpoints
  save_interval: 500
  non_persistent_save_interval: 20
  non_persistent_ckpt_type: local
  custom_manager_class: ml_flashpoint.adapter.megatron_bridge.MLFlashpointBridgeCheckpointManager

Making resume find the checkpoint

Bridge decides whether to attempt a resume by asking the checkpoint manager's
context for a local checkpoint manager:

# megatron/bridge/training/setup.py::_should_load_checkpoint
checkpointing_context = getattr(checkpoint_manager, "checkpointing_context", {})
has_local_checkpoint = (
    "local_checkpoint_manager" in checkpointing_context
    and checkpointing_context["local_checkpoint_manager"].find_latest() != -1
)

So the adapter publishes MLFlashpointLocalCheckpointIndex under that key. It
implements only find_latest() and local_ckpt_dir — deliberately not load(),
because Bridge's own local-load path expects an NVRx MCoreTensorAwareStateDict
container, which ML Flashpoint does not produce. Before falling back to Bridge, the
index is disabled so Bridge is never offered a container it can't read:

if self._local_index is not None:
    self._local_index.disable()   # find_latest() now reports -1
return load_checkpoint(...)       # Bridge's own path

Load order is: newest recoverable ML Flashpoint container first; Bridge's durable
path if there is none, or if the in-memory read fails.

NeMo RL: the config hook doesn't reach it

NeMo RL never calls create_checkpoint_manager. A grep for
custom_manager_class across the entire NeMo RL repo returns zero hits. It builds
its Megatron state with Bridge, but drives checkpointing itself:

# nemo_rl/models/policy/workers/megatron_policy_worker.py::save_checkpoint
maybe_finalize_async_save(self.mcore_state, ckpt_cfg=..., blocking=True)
self.mcore_state.cfg.checkpoint.save = weights_path
save_checkpoint(                       # Bridge's *function*, not its manager
    state=self.mcore_state,
    model=[self.model],
    checkpointing_context=self.checkpointing_context,
)

So a config-only integration is impossible here. Instead, install_into_worker
wraps that one method on a worker instance:

from ml_flashpoint.adapter import nemo_rl as mlf_nemo_rl

mlf_nemo_rl.install_into_worker(worker, mode=mlf_nemo_rl.MODE_AUGMENT)

The wrapper is small enough to read in full:

@functools.wraps(original_save)
@log_execution_time(logger=_LOGGER, name="nemo_rl.save_checkpoint", level=logging.INFO)
def save_checkpoint(weights_path: str, optimizer_path: Optional[str] = None, **kwargs):
    save_durable = hooks.should_save_durable()
    _save_ml_flashpoint(worker, hooks, optimizer_path is not None)
    if not save_durable:
        return None
    return original_save(weights_path, optimizer_path=optimizer_path, **kwargs)

Two modes:

• augment (default) — every durable write still happens, with an ML Flashpoint
  checkpoint alongside. Faster recovery, unchanged durability.
• replace with durable_every_n_saves=N — only every Nth checkpoint stays
  durable. This is what removes checkpoint stalls from the RL loop.

For A/B runs, install_from_env(worker) reads
MLFLASHPOINT_NEMO_RL_ENABLED / _MODE / _DURABLE_EVERY_N_SAVES, so both arms
share one launch command.

Why not just inject a save strategy?

Seeding checkpointing_context["save_strategy"] would be honoured by Bridge's
save_checkpoint, and it's tempting because it needs no worker changes. It's wrong
here for two reasons: dist_checkpointing.save writes common.pt on global rank 0
only, so other nodes couldn't recover locally; and it writes into the caller's
weights_path, putting a node-local memory container inside the durable directory.

Measuring it

Where the numbers come from

Bridge wraps every checkpoint in barriers and then logs it, so the recorded value is
the wall-clock time the whole job stalled:

# megatron/bridge/training/train.py::save_checkpoint_and_time
timer_key = "save-checkpoint-non-persistent" if non_persistent_ckpt else "save-checkpoint"
timers(timer_key, log_level=0).start(barrier=True)
checkpoint_manager.save(CheckpointSaveContext(...), callback_manager)
timers(timer_key).stop(barrier=True)
timers.log([timer_key])

That produces lines like this, in milliseconds:

(min, max) time across ranks (ms):
    save-checkpoint ................................: (18450.20, 18512.90)

The pair is (min, max) across ranks, not across steps — so this is one
checkpoint, where the fastest rank finished in 18.45 s and the slowest in 18.51 s.
The parser keeps the max, because the barrier means every rank waits for the
slowest one. This line contributes a single sample of 18.51 s.

Turning that into a comparison

scripts/benchmarks/parse_checkpoint_timings.py --label baseline logs/baseline.log --output baseline.json
scripts/benchmarks/parse_checkpoint_timings.py --label flashpoint logs/mlf.log --output flashpoint.json
scripts/benchmarks/compare_checkpoint_timings.py --baseline baseline.json --candidate flashpoint.json

timer                                           n      baseline    flashpoint  mean delta
----------------------------------------------------------------------------------------
save-checkpoint                                 3       18.306s       18.402s  +0.096s (+0.5%, 0.99x)
save-checkpoint-non-persistent                 10             -        0.621s  n/a

n is the number of checkpoints observed in that arm; the two middle columns are
mean seconds per checkpoint.

Row 1 is the control, not the result. Durable checkpoints cost the same in both
arms (18.306 s vs 18.402 s) — which is exactly right, because this change doesn't
touch the durable path. A real difference here would mean the two runs differed in
something else, and would invalidate the rest of the table.

Row 2 is the new work. It shows - and n/a because the baseline has no
non-persistent cadence at all, so there is no same-named timer to subtract from.

The headline number is the cross-row comparison, which the tool deliberately
does not compute for you:

baseline   save-checkpoint                 18.306 s   <- what a checkpoint used to cost
flashpoint save-checkpoint-non-persistent   0.621 s   <- what the substituted checkpoint costs now
                                           ---------
                                           ~29x faster, ~17.7 s off each substituted checkpoint

On steps where ML Flashpoint now holds the checkpoint, the loop stalls for ~0.6 s
instead of ~18 s. That is only a real saving in replace mode, where those steps
genuinely skip the durable write — in augment mode the durable write still happens,
so row 2 is added cost rather than saved time.

The numbers above are illustrative, to explain the table's shape. See
What is not covered — this experiment has not been run.

docs/checkpoint-timing-experiment.md is the runbook: cluster requirements (≥2 nodes
so replication is exercised, /dev/shm sized for two shards), a ~25-step run config,
and how to read the result including the sample-size caveat.

Testing

184 passed
96% coverage on the new packages

Two of the tests caught real bugs while being written: a disabled manager still
handing back its runtime instead of None, and the resume index never getting
installed because the context property waited on a runtime that Bridge builds later.

tests/adapter/conftest.py stubs megatron.bridge.* when it isn't installed, so the
suite runs without the full NVIDIA stack (Transformer Engine, ModelOpt, NVRx). When
the real package is present, it is used instead.

What is not covered

• The timing experiment has not been run yet. It needs a multi-node GPU cluster. The
  results table in the runbook is intentionally empty.
• No integration path has executed against real Megatron Bridge. The tests stub
  it, so a first real run should be expected to surface friction.
• pyproject.toml still pins megatron_core==0.13.1 for the megatron extra,
  which is older than what megatron-bridge==0.6.0 requires. Resolution between the
  two extras is the most likely first snag. Three symbols that moved between mcore
  releases (unwrap_model, is_graph_safe_cuda_rng_tracker,
  convert_cuda_rng_state) are already resolved defensively.

Notes for reviewers

• notes/checkpointing-integration-research.md documents everything read out of
  Megatron Bridge and NeMo RL source — the protocol, the allowlist, which hooks are
  injectable and which aren't, and three places the Bridge docs are stale. GEMINI.md
  points at it. It's the file to read first if any of the above looks surprising.
• bridge_state.py rebuilds the state dict that Bridge assembles inside
  save_checkpoint, using Bridge's own helpers. Three of those are private
  (_build_sharded_state_dict_metadata, _load_model_state_dict,
  _clean_metadata_for_serialization) and are resolved with getattr so a Bridge
  upgrade produces a clear runtime error rather than an import failure.
• ML Flashpoint gets its own AsyncCallsQueue, separate from the durable one. This
  isn't tidiness: a single queue finalizes in scheduling order, so fast ML Flashpoint
  finalizations would stack up behind a slow durable save and pin their buffers until
  the pool is exhausted. Same reasoning as MLFlashpointAsyncFinalizableCheckpointIO
  in the NeMo adapter.

Checklist

• Tests pass
• Appropriate changes to documentation are included in the PR