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

GEMINI.md

@@ -3,6 +3,13 @@
 1. If a virtual environment does not exist at `.venv` or is not already activated, create one in `.venv`, prompting the user for confirmation of the command first.
 1. Activate the venv via `source .venv/bin/activate`.
 
+## Reference Notes
+
+1. `notes/checkpointing-integration-research.md` holds source-verified notes on how Megatron Bridge and NeMo RL do
+checkpointing, and where ML Flashpoint hooks into each. Read it before changing
+`src/ml_flashpoint/adapter/megatron_bridge` or `src/ml_flashpoint/adapter/nemo_rl`, and update it when a finding there
+turns out to be wrong or stale.
+
 ## General Rules
 1. Always, always, always reread the code in case other external changes have been made to it.
 Do not assume that it is in the exact same state as it was the last time you read or edited it.

README.md

@@ -12,7 +12,7 @@
 
 A memory-first, lightning-fast, ready-to-use ML checkpointing library.
 
-Adapters for PyTorch DCP, Megatron-LM and NeMo 2.0 are readily available for seamless integration.
+Adapters for PyTorch DCP, Megatron-LM, Megatron Bridge, NeMo 2.0 and NeMo RL are readily available for seamless integration.
 They are built on top of the core checkpointing APIs, which can also be used directly for custom integrations.
 
 If interested in a native integration with another framework, please let us know by creating a [feature request](https://github.com/google/ml-flashpoint/issues/new?template=feature_request.md) or upvoting an [existing one](https://github.com/google/ml-flashpoint/issues?q=is%3Aissue%20state%3Aopen%20label%3Aenhancement)!
@@ -46,6 +46,12 @@ pip install -e .[pytorch]
 # Megatron-LM
 pip install -e .[megatron]
 
+# Megatron Bridge
+pip install -e .[megatron-bridge]
+
+# NeMo RL (installs the Megatron Bridge stack; NeMo RL itself is installed from source)
+pip install -e .[nemo-rl]
+
 # Multiple
 pip install -e .[pytorch,megatron]
 ```

docs/checkpoint-timing-experiment.md

@@ -0,0 +1,186 @@
+# Measuring the checkpoint-time difference
+
+This is the procedure for validating an ML Flashpoint integration by measuring what checkpointing costs the training
+loop, with and without it. It is deliberately short: a handful of steps is enough, because the quantity of interest is
+per-checkpoint wall clock, not convergence.
+
+!!! note
+
+    The numbers below are placeholders. This page describes how to produce them; it does not report a result. Fill in
+    the results table from your own run.
+
+## What is being measured
+
+Megatron Bridge brackets each checkpoint with barriers and logs the elapsed time from
+`megatron.bridge.training.train.save_checkpoint_and_time`:
+
+* `save-checkpoint` — a durable checkpoint.
+* `save-checkpoint-non-persistent` — a non-persistent one, which is the ML Flashpoint checkpoint when the adapter is
+  enabled.
+
+Both are logged through Megatron's timers in **milliseconds**, in one of two shapes depending on
+`logger.timing_log_option`:
+
+```
+(min, max) time across ranks (ms):
+    save-checkpoint ................................: (18450.20, 18512.90)   # minmax (default)
+    save-checkpoint ................................: 18512.90               # max
+```
+
+The parser reads either shape and always keeps the **max** — the slowest rank — because the save is barrier-bracketed
+on both sides, so the whole job waits for that rank. `(18450.20, 18512.90)` therefore contributes one sample of
+**18.51 s**.
+
+For NeMo RL, Bridge's timers never fire at all: NeMo RL calls `save_checkpoint` directly rather than through
+`train.py`. The adapter emits `nemo_rl.save_checkpoint` instead, covering the whole worker save including the blocking
+`maybe_finalize_async_save` that precedes it.
+
+The headline comparison is the mean and max of those timers between two runs that differ only in whether ML Flashpoint
+is enabled.
+
+## Cluster setup
+
+The workload is the prebuilt NVIDIA NeMo RL job for Google Cloud training clusters, described in
+[Run prebuilt workloads](https://docs.cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/training-clusters/run-prebuilt-workloads#nvidia-nemo-rl).
+Follow that page to create the cluster and get a working NeMo RL run first; do not add ML Flashpoint until an unmodified
+run completes and logs checkpoint timings.
+
+Requirements specific to this experiment:
+
+* **At least two nodes.** ML Flashpoint replicates each node's checkpoint objects to a peer, and a single-node run does
+  not exercise that path.
+* **`/dev/shm` sized for the checkpoint.** Each node holds its own shard plus a peer's replica, so size shared memory to
+  at least twice the per-node checkpoint plus headroom. The pods' shared-memory volume default is usually too small.
+* **A durable checkpoint destination** (a GCS mount or a network filesystem) for the baseline arm and for the durable
+  cadence of the ML Flashpoint arm. Both arms must write durable checkpoints to the same kind of destination, or the
+  comparison measures the storage backend rather than the adapter.
+* **The same node pool, model, parallelism and batch size across both arms.** Run them back to back.
+
+## Run configuration
+
+Keep the run short and make it checkpoint often enough to collect several samples:
+
+* 20–30 training steps.
+* Durable checkpoints every 10 steps, giving 2–3 `save-checkpoint` samples per arm.
+* ML Flashpoint checkpoints every 2 steps in the candidate arm, giving ~10 non-persistent samples.
+* `logger.timing_log_level: 0` or higher, so Megatron logs its timers.
+* ML Flashpoint logging at `INFO`.
+
+A run of this length produces single-digit sample counts for the durable timer. Report the max alongside the mean, and
+do not read a small mean difference as significant.
+
+## Arm A — baseline
+
+Run the workload unchanged. For a Megatron Bridge run, leave `custom_manager_class` unset. For a NeMo RL run, leave
+`MLFLASHPOINT_NEMO_RL_ENABLED` unset, so `install_from_env` is a no-op and the worker behaves exactly as upstream.
+
+Capture stdout from every rank; rank 0 carries the timer lines.
+
+```bash
+kubectl logs -f job/<baseline-job> --all-containers --prefix > logs/baseline.log
+```
+
+## Arm B — ML Flashpoint
+
+Same job, same everything, with the adapter enabled.
+
+Megatron Bridge:
+
+```yaml
+checkpoint:
+  save: /gcs/my-run/checkpoints
+  save_interval: 10
+  non_persistent_save_interval: 2
+  non_persistent_ckpt_type: local
+  custom_manager_class: ml_flashpoint.adapter.megatron_bridge.MLFlashpointBridgeCheckpointManager
+```
+
+NeMo RL, as environment variables on the worker pods:
+
+```bash
+MLFLASHPOINT_NEMO_RL_ENABLED=true
+MLFLASHPOINT_NEMO_RL_MODE=replace
+MLFLASHPOINT_NEMO_RL_DURABLE_EVERY_N_SAVES=5
+MLFLASHPOINT_BASE_CONTAINER=/dev/shm/ml_flashpoint/${JOB_ID}
+```
+
+`replace` is the mode that shows the difference: it lets most checkpoints go to memory alone. `augment` keeps every
+durable write and therefore cannot make the loop faster — use it to check correctness, not to measure a speedup.
+
+```bash
+kubectl logs -f job/<flashpoint-job> --all-containers --prefix > logs/flashpoint.log
+```
+
+## Comparing
+
+```bash
+scripts/benchmarks/parse_checkpoint_timings.py \
+    --label baseline logs/baseline.log --output baseline.json
+
+scripts/benchmarks/parse_checkpoint_timings.py \
+    --label flashpoint logs/flashpoint.log --output flashpoint.json
+
+scripts/benchmarks/compare_checkpoint_timings.py \
+    --baseline baseline.json --candidate flashpoint.json
+```
+
+Which prints one row per timer name, with the sample count, the mean in each arm, and the delta as an absolute change,
+a percentage and a speedup:
+
+```
+Checkpoint timing: flashpoint vs baseline
+
+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
+```
+
+Add `--json` for a machine-readable form.
+
+## Reading the result
+
+Read the table **down the rows, then across the arms** — and be aware that the tool's own `mean delta` column does not
+show the headline number, for the reason below.
+
+**Row 1, `save-checkpoint`, is the control.** 18.306 s vs 18.402 s: durable checkpoints cost the same in both arms.
+That is the expected and desired result. The adapter does not touch the durable path, so any real difference here means
+the two runs differed in something else — a different node pool, a cold storage cache, contention — and everything else
+in the table should be distrusted until that is explained.
+
+**Row 2, `save-checkpoint-non-persistent`, is the new work.** It exists only in the ML Flashpoint arm, because the
+baseline has no non-persistent cadence at all. Hence `-` for baseline and `n/a` for the delta: the tool compares
+like-named timers across arms, and there is nothing to subtract from.
+
+**The headline number is the cross-row comparison the tool cannot 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
+```
+
+That is the claim: on the steps where ML Flashpoint now holds the checkpoint, the loop stalls for ~0.6 s instead of
+~18 s. It 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 pure added cost — correct, but not faster.
+
+Two further checks before treating the integration as validated:
+
+* **Sample counts are small.** Three durable samples per arm is enough to spot an order-of-magnitude difference, not a
+  5% one. Read `max_s` in the JSON alongside the mean.
+* **A fast checkpoint is not the whole claim — recovery has to work.** Confirm the ML Flashpoint arm logs
+  `Recovered from ML Flashpoint checkpoint` after a deliberate restart.
+
+## Results
+
+Fill this in from your own run.
+
+| Timer | Baseline mean | ML Flashpoint mean | Delta |
+|---|---|---|---|
+| `save-checkpoint` | | | |
+| `save-checkpoint-non-persistent` | n/a | | |
+| `nemo_rl.save_checkpoint` | | | |
+
+Record alongside it: node count, GPUs per node, model, parallelism, per-rank checkpoint size, durable destination, and
+the commit of each repository involved.

docs/user-guide.md

@@ -232,6 +232,153 @@ else:
     )
 ```
 
+### Megatron Bridge
+
+Code: See the [`ml_flashpoint.adapter.megatron_bridge`](https://github.com/google/ml-flashpoint/tree/main/src/ml_flashpoint/adapter/megatron_bridge) package.
+
+Megatron Bridge lets a run replace its checkpointing implementation through
+[`CheckpointConfig.custom_manager_class`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/training/checkpointing.md#custom-checkpoint-manager).
+ML Flashpoint ships a manager for that hook which splits the two cadences Megatron Bridge already distinguishes:
+
+| Checkpoint | Cadence | Written by | Durability |
+|---|---|---|---|
+| Persistent | `save_interval` | Megatron Bridge, unchanged | Durable, wherever `save` points |
+| Non-persistent | `non_persistent_save_interval` | ML Flashpoint | Node-local memory, replicated to a peer node |
+
+Bridge takes the non-persistent branch only on steps that are *not* also persistent-checkpoint steps, so the two never collide.
+
+!!! warning
+
+    ML Flashpoint checkpoints are a fast recovery tier, not a replacement for durable ones.
+    They survive a process or node failure inside a run; they do not survive losing the cluster.
+    Keep `save` and `save_interval` configured.
+
+#### Configuration
+
+```python
+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,
+    async_save=True,
+    ckpt_format="torch_dist",  # The only format the ML Flashpoint strategies support.
+)
+
+# Sets custom_manager_class, non_persistent_ckpt_type="local" and
+# non_persistent_save_interval, and registers `ml_flashpoint` with the Megatron
+# Bridge import allowlist.
+mlf_bridge.enable(checkpoint, non_persistent_save_interval=20)
+```
+
+To configure it from YAML instead, set the fields directly and register the allowlist prefix before
+`megatron.bridge.training.setup` runs — Megatron Bridge rejects a `custom_manager_class` outside its allowlist:
+
+```yaml
+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
+```
+
+```python
+import ml_flashpoint.adapter.megatron_bridge as mlf_bridge
+
+mlf_bridge.register_with_megatron_bridge()
+```
+
+#### ML Flashpoint settings
+
+Megatron Bridge constructs the manager with only its own `CheckpointConfig`, so ML Flashpoint's own knobs come from
+either an explicit registration or `MLFLASHPOINT_*` environment variables:
+
+```python
+from ml_flashpoint.adapter.megatron_bridge import MLFlashpointBridgeConfig, configure
+
+configure(
+    MLFlashpointBridgeConfig(
+        base_container="/dev/shm/ml_flashpoint/job-145",
+        write_thread_count=2,
+    )
+)
+```
+
+| Environment variable | Default | Meaning |
+|---|---|---|
+| `MLFLASHPOINT_BRIDGE_ENABLED` | `True` | Set to `false` to make the manager a pass-through to Megatron Bridge. |
+| `MLFLASHPOINT_BASE_CONTAINER` | `/dev/shm/ml_flashpoint` | Node-local, memory-backed base directory holding one child container per checkpoint. |
+| `MLFLASHPOINT_ASYNC_SAVE` | `True` | Keep saves off the training critical path. |
+| `MLFLASHPOINT_WRITE_THREAD_COUNT` | `1` | Writer threads per rank. |
+| `MLFLASHPOINT_INITIAL_WRITE_BUFFER_SIZE_BYTES` | 16 GiB | Initial per-buffer size. Raise it if per-rank checkpoint data is larger. |
+| `MLFLASHPOINT_USE_OPTIMIZED_SAVE` | `True` | Zero-copy tensor writes. |
+| `MLFLASHPOINT_USE_CACHED_CKPT_STRUCTURE` | `False` | Reuse the save plan across steps. Only safe with a constant checkpoint structure. |
+| `MLFLASHPOINT_USE_FULLY_PARALLEL_WRAPPER` | `True` | Spread checkpoint data evenly across ranks. |
+| `MLFLASHPOINT_KEEP_CHECKPOINTS_ON_FINALIZE` | `False` | Keep the container after training ends instead of releasing node memory. |
+
+The base container should be unique per job run but sticky across restarts of the same job, exactly as for the NeMo
+adapter above.
+
+#### Recovery
+
+The manager registers itself in Megatron Bridge's `checkpointing_context` under `local_checkpoint_manager`, which is
+what `megatron.bridge.training.setup` consults to decide whether to attempt a resume. On resume it prefers the newest
+recoverable ML Flashpoint container and falls back to Megatron Bridge's own load path when there is none, or when the
+in-memory read fails.
+
+Because ML Flashpoint containers are node-local, recovery expects the same nodes; missing objects are pulled from the
+peer that holds the replica.
+
+### NeMo RL
+
+Code: See the [`ml_flashpoint.adapter.nemo_rl`](https://github.com/google/ml-flashpoint/tree/main/src/ml_flashpoint/adapter/nemo_rl) package.
+
+!!! note
+
+    NeMo RL builds its Megatron training state with Megatron Bridge but drives checkpointing itself: its
+    `MegatronPolicyWorker` calls `megatron.bridge.training.checkpointing.save_checkpoint` directly rather than going
+    through `create_checkpoint_manager`, so `custom_manager_class` is never consulted. There is therefore no
+    configuration-only way to enable ML Flashpoint for a NeMo RL run; the adapter attaches to the worker instead.
+
+`install_into_worker` wraps `MegatronPolicyWorker.save_checkpoint` on a worker instance, after the worker has finished
+initializing (`mcore_state`, `model` and the process group are all live):
+
+```python
+from ml_flashpoint.adapter import nemo_rl as mlf_nemo_rl
+
+mlf_nemo_rl.install_into_worker(worker, mode=mlf_nemo_rl.MODE_AUGMENT)
+```
+
+Two modes are available:
+
+* `MODE_AUGMENT` (default) — every NeMo RL checkpoint is still written durably, and an ML Flashpoint checkpoint is
+  written alongside it. Faster recovery, unchanged durability.
+* `MODE_REPLACE` with `durable_every_n_saves=N` — only every N-th checkpoint is written durably; the rest go to
+  ML Flashpoint alone. This is what removes checkpoint stalls from the RL loop.
+
+An ML Flashpoint failure never fails the durable write: it is logged and the original save proceeds.
+
+For A/B experiments, `install_from_env` makes both arms share one launch command:
+
+```python
+from ml_flashpoint.adapter import nemo_rl as mlf_nemo_rl
+
+mlf_nemo_rl.install_from_env(worker)
+```
+
+| Environment variable | Default | Meaning |
+|---|---|---|
+| `MLFLASHPOINT_NEMO_RL_ENABLED` | `False` | Master switch. When unset, the run is a plain NeMo RL run. |
+| `MLFLASHPOINT_NEMO_RL_MODE` | `augment` | `augment` or `replace`. |
+| `MLFLASHPOINT_NEMO_RL_DURABLE_EVERY_N_SAVES` | `1` | In `replace` mode, how often to still write durably. |
+
+The `MLFLASHPOINT_*` settings from the Megatron Bridge section above apply here too.
+
+To restore, call `MLFlashpointNeMoRLCheckpointer.load(...)` from the worker's setup path. Unlike the Megatron Bridge
+manager, it never falls back to the durable load path — NeMo RL owns that decision.
+
 ### PyTorch DCP
 
 Code: See the [`ml_flashpoint.adapter.pytorch`](https://github.com/google/ml-flashpoint/tree/main/src/ml_flashpoint/adapter/pytorch) package.

mkdocs.yml

@@ -21,6 +21,7 @@ nav:
   - "Home": README.md
   - "Getting Started": user-guide.md
   - "Overview": overview.md
+  - "Checkpoint Timing Experiment": checkpoint-timing-experiment.md
   # Keep Troubleshooting at the end.
   - "Troubleshooting": troubleshooting.md