diff --git a/GEMINI.md b/GEMINI.md index ae845bc..73ec5c4 100644 --- a/GEMINI.md +++ b/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. diff --git a/README.md b/README.md index 1b2ebc6..46b075e 100644 --- a/README.md +++ b/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] ``` diff --git a/docs/checkpoint-timing-experiment.md b/docs/checkpoint-timing-experiment.md new file mode 100644 index 0000000..7de1a79 --- /dev/null +++ b/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/ --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/ --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. diff --git a/docs/user-guide.md b/docs/user-guide.md index 3d50b2d..e385259 100644 --- a/docs/user-guide.md +++ b/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. diff --git a/mkdocs.yml b/mkdocs.yml index a651afe..dcb377e 100644 --- a/mkdocs.yml +++ b/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 diff --git a/notes/checkpointing-integration-research.md b/notes/checkpointing-integration-research.md new file mode 100644 index 0000000..46ae755 --- /dev/null +++ b/notes/checkpointing-integration-research.md @@ -0,0 +1,312 @@ +# Checkpointing integration research notes + +Reference notes for anyone (human or agent) extending ML Flashpoint into the Megatron Bridge / NeMo RL stack. +Everything here was read out of source, not documentation, on the dates below. Docs for these projects lag the code; +where they disagreed, the code won, and the disagreement is called out. + +**Sources read** + +| Repo | Ref | What was read | +|---|---|---| +| `google/ml-flashpoint` | `97f5c54` (main) | `src/ml_flashpoint/adapter/{megatron,nemo,pytorch}`, `src/ml_flashpoint/core` | +| `NVIDIA-NeMo/Megatron-Bridge` | `main` @ 2026-09-05 | `src/megatron/bridge/training/{checkpointing,train,setup,config}.py`, `utils/instantiate_utils.py`, `docs/training/checkpointing.md` | +| `NVIDIA-NeMo/RL` | `main` @ 2026-09-05 | `nemo_rl/models/megatron/setup.py`, `nemo_rl/models/policy/workers/megatron_policy_worker.py`, `nemo_rl/utils/checkpoint.py`, `nemo_rl/algorithms/*` | + +Line numbers are from the refs above and will drift. Treat them as "look near here", not as addresses. + +--- + +## 1. The one-paragraph version + +Megatron Bridge added a public `CheckpointManager` protocol that a run can swap in via +`CheckpointConfig.custom_manager_class`. That is the right hook for ML Flashpoint, and the Megatron Bridge adapter uses +it. **NeMo RL does not use it.** NeMo RL builds its Megatron state with Bridge but calls Bridge's *functional* +`save_checkpoint` directly from its policy worker, so `custom_manager_class` is never read anywhere in that repo. A +NeMo RL integration therefore has to attach to the worker, not to the config. + +--- + +## 2. Megatron Bridge + +### 2.1 The custom checkpoint manager hook + +`megatron/bridge/training/checkpointing.py`: + +| Symbol | ~Line | Notes | +|---|---|---| +| `CheckpointSaveContext` | 622 | `state, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far, train_data_iterator, non_persistent_ckpt, pg_collection, module_name` | +| `CheckpointLoadContext` | 648 | `state, model, optimizer, opt_param_scheduler, strict, skip_load_to_model_and_opt, pg_collection, module_name` | +| `CheckpointManager` (Protocol) | 673 | `__init__(checkpoint_config)`, `save(ctx, callback_manager)`, `load(ctx) -> (int, int)`, `finalize_async_saves(state, blocking, terminate)` | +| `DefaultCheckpointManager` | 720 | Wraps the functional `save_checkpoint` / `load_checkpoint`. Owns `checkpointing_context`. | +| `create_checkpoint_manager` | 815 | Factory. Imports `custom_manager_class`, checks the protocol, and checks `save` accepts `(ctx, callback_manager)`. | + +**Gotcha — the import allowlist.** `create_checkpoint_manager` calls +`megatron.bridge.utils.instantiate_utils._validate_target_prefix`, which rejects any target whose module prefix is not +allowlisted. The defaults cover `megatron.*`, `torch.*`, `transformers.*`, `nvidia.*`, `numpy.*`, `nemo.*` — **not** +`ml_flashpoint`. Call `megatron.bridge.utils.instantiate_utils.register_allowed_target_prefix("ml_flashpoint")` before +the factory runs or instantiation raises `InstantiationException`. The same validator also rejects any target with a +private (underscore-prefixed) path segment. + +### 2.2 The two checkpoint cadences + +`megatron/bridge/training/train.py::checkpoint_and_decide_exit` (~1425): + +```python +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) +``` + +It is an `elif`, so a step that is a durable-checkpoint step never also takes the non-persistent branch. This is what +makes "durable on `save_interval`, ML Flashpoint on `non_persistent_save_interval`" collision-free without any extra +skip logic (contrast the NeMo 2.0 adapter, which needs `skip_every_n_steps`). + +`save_checkpoint_and_time` (~1302) is also where the timings come from: + +```python +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]) +``` + +Barriers on both sides, so the logged value is the slowest rank — i.e. what the loop actually paid. These two timer +names are the measurement surface for any checkpoint-time experiment. Megatron logs them in **milliseconds** as +`name ....: (min, max)`. + +`finalize_async_saves` is called from the train loop at ~394 (`blocking=False`, every iteration), ~816 +(`blocking=True, terminate=False`) and ~845 / `_finish_train` (`blocking=True, terminate=True`). + +### 2.3 `non_persistent_ckpt_type` and the local-checkpoint duck type + +`save_checkpoint` (~1201) decides the checkpoint type: + +* `non_persistent_ckpt and non_persistent_ckpt_type == "local"` → `CheckpointType.LOCAL`, and + `save_dir = checkpointing_context["local_checkpoint_manager"].local_ckpt_dir` +* `non_persistent_ckpt and non_persistent_ckpt_type == "global"` → `CheckpointType.GLOBAL` into a `non_persistent/` + subdir +* otherwise → `CheckpointType.GLOBAL` + +Bridge only ever touches `checkpointing_context["local_checkpoint_manager"]` when +`non_persistent_ckpt_type == "local"`. The places it does: + +| Where | ~Line | Call | +|---|---|---| +| `save_checkpoint` | 1209 | `.local_ckpt_dir` | +| `save_checkpoint` | 1496 | `.save(state_dict_for_save, step, is_async=...)` | +| `load_checkpoint` | 3279 | `.local_ckpt_dir` (LayerWise optimizer only) | +| `_get_non_persistent_iteration` | 3592 | `.find_latest()` | +| `_load_non_persistent_base_checkpoint` | 3640 | `.load()` | +| `setup._should_load_checkpoint` | 160 | `.find_latest() != -1` | + +That last one is the important one: **`megatron/bridge/training/setup.py::_should_load_checkpoint` is how a run decides +whether to attempt a resume at all.** It reads `getattr(checkpoint_manager, "checkpointing_context", {})`. A custom +manager that wants its own checkpoints to trigger a resume must expose a `checkpointing_context` property containing a +`local_checkpoint_manager` with a `find_latest()` that returns a step (or `-1`). ML Flashpoint's +`MLFlashpointLocalCheckpointIndex` is exactly that duck type, and nothing more — it deliberately does not implement +`save()`/`load()`, because Bridge's own local paths expect an NVRx `MCoreTensorAwareStateDict` container that +ML Flashpoint does not produce. Before delegating to Bridge, `disable()` it so `find_latest()` reports `-1`. + +`init_checkpointing_context` (~3409) raises if `non_persistent_ckpt_type == "local"` and `nvidia_resiliency_ext` is not +installed. A custom manager that provides its own local checkpointing should catch that and continue with `{}`. + +### 2.4 Where strategies can and cannot be injected + +* **Save: injectable.** `save_checkpoint` (~1402) does + `if checkpointing_context is not None and "save_strategy" in checkpointing_context: save_strategy = ...` before + building a `TorchDistSaveShardedStrategy`. Pre-seeding `checkpointing_context["save_strategy"]` works, and works for + *any* caller of `save_checkpoint`, including NeMo RL. +* **Load: not injectable.** Both `_load_global_dist_base_checkpoint` (~3689) and `_load_model_weights_from_checkpoint` + (~2420) construct `TorchDistLoadShardedStrategy()` unconditionally. The `checkpointing_context["load_strategy"]` key + is *written* there but never read back as an override. A custom load path has to call + `megatron.core.dist_checkpointing.load(...)` itself with its own `sharded_strategy`. + +**Why the save-strategy hook is not enough on its own for a local checkpointer.** `dist_checkpointing.save` writes +`common.pt` via the common strategy on **global rank 0 only**, and writes into the directory the caller passed. For a +node-local memory checkpoint you need common state on *every node* (or no node but rank 0's can recover) and you need +the container under the node-local base path, not the durable one. ML Flashpoint's +`adapter/megatron/save_utils.py::save_local_aware_megatron_checkpoint` exists precisely to solve the first half: it +splits with `mcore_state_dict_utils.save_preprocess` and `torch.save`s the common part on every +`torch.distributed.get_node_local_rank() == 0`. + +### 2.5 Rebuilding the state dict outside `save_checkpoint` + +A custom manager that does not delegate has to assemble what `save_checkpoint` assembles. The pieces, all in +`checkpointing.py`: + +* `get_rng_state(data_parallel_random_init, ckpt_format, *, pg_collection, module_name)` (~525) — gated on + `ckpt_cfg.save_rng`. +* `get_rerun_state_machine().state_dict(data_iterator=..., ckpt_format=...)`. +* `_build_sharded_state_dict_metadata(use_distributed_optimizer, ckpt_cfg)` (~3991) — **private**. Then + `metadata["dp_cp_group"] = pg_collection.dp_cp`. +* `generate_state_dict(ckpt_cfg, model, optimizer, opt_param_scheduler, rng_state, iteration=..., optim_sd_kwargs=dict(metadata=...), model_sd_kwargs=dict(metadata=...), rerun_state=..., pg_collection=...)` (~2180). +* For a **load** skeleton, the same call with `optim_sd_kwargs=dict(metadata=..., is_loading=True)` and no `iteration`. + +`dp_cp_group` is a `ProcessGroup` and cannot be pickled — strip it before persisting the metadata. Do not rely on +`megatron.core.dist_checkpointing.utils._clean_metadata_for_serialization` to catch it; drop the key explicitly. + +Local checkpoints have no `latest_train_state.pt` next to them, so Bridge embeds `state_dict["train_state_metadata"] = +train_state.state_dict()` for `CheckpointType.LOCAL` (~1487) and restores from it on load (~3175). Any custom local +format should do the same, plus carry cumulative FLOPs, since `load` must return +`(step, num_floating_point_operations_so_far)`. + +Applying a loaded state dict mirrors `load_checkpoint` ~3170–3390: `set_checkpoint_version`, restore `TrainState`, +`update_num_microbatches`, `_load_model_state_dict` (**private**, ~2779) per chunk (`"model"` for one chunk, `"model%d"` +for many, skipping absent keys = empty PP stages), `optimizer.load_state_dict` under `torch.no_grad()`, scheduler from +`"lr_scheduler"` or `"opt_param_scheduler"`, rerun state, then RNG. + +**Version drift to guard for.** These moved between Megatron Core releases and should be resolved defensively: + +* `unwrap_model` — in `megatron.core.utils` on the mcore that Bridge 0.6 requires; in `megatron.training.utils` on + older ones (e.g. mcore 0.13.1 has it in neither). +* `tensor_parallel.is_graph_safe_cuda_rng_tracker` / `tensor_parallel.convert_cuda_rng_state` — absent on older mcore. + Fall back to setting tracker states directly. +* `_build_sharded_state_dict_metadata`, `_load_model_state_dict`, `_clean_metadata_for_serialization` are private and + can move; `getattr` them and raise a clear error rather than failing at import. + +### 2.6 Config surface + +`CheckpointConfig` (`megatron/bridge/training/config.py` ~498) extends Megatron-LM's. Fields that matter here: +`save`, `load`, `save_interval`, `non_persistent_save_interval`, `non_persistent_ckpt_type`, +`non_persistent_local_ckpt_dir`, `async_save`, `async_strategy` (`"nvrx"` | `"mcore"`), `ckpt_format` +(`"torch_dist"` | `"fsdp_dtensor"` | `"torch"`), `ckpt_assume_constant_structure`, `most_recent_k`, +`fully_parallel_save`, `custom_manager_class`. + +**The factory passes only `CheckpointConfig` to the custom manager.** There is nowhere in Bridge config to put +adapter-specific settings, so they have to come from a module-level registration call or environment variables. + +### 2.7 Where the docs are stale + +`docs/training/checkpointing.md` §"Implementing a Custom Manager" is broadly right but: + +* it does not mention the `_validate_target_prefix` allowlist at all, which is the first thing a third-party manager + hits; +* its `CheckpointSaveContext` / `CheckpointLoadContext` tables omit `pg_collection` and `module_name`, both of which + the real dataclasses carry and `DefaultCheckpointManager` forwards; +* its example `save()` signature omits `pg_collection=` and `module_name=` in the `save_checkpoint` call. + +--- + +## 3. NeMo RL + +### 3.1 The headline finding + +`grep -r "custom_manager_class\|create_checkpoint_manager"` across the whole NeMo RL repo (`.py`, `.yaml`, `.md`) +returns **zero hits**. NeMo RL uses Megatron Bridge for model/optimizer construction and for the checkpoint *functions*, +but never for the checkpoint *manager*. + +What it actually does: + +* `nemo_rl/models/megatron/setup.py:1858` — `checkpointing_context = init_checkpointing_context(megatron_cfg.checkpoint)`, + stored on `ModelAndOptimizerState` and then on the worker as `self.checkpointing_context` + (`megatron_policy_worker.py:586`). +* `nemo_rl/models/policy/workers/megatron_policy_worker.py:3714` — `MegatronPolicyWorker.save_checkpoint(weights_path, + optimizer_path=None)`: + 1. `maybe_finalize_async_save(..., blocking=True)` — blocks on the *previous* save, + 2. onloads model (and optimizer) to CUDA, `torch.cuda.synchronize()`, + 3. temporarily overwrites `self.mcore_state.cfg.checkpoint.save = weights_path`, + 4. calls Bridge's functional `save_checkpoint(state=..., model=[self.model], ..., checkpointing_context=self.checkpointing_context)`, + 5. if sync, `maybe_finalize_async_save(..., blocking=True)` again, + 6. restores `cfg.checkpoint.save` in a `finally`. +* `MegatronPolicyWorker.load_checkpoint` (~3871) **raises `NotImplementedError`** — resume happens only through the + worker's init path. + +Consequences for an integration: + +1. `custom_manager_class` is inert. Do not ship a NeMo RL story that depends on it. +2. `checkpointing_context["save_strategy"]` *would* be honoured (it flows into Bridge's `save_checkpoint`), but see + §2.4 — that puts the container under NeMo RL's durable `weights_path` and writes `common.pt` on rank 0 only, which + is wrong for a node-local checkpointer. +3. The worker attributes needed to build a Bridge `CheckpointSaveContext` are all present and stable: + `worker.mcore_state` (a Bridge `GlobalState`), `worker.model`, `worker.optimizer`, `worker.scheduler`. Wrapping + `worker.save_checkpoint` and driving a Bridge-shaped context from those is the least-coupled option, and is what + `ml_flashpoint.adapter.nemo_rl.install_into_worker` does. + +### 3.2 NeMo RL has no local/non-persistent cadence + +`nemo_rl/utils/checkpoint.py::CheckpointingConfig` has `save_period`, `ft_save_period`, `ft_keep_latest_k`, +`keep_top_k`, `checkpoint_must_save_by`. In the algorithm loops (`grpo.py` ~3765, and the same shape in `ppo.py`, +`dpo.py`, `distillation.py`, `single_controller.py`): + +```python +should_save_by_step = ( + is_last_step + or early_stop_message is not None + or (total_steps + 1) % checkpointing["save_period"] == 0 + or (ft_save_period is not None and (total_steps + 1) % ft_save_period == 0) +) +``` + +`ft_save_period` checkpoints go through the **same** `save_checkpoint` to the **same** `step_N` layout; they differ only +in retention. So there is no upstream notion of "cheap crash-recovery checkpoint" to map ML Flashpoint onto — unlike +Bridge's `non_persistent_save_interval`. Adding one means either a fork/upstream change, or deciding per-save on the +adapter side (which is why the adapter offers `augment` vs `replace` + `durable_every_n_saves`). + +### 3.3 Measurement surface for NeMo RL + +Bridge's `save_checkpoint_and_time` timers do **not** fire for NeMo RL, because NeMo RL calls `save_checkpoint` +directly and not through `train.py`. So `save-checkpoint` / `save-checkpoint-non-persistent` will be absent from a NeMo +RL log. The adapter therefore wraps the worker method with `log_execution_time(name="nemo_rl.save_checkpoint")`, and +that is the timer to compare across arms. Note it includes the blocking `maybe_finalize_async_save` for the *previous* +save, which is a fair thing to measure — that stall is real and is charged to the RL loop. + +--- + +## 4. ML Flashpoint side + +### 4.1 What already existed + +`src/ml_flashpoint/adapter/megatron/` is framework-agnostic and reusable as-is: + +* `save_strategies.MLFlashpointMegatronAsyncSaveStrategy` — an `AsyncSaveShardedStrategy`. Takes a + `MemoryStorageWriter`; `async_save(sharded_state_dict, checkpoint_dir)` returns a Megatron `AsyncRequest` that the + **caller must schedule**, and whose `preload_fn` must run before/at scheduling. It writes a stub + `metadata.json` = `{"sharded_backend": ""}` into `checkpoint_dir` purely to satisfy Megatron's loader validation + (all the checks against it are no-ops). +* `load_strategies.MLFlashpointMegatronLoadStrategy` — a `LoadShardedStrategy` used via + `mcore_dist_checkpointing.load(..., sharded_strategy=, common_strategy=TorchCommonLoadStrategy())`. +* `save_utils.save_local_aware_megatron_checkpoint` — the node-local `common.pt` writer described in §2.4. Swallows + save exceptions and returns `None`. + +Both strategies only handle `torch_dist`-shaped sharded state dicts. `fsdp_dtensor` and legacy `torch` are out. + +### 4.2 Per-rank singletons + +The NeMo 2.0 adapter's `wrapper_util.py` is the reference for what has to be constructed once per rank, and the +Megatron Bridge adapter's `runtime.py` mirrors it: + +`BufferPoolConfig(pool_dir_path=/buffer_pool, rank, num_buffers=threads*2, buffer_size)` → +`CheckpointObjectManager` → `ReplicationManager().initialize(...)` (a **collective**: needs +`torch.distributed` up on all ranks) → `DefaultMLFlashpointCheckpointSaver` → `MemoryStorageWriter` → +save strategy; and `DefaultMLFlashpointCheckpointLoader` → load strategy. Optionally wrapped in +`FullyParallel{Save,Load}StrategyWrapper`. Plus an `AsyncCallsQueue(persistent=True)`. + +Two non-obvious details carried over from the NeMo adapter, both load-bearing: + +* The `torch.multiprocessing` context must be **`spawn`**, not `fork`. A forked `SyncManager` inherits the CUDA + context; if the trainer is SIGKILLed (NVRx in-job restart) the orphan pins GPU memory and the restart OOMs. +* On teardown, after closing the queue, monkeypatch `queue.persistent_caller.close = lambda: None`. + `PersistentAsyncCaller.__del__` calls `close()` → `torch.distributed.get_rank()`, which crashes at interpreter + shutdown once the process group is gone. + +### 4.3 Separate async queues + +`MLFlashpointAsyncFinalizableCheckpointIO` in the NeMo adapter keeps ML Flashpoint's `AsyncCallsQueue` separate from +the durable one, and the Bridge adapter does the same. The reason is not tidiness: a single queue finalizes in +*scheduling* order, not completion order, so fast ML Flashpoint finalizations queue behind a slow durable save while +new ones keep getting scheduled — the buffers stay pinned and the pool OOMs. + +--- + +## 5. Things worth re-checking before trusting this + +* None of the Megatron Bridge or NeMo RL integration paths have been executed. The unit tests stub + `megatron.bridge.*` (see `tests/adapter/conftest.py`) and megatron-core 0.13.1 — the version this repo pins — is + older than what Bridge 0.6 needs, so a real run needs the `megatron-bridge` extra and a matching mcore. +* `FullyParallelSaveStrategyWrapper` is constructed without a parallelization group (matching the NeMo adapter, which + works). Bridge passes `pg_collection.dp_cp`. If sharding-distribution behaviour looks wrong, this is the first thing + to look at. +* `megatron-bridge` on PyPI is at 0.6.0/0.6.1; `nemo-rl` on PyPI is a `0.0.0` placeholder — NeMo RL is installed from + source, so the `nemo-rl` extra here only pulls the Bridge stack. diff --git a/pyproject.toml b/pyproject.toml index cd44f61..49ee809 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,20 @@ nemo = [ "ml_flashpoint[pytorch]", "nemo_toolkit[all]==2.4.0", ] +# An extra for users who want to use this library with its Megatron Bridge adapter. +# Installed via: `pip install ml-flashpoint[megatron-bridge]` +megatron-bridge = [ + "megatron-bridge==0.6.0", + "ml_flashpoint[megatron,pytorch]", +] +# An extra for users who want to use this library with NeMo RL. +# NeMo RL itself is not on PyPI as an installable distribution (the PyPI +# `nemo-rl` entry is a placeholder); it is installed from source, so this extra +# only pulls the Megatron Bridge stack that the NeMo RL adapter builds on. +# Installed via: `pip install ml-flashpoint[nemo-rl]` +nemo-rl = [ + "ml_flashpoint[megatron-bridge]", +] # An extra for generating the documentation site. # Installed via: `pip install ml-flashpoint[docs]` @@ -106,9 +120,7 @@ dev-nemo = [ # Defines a "dev-nemo-rl" extra for NeMo RL development (typically uses Python 3.12+). # Installed via: `pip install -e .[dev-nemo-rl]` dev-nemo-rl = [ - "ml-flashpoint[dev-nemo]", - # TODO: uncomment below and remove line above when nemo-rl profile is added - #"ml-flashpoint[dev-base,nemo-rl]", + "ml-flashpoint[dev-base,nemo-rl,docs]", ] # Defines a "dev" extra for setting up a development environment. diff --git a/scripts/benchmarks/compare_checkpoint_timings.py b/scripts/benchmarks/compare_checkpoint_timings.py new file mode 100755 index 0000000..089b8dd --- /dev/null +++ b/scripts/benchmarks/compare_checkpoint_timings.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compares two checkpoint-timing reports and prints the delta. + +Consumes the JSON produced by ``parse_checkpoint_timings.py``:: + + compare_checkpoint_timings.py --baseline baseline.json --candidate flashpoint.json + +The headline number is the mean wall-clock cost of ``save-checkpoint`` (or +``nemo_rl.save_checkpoint`` for a NeMo RL run), because that is the part of the +step the training loop blocks on. +""" + +import argparse +import json +import sys +from typing import Any, Optional + +HEADLINE_TIMERS = ( + "save-checkpoint", + "save-checkpoint-non-persistent", + "nemo_rl.save_checkpoint", +) +"""Timers reported first, in this order, when present.""" + + +def _load(path: str) -> dict[str, Any]: + """Reads a timing report. + + Args: + path: Path to a report produced by ``parse_checkpoint_timings.py``. + + Returns: + The parsed report. + """ + with open(path, "r") as handle: + return json.load(handle) + + +def _fmt_delta(baseline: Optional[float], candidate: Optional[float]) -> str: + """Formats the change between two measurements. + + Args: + baseline: The baseline value, or None when the timer is absent. + candidate: The candidate value, or None when the timer is absent. + + Returns: + A human-readable delta. + """ + if baseline is None or candidate is None: + return "n/a" + if baseline == 0: + return "n/a (baseline is 0)" + delta = candidate - baseline + pct = delta / baseline * 100.0 + speedup = baseline / candidate if candidate > 0 else float("inf") + return f"{delta:+.3f}s ({pct:+.1f}%, {speedup:.2f}x)" + + +def compare(baseline: dict[str, Any], candidate: dict[str, Any]) -> list[dict[str, Any]]: + """Builds a per-timer comparison. + + Args: + baseline: The baseline report. + candidate: The candidate report. + + Returns: + One row per timer seen in either report, headline timers first. + """ + base_timers = baseline.get("timers", {}) + cand_timers = candidate.get("timers", {}) + names = sorted(set(base_timers) | set(cand_timers)) + names.sort(key=lambda n: (HEADLINE_TIMERS.index(n) if n in HEADLINE_TIMERS else len(HEADLINE_TIMERS), n)) + + rows = [] + for name in names: + base = base_timers.get(name) + cand = cand_timers.get(name) + rows.append( + { + "timer": name, + "baseline_mean_s": base["mean_s"] if base else None, + "candidate_mean_s": cand["mean_s"] if cand else None, + "baseline_max_s": base["max_s"] if base else None, + "candidate_max_s": cand["max_s"] if cand else None, + "baseline_count": base["count"] if base else 0, + "candidate_count": cand["count"] if cand else 0, + "mean_delta": _fmt_delta(base["mean_s"] if base else None, cand["mean_s"] if cand else None), + "max_delta": _fmt_delta(base["max_s"] if base else None, cand["max_s"] if cand else None), + } + ) + return rows + + +def render(baseline_label: str, candidate_label: str, rows: list[dict[str, Any]]) -> str: + """Renders the comparison as a plain-text table. + + Args: + baseline_label: Name of the baseline arm. + candidate_label: Name of the candidate arm. + rows: Rows from :func:`compare`. + + Returns: + The rendered table. + """ + lines = [ + f"Checkpoint timing: {candidate_label} vs {baseline_label}", + "", + f"{'timer':<44} {'n':>4} {baseline_label[:12]:>13} {candidate_label[:12]:>13} {'mean delta':<28}", + "-" * 100, + ] + for row in rows: + base = "-" if row["baseline_mean_s"] is None else f"{row['baseline_mean_s']:.3f}s" + cand = "-" if row["candidate_mean_s"] is None else f"{row['candidate_mean_s']:.3f}s" + count = max(row["baseline_count"], row["candidate_count"]) + lines.append(f"{row['timer']:<44} {count:>4} {base:>13} {cand:>13} {row['mean_delta']:<28}") + lines.append("") + lines.append("Means are per checkpoint, over the samples found in the logs. A run of only a few steps") + lines.append("produces few samples, so read the max column alongside the mean before concluding.") + return "\n".join(lines) + + +def main(argv: Optional[list[str]] = None) -> int: + """Entry point. + + Args: + argv: Command line arguments, excluding the program name. + + Returns: + Process exit code. + """ + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--baseline", required=True, help="Baseline report JSON.") + parser.add_argument("--candidate", required=True, help="Candidate report JSON.") + parser.add_argument("--json", action="store_true", help="Emit JSON instead of a table.") + args = parser.parse_args(argv) + + baseline = _load(args.baseline) + candidate = _load(args.candidate) + rows = compare(baseline, candidate) + + if args.json: + print( + json.dumps( + { + "baseline_label": baseline.get("label", "baseline"), + "candidate_label": candidate.get("label", "candidate"), + "rows": rows, + }, + indent=2, + ) + ) + else: + print(render(baseline.get("label", "baseline"), candidate.get("label", "candidate"), rows)) + + if not rows: + print("No timers in either report.", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/benchmarks/parse_checkpoint_timings.py b/scripts/benchmarks/parse_checkpoint_timings.py new file mode 100755 index 0000000..ca08020 --- /dev/null +++ b/scripts/benchmarks/parse_checkpoint_timings.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Extracts checkpoint timings from Megatron Bridge / NeMo RL training logs. + +Three sources are recognized, in decreasing order of preference: + +``megatron-timer`` + Megatron's own timer line, emitted by ``save_checkpoint_and_time``. It brackets + the save with barriers, so it is the number that reflects what the training + loop actually paid. Values are milliseconds; both ``timing_log_option`` shapes + are read, and the max (slowest rank) is the sample kept:: + + save-checkpoint ................................: (1234.56, 2345.67) + save-checkpoint ................................: 2345.67 + +``mlf-timer`` + ML Flashpoint's ``log_execution_time`` output, which isolates the adapter's + own cost:: + + MLFlashpointBridgeCheckpointManager.save took 0.1234s + +``nemo-rl-timer`` + NeMo RL's wrapper around the worker save, installed by the adapter:: + + nemo_rl.save_checkpoint took 12.3456s + +Usage:: + + parse_checkpoint_timings.py --label baseline logs/baseline/*.log > baseline.json + parse_checkpoint_timings.py --label flashpoint logs/mlf/*.log > flashpoint.json +""" + +import argparse +import json +import re +import statistics +import sys +from typing import Iterable, Optional + +# Megatron's timers report milliseconds, in one of two shapes depending on +# `logger.timing_log_option`: +# +# minmax (the default) save-checkpoint ....: (1234.56, 2345.67) -> (min, max) across ranks +# max save-checkpoint ....: 2345.67 -> max across ranks +# +# Either way the value taken below is the max, because that is the rank the +# training loop actually waited for. The `all` option prints a different, +# per-rank layout and is not parsed. +_MEGATRON_TIMER = re.compile( + r"(?Psave-checkpoint(?:-non-persistent)?|load-checkpoint)\s*\.*\s*:\s*" + r"(?:\(\s*[0-9.]+\s*,\s*(?P[0-9.]+)\s*\)|(?P[0-9.]+))" +) + +# ml_flashpoint.core.utils.log_execution_time: " took 1.2345s" +_MLF_TIMER = re.compile(r"(?P[A-Za-z_][\w.]*) took (?P[0-9.]+)s") + +_MLF_NAMES_OF_INTEREST = ( + "MLFlashpointBridgeCheckpointManager.save", + "MLFlashpointBridgeCheckpointManager.mlf_save", + "MLFlashpointBridgeCheckpointManager.load", + "MLFlashpointBridgeCheckpointManager.mlf_load", + "MLFlashpointBridgeCheckpointManager.finalize_async_saves", + "MLFlashpointNeMoRLCheckpointer.save", + "nemo_rl.save_checkpoint", + "async_save", +) + + +def parse_lines(lines: Iterable[str]) -> dict[str, list[float]]: + """Collects every recognized timing from a stream of log lines. + + Args: + lines: The log lines to scan. + + Returns: + A mapping of timer name to the observed durations, in seconds. + """ + samples: dict[str, list[float]] = {} + for line in lines: + match = _MEGATRON_TIMER.search(line) + if match: + # Megatron timers are reported in milliseconds, under whichever of the + # two `timing_log_option` shapes the run was configured for. + max_ms = match.group("minmax") or match.group("maxonly") + samples.setdefault(match.group("name"), []).append(float(max_ms) / 1000.0) + continue + match = _MLF_TIMER.search(line) + if match and match.group("name") in _MLF_NAMES_OF_INTEREST: + samples.setdefault(match.group("name"), []).append(float(match.group("seconds"))) + return samples + + +def summarize(samples: dict[str, list[float]]) -> dict[str, dict[str, float]]: + """Reduces raw samples to per-timer statistics. + + Args: + samples: Timer name to durations in seconds. + + Returns: + Timer name to a statistics dictionary. + """ + summary = {} + for name, values in sorted(samples.items()): + ordered = sorted(values) + summary[name] = { + "count": len(ordered), + "total_s": sum(ordered), + "mean_s": statistics.fmean(ordered), + "median_s": statistics.median(ordered), + "min_s": ordered[0], + "max_s": ordered[-1], + # With a handful of steps there are too few samples for a real p95, + # so report the worst observed value alongside the mean instead. + "stdev_s": statistics.stdev(ordered) if len(ordered) > 1 else 0.0, + } + return summary + + +def _read(paths: list[str]) -> Iterable[str]: + """Yields lines from the given files, or stdin when no file is given. + + Args: + paths: Log file paths. + + Yields: + Individual log lines. + """ + if not paths: + yield from sys.stdin + return + for path in paths: + with open(path, "r", errors="replace") as handle: + yield from handle + + +def main(argv: Optional[list[str]] = None) -> int: + """Entry point. + + Args: + argv: Command line arguments, excluding the program name. + + Returns: + Process exit code. + """ + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("logs", nargs="*", help="Log files to parse. Reads stdin when omitted.") + parser.add_argument("--label", required=True, help="Name for this arm of the experiment, e.g. 'baseline'.") + parser.add_argument("--output", help="Write JSON here instead of stdout.") + args = parser.parse_args(argv) + + samples = parse_lines(_read(args.logs)) + if not samples: + print( + "No checkpoint timings found. Confirm the run had logger.timing_log_level >= 0 and that " + "ML Flashpoint logging is at INFO.", + file=sys.stderr, + ) + report = {"label": args.label, "timers": summarize(samples), "raw_samples": samples} + + serialized = json.dumps(report, indent=2) + if args.output: + with open(args.output, "w") as handle: + handle.write(serialized + "\n") + else: + print(serialized) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/ml_flashpoint/adapter/megatron_bridge/__init__.py b/src/ml_flashpoint/adapter/megatron_bridge/__init__.py new file mode 100644 index 0000000..09452d5 --- /dev/null +++ b/src/ml_flashpoint/adapter/megatron_bridge/__init__.py @@ -0,0 +1,106 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ML Flashpoint adapter for Megatron Bridge.""" + +from typing import Optional + +from ml_flashpoint.adapter.megatron_bridge.checkpoint_manager import ( + MLFlashpointBridgeCheckpointManager as MLFlashpointBridgeCheckpointManager, +) +from ml_flashpoint.adapter.megatron_bridge.config import ( + DEFAULT_BASE_CONTAINER as DEFAULT_BASE_CONTAINER, +) +from ml_flashpoint.adapter.megatron_bridge.config import ( + MLFlashpointBridgeConfig as MLFlashpointBridgeConfig, +) +from ml_flashpoint.adapter.megatron_bridge.config import configure as configure +from ml_flashpoint.adapter.megatron_bridge.config import get_config as get_config +from ml_flashpoint.adapter.megatron_bridge.config import reset_configuration as reset_configuration +from ml_flashpoint.adapter.megatron_bridge.runtime import ( + MLFlashpointBridgeRuntime as MLFlashpointBridgeRuntime, +) +from ml_flashpoint.adapter.megatron_bridge.runtime import get_runtime as get_runtime +from ml_flashpoint.adapter.megatron_bridge.runtime import shutdown_runtime as shutdown_runtime +from ml_flashpoint.core.mlf_logging import get_logger + +_LOGGER = get_logger(__name__) + +CUSTOM_MANAGER_CLASS = "ml_flashpoint.adapter.megatron_bridge.MLFlashpointBridgeCheckpointManager" +"""Value to assign to ``CheckpointConfig.custom_manager_class``.""" + +_ALLOWLIST_PREFIX = "ml_flashpoint" + + +def register_with_megatron_bridge() -> None: + """Allows Megatron Bridge to import the ML Flashpoint checkpoint manager. + + Bridge validates ``custom_manager_class`` against an import allowlist that + only covers a fixed set of prefixes, so ``ml_flashpoint`` has to be added + before the manager can be instantiated from config. Idempotent. + """ + try: + from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix + except ImportError: + _LOGGER.warning( + "Could not import megatron.bridge.utils.instantiate_utils.register_allowed_target_prefix; " + "this Megatron Bridge version may reject custom_manager_class='%s'.", + CUSTOM_MANAGER_CLASS, + ) + return + register_allowed_target_prefix(_ALLOWLIST_PREFIX) + _LOGGER.debug("Registered '%s' as an allowed Megatron Bridge target prefix.", _ALLOWLIST_PREFIX) + + +def enable( + checkpoint_config, + non_persistent_save_interval: int, + mlf_config: Optional[MLFlashpointBridgeConfig] = None, +) -> None: + """Points a Megatron Bridge ``CheckpointConfig`` at ML Flashpoint. + + Mutates ``checkpoint_config`` in place so that Bridge saves a fast, + node-local ML Flashpoint checkpoint every ``non_persistent_save_interval`` + steps, while durable checkpoints keep going wherever ``save`` points, on the + existing ``save_interval`` cadence. + + Bridge only takes the non-persistent branch on steps that are *not* also + durable-checkpoint steps, so the two cadences do not collide. + + Args: + checkpoint_config: The Bridge ``CheckpointConfig`` to modify. + non_persistent_save_interval: How often, in steps, to write an ML + Flashpoint checkpoint. Must be positive. + mlf_config: ML Flashpoint settings to register. When omitted, settings + are read from ``MLFLASHPOINT_*`` environment variables. + + Raises: + ValueError: If ``non_persistent_save_interval`` is not positive. + """ + if non_persistent_save_interval < 1: + raise ValueError( + f"non_persistent_save_interval must be a positive integer, got {non_persistent_save_interval}." + ) + if mlf_config is not None: + configure(mlf_config) + + register_with_megatron_bridge() + checkpoint_config.custom_manager_class = CUSTOM_MANAGER_CLASS + checkpoint_config.non_persistent_ckpt_type = "local" + checkpoint_config.non_persistent_save_interval = non_persistent_save_interval + _LOGGER.info( + "Enabled ML Flashpoint for Megatron Bridge: non_persistent_save_interval=%d, config=%s", + non_persistent_save_interval, + get_config(), + ) diff --git a/src/ml_flashpoint/adapter/megatron_bridge/bridge_state.py b/src/ml_flashpoint/adapter/megatron_bridge/bridge_state.py new file mode 100644 index 0000000..2ec65d7 --- /dev/null +++ b/src/ml_flashpoint/adapter/megatron_bridge/bridge_state.py @@ -0,0 +1,446 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Translation between Megatron Bridge contexts and Megatron sharded state dicts. + +Megatron Bridge builds its state dict inside ``save_checkpoint`` / +``load_checkpoint``, which also decide where the checkpoint goes. ML Flashpoint +needs the state dict but not the destination, so the pieces are rebuilt here +using Bridge's own helpers. Anything imported from Bridge that is private is +resolved defensively so that a Bridge upgrade degrades to a clear runtime error +rather than an import failure at module load. +""" + +import importlib +import random +from typing import Any, Optional + +import numpy as np +import torch +from megatron.bridge.training.checkpointing import ( + generate_state_dict, + get_rng_state, + set_checkpoint_version, +) +from megatron.bridge.training.state import TrainState +from megatron.bridge.training.utils.pg_utils import get_pg_collection +from megatron.core import tensor_parallel +from megatron.core.num_microbatches_calculator import update_num_microbatches +from megatron.core.rerun_state_machine import get_rerun_state_machine + +from ml_flashpoint.core.mlf_logging import get_logger + +_LOGGER = get_logger(__name__) + +TRAIN_STATE_KEY = "train_state_metadata" +"""Key under which the Bridge ``TrainState`` is embedded in the state dict. + +Mirrors what Bridge itself does for local (non-persistent) checkpoints: local +checkpoints have no ``latest_train_state.pt`` tracker next to them, so the +counters have to travel inside the checkpoint. +""" + +CONTENT_METADATA_KEY = "content_metadata" +"""Key under which the sharded-state-dict metadata is embedded.""" + +FLOPS_KEY = "num_floating_point_operations_so_far" +"""Key under which cumulative FLOPs are embedded.""" + +DP_CP_GROUP_KEY = "dp_cp_group" +"""Metadata key carrying the data/context-parallel process group. + +Megatron's ``sharded_state_dict`` methods need it, but a process group cannot be +pickled, so it is stripped before the metadata is persisted. +""" + + +def unwrap_model(model): + """Strips the DDP / Float16Module wrappers off each model chunk. + + ``unwrap_model`` moved into ``megatron.core.utils`` only in the Megatron Core + releases that Megatron Bridge 0.6 requires, so it is resolved at call time and + the older ``megatron.training.utils`` location is accepted as well. + + Args: + model: A module or list of module chunks. + + Returns: + The unwrapped chunks, always as a list. + + Raises: + RuntimeError: If neither location provides the helper. + """ + for module_path in ("megatron.core.utils", "megatron.training.utils"): + try: + module = importlib.import_module(module_path) + except ImportError: + continue + impl = getattr(module, "unwrap_model", None) + if impl is not None: + unwrapped = impl(model) + return unwrapped if isinstance(unwrapped, list) else [unwrapped] + raise RuntimeError( + "Could not resolve unwrap_model from megatron.core.utils or megatron.training.utils. " + "Install a Megatron Core version compatible with the Megatron Bridge release in use." + ) + + +def _import_optional(name: str): + """Imports a Bridge symbol that is private and therefore may move. + + Args: + name: Attribute name on ``megatron.bridge.training.checkpointing``. + + Returns: + The attribute, or None when this Bridge version does not expose it. + """ + import megatron.bridge.training.checkpointing as bridge_checkpointing + + return getattr(bridge_checkpointing, name, None) + + +def build_sharded_state_dict_metadata(use_distributed_optimizer: bool, ckpt_cfg) -> dict[str, Any]: + """Builds the metadata Bridge passes to ``sharded_state_dict`` methods. + + Args: + use_distributed_optimizer: Whether the run uses the distributed optimizer. + ckpt_cfg: The Bridge ``CheckpointConfig``. + + Returns: + The metadata dictionary. + + Raises: + RuntimeError: If the installed Megatron Bridge does not expose the helper. + """ + builder = _import_optional("_build_sharded_state_dict_metadata") + if builder is None: + raise RuntimeError( + "megatron.bridge.training.checkpointing._build_sharded_state_dict_metadata is not available. " + "This Megatron Bridge version is not supported by the ML Flashpoint adapter." + ) + return builder(use_distributed_optimizer, ckpt_cfg) + + +def clean_metadata_for_serialization(metadata: dict[str, Any]) -> dict[str, Any]: + """Strips non-serializable entries (e.g. process groups) from metadata. + + Args: + metadata: The sharded-state-dict metadata. + + Returns: + A copy safe to persist, or the input unchanged when the Megatron helper + is unavailable. + """ + try: + from megatron.core.dist_checkpointing.utils import _clean_metadata_for_serialization + + return _clean_metadata_for_serialization(metadata) + except ImportError: + _LOGGER.warning("Could not import _clean_metadata_for_serialization; dropping process groups manually.") + return {k: v for k, v in metadata.items() if not isinstance(v, torch.distributed.ProcessGroup)} + + +def resolve_pg_collection(model: list, pg_collection=None): + """Returns the process group collection to use. + + Args: + model: The (unwrapped) model modules. + pg_collection: An explicit collection from the Bridge context, if any. + + Returns: + The resolved ``ProcessGroupCollection``. + """ + return pg_collection if pg_collection is not None else get_pg_collection(model) + + +def build_save_state_dict( + ctx, + step: int, + num_floating_point_operations_so_far: int, +) -> dict[str, Any]: + """Builds the sharded state dict for an ML Flashpoint save. + + Reproduces the portion of Bridge's ``save_checkpoint`` that assembles state, + and additionally embeds the train state, cumulative FLOPs and content + metadata, none of which have a home on disk for an ML Flashpoint container. + + Args: + ctx: The Bridge ``CheckpointSaveContext``. + step: The training step this checkpoint represents. + num_floating_point_operations_so_far: Cumulative FLOPs to embed. + + Returns: + The state dict to hand to the ML Flashpoint save strategy. + """ + cfg = ctx.state.cfg + ckpt_cfg = cfg.checkpoint + model = unwrap_model(ctx.model) + pg_collection = resolve_pg_collection(model, getattr(ctx, "pg_collection", None)) + + rng_state = None + if ckpt_cfg.save_rng: + rng_state = get_rng_state( + data_parallel_random_init=cfg.rng.data_parallel_random_init, + ckpt_format=ckpt_cfg.ckpt_format, + pg_collection=pg_collection, + module_name=getattr(ctx, "module_name", None), + ) + + rerun_state = get_rerun_state_machine().state_dict( + data_iterator=ctx.train_data_iterator, + ckpt_format=ckpt_cfg.ckpt_format, + ) + + sharded_sd_metadata = build_sharded_state_dict_metadata(cfg.optimizer.use_distributed_optimizer, ckpt_cfg) + # The process group is needed by sharded_state_dict() but must not be persisted. + sharded_sd_metadata[DP_CP_GROUP_KEY] = pg_collection.dp_cp + + state_dict = generate_state_dict( + ckpt_cfg, + model, + ctx.optimizer, + ctx.opt_param_scheduler, + rng_state, + iteration=step, + optim_sd_kwargs=dict(metadata=sharded_sd_metadata), + model_sd_kwargs=dict(metadata=sharded_sd_metadata), + rerun_state=rerun_state, + pg_collection=pg_collection, + ) + + state_dict[TRAIN_STATE_KEY] = ctx.state.train_state.state_dict() + state_dict[FLOPS_KEY] = int(num_floating_point_operations_so_far) + # Drop the process group explicitly rather than trusting the Megatron cleaner + # to recognize it: it is the one entry this adapter adds, and it is the one + # entry that cannot be pickled into common.pt. + persistable_metadata = {k: v for k, v in sharded_sd_metadata.items() if k != DP_CP_GROUP_KEY} + state_dict[CONTENT_METADATA_KEY] = clean_metadata_for_serialization(persistable_metadata) + return state_dict + + +def build_load_state_dict(ctx) -> dict[str, Any]: + """Builds the sharded state dict skeleton used to load an ML Flashpoint save. + + The skeleton mirrors Bridge's local-checkpoint load path: TP/PP are assumed + unchanged (an ML Flashpoint checkpoint is only ever recovered by the same + job with the same parallelism), so no run-config comparison is performed. + + Args: + ctx: The Bridge ``CheckpointLoadContext``. + + Returns: + The sharded state dict to pass to ``dist_checkpointing.load``. + """ + cfg = ctx.state.cfg + ckpt_cfg = cfg.checkpoint + model = unwrap_model(ctx.model) + pg_collection = resolve_pg_collection(model, getattr(ctx, "pg_collection", None)) + + rng_state = None + if ckpt_cfg.load_rng: + rng_state = get_rng_state( + data_parallel_random_init=cfg.rng.data_parallel_random_init, + ckpt_format=ckpt_cfg.ckpt_format, + pg_collection=pg_collection, + module_name=getattr(ctx, "module_name", None), + ) + + rerun_state = get_rerun_state_machine().state_dict( + data_iterator=None, + ckpt_format=ckpt_cfg.ckpt_format, + force=True, + ) + + sharded_sd_metadata = build_sharded_state_dict_metadata(cfg.optimizer.use_distributed_optimizer, ckpt_cfg) + sharded_sd_metadata[DP_CP_GROUP_KEY] = pg_collection.dp_cp + + load_optimizer = ckpt_cfg.load_optim and not ckpt_cfg.finetune + return generate_state_dict( + ckpt_cfg, + model, + ctx.optimizer if load_optimizer else None, + ctx.opt_param_scheduler if load_optimizer else None, + rng_state, + optim_sd_kwargs=dict(metadata=sharded_sd_metadata, is_loading=True), + model_sd_kwargs=dict(metadata=sharded_sd_metadata), + rerun_state=rerun_state, + pg_collection=pg_collection, + ) + + +def restore_train_state(state, state_dict: dict[str, Any]) -> None: + """Restores the Bridge ``TrainState`` from a loaded ML Flashpoint state dict. + + Args: + state: The Bridge ``GlobalState`` to mutate. + state_dict: The loaded state dict. + """ + if TRAIN_STATE_KEY in state_dict: + state.train_state = TrainState() + state.train_state.load_state_dict(state_dict[TRAIN_STATE_KEY]) + else: + _LOGGER.warning("'%s' missing from the checkpoint; training counters reset.", TRAIN_STATE_KEY) + state.train_state = TrainState() + state.train_state.step = state_dict.get("iteration", 0) + + if FLOPS_KEY in state_dict: + state.train_state.floating_point_operations_so_far = state_dict[FLOPS_KEY] + + +def restore_model(model: list, state_dict: dict[str, Any], strict: bool) -> None: + """Loads model weights from a loaded state dict. + + Args: + model: The unwrapped model modules. + state_dict: The loaded state dict. + strict: Whether to enforce strict key matching. + + Raises: + RuntimeError: If the installed Megatron Bridge does not expose the helper. + """ + loader = _import_optional("_load_model_state_dict") + if loader is None: + raise RuntimeError( + "megatron.bridge.training.checkpointing._load_model_state_dict is not available. " + "This Megatron Bridge version is not supported by the ML Flashpoint adapter." + ) + if len(model) == 1: + loader(model[0], state_dict["model"], strict) + return + for i in range(len(model)): + model_key = "model%d" % i + if model_key not in state_dict: + # Empty pipeline stage. + continue + loader(model[i], state_dict[model_key], strict) + + +def restore_optimizer(ctx, state_dict: dict[str, Any]) -> None: + """Loads optimizer and scheduler state from a loaded state dict. + + Args: + ctx: The Bridge ``CheckpointLoadContext``. + state_dict: The loaded state dict. + """ + ckpt_cfg = ctx.state.cfg.checkpoint + if ckpt_cfg.finetune or not ckpt_cfg.load_optim: + return + + optimizer = ctx.optimizer + if ( + not ctx.skip_load_to_model_and_opt + and optimizer is not None + and not getattr(optimizer, "is_stub_optimizer", False) + and "optimizer" in state_dict + ): + # no_grad is required because DistributedOptimizer copies the loaded + # tensors into main params with .copy_(), which rejects leaf Variables + # that require grad. + with torch.no_grad(): + optimizer.load_state_dict(state_dict["optimizer"]) + + if ctx.opt_param_scheduler is not None: + scheduler_state = state_dict.get("lr_scheduler", state_dict.get("opt_param_scheduler")) + if scheduler_state is not None: + ctx.opt_param_scheduler.load_state_dict(scheduler_state) + + +def restore_rerun_state(state_dict: dict[str, Any]) -> None: + """Restores the rerun state machine, logging and continuing on failure. + + Args: + state_dict: The loaded state dict. + """ + if "rerun_state_machine" not in state_dict: + return + try: + get_rerun_state_machine().load_state_dict(state_dict["rerun_state_machine"]) + except Exception: + _LOGGER.exception("Unable to restore the rerun state machine. Skipping.") + + +def restore_rng_state(ctx, state_dict: dict[str, Any], pg_collection) -> None: + """Restores RNG state, logging and continuing on failure. + + Args: + ctx: The Bridge ``CheckpointLoadContext``. + state_dict: The loaded state dict. + pg_collection: The process group collection. + """ + cfg = ctx.state.cfg + if cfg.checkpoint.finetune or not cfg.checkpoint.load_rng or "rng_state" not in state_dict: + return + try: + rng_states = state_dict["rng_state"] + rng_state = rng_states[pg_collection.dp.rank()] if cfg.rng.data_parallel_random_init else rng_states[0] + random.setstate(rng_state["random_rng_state"]) + np.random.set_state(rng_state["np_rng_state"]) + torch.set_rng_state(rng_state["torch_rng_state"]) + torch.cuda.set_rng_state(rng_state["cuda_rng_state"]) + tracker_states = rng_state["rng_tracker_states"] + if not tracker_states: + raise KeyError("rng_tracker_states is empty") + cuda_rng_tracker = tensor_parallel.get_cuda_rng_tracker() + # Graph-safe RNG conversion only exists on the Megatron Core releases that + # ship CUDA-graph-capturable trackers; older builds store states directly. + is_graph_safe = getattr(tensor_parallel, "is_graph_safe_cuda_rng_tracker", None) + convert = getattr(tensor_parallel, "convert_cuda_rng_state", None) + if is_graph_safe is not None and convert is not None: + graph_safe_rng = is_graph_safe(cuda_rng_tracker) + tracker_states = {k: convert(v, to_graphable=graph_safe_rng) for k, v in tracker_states.items()} + cuda_rng_tracker.set_states(tracker_states) + except Exception: + _LOGGER.exception("Unable to restore RNG state from the ML Flashpoint checkpoint. Continuing without it.") + + +def apply_loaded_state(ctx, state_dict: dict[str, Any]) -> tuple[int, int]: + """Applies a loaded ML Flashpoint state dict to the training state. + + Args: + ctx: The Bridge ``CheckpointLoadContext``. + state_dict: The state dict returned by ``dist_checkpointing.load``. + + Returns: + A tuple of (step, cumulative floating point operations). + """ + state = ctx.state + model = unwrap_model(ctx.model) + pg_collection = resolve_pg_collection(model, getattr(ctx, "pg_collection", None)) + + set_checkpoint_version(state_dict.get("checkpoint_version", 0)) + restore_train_state(state, state_dict) + update_num_microbatches(consumed_samples=state.train_state.consumed_train_samples, verbose=True) + + if not ctx.skip_load_to_model_and_opt: + restore_model(model, state_dict, ctx.strict) + restore_optimizer(ctx, state_dict) + restore_rerun_state(state_dict) + restore_rng_state(ctx, state_dict, pg_collection) + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + return state.train_state.step, state.train_state.floating_point_operations_so_far + + +def get_content_metadata(state_dict: dict[str, Any]) -> Optional[dict[str, Any]]: + """Returns the embedded content metadata, if present. + + Args: + state_dict: A loaded state dict. + + Returns: + The content metadata, or None. + """ + return state_dict.get(CONTENT_METADATA_KEY) diff --git a/src/ml_flashpoint/adapter/megatron_bridge/checkpoint_manager.py b/src/ml_flashpoint/adapter/megatron_bridge/checkpoint_manager.py new file mode 100644 index 0000000..289da41 --- /dev/null +++ b/src/ml_flashpoint/adapter/megatron_bridge/checkpoint_manager.py @@ -0,0 +1,464 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A Megatron Bridge ``CheckpointManager`` backed by ML Flashpoint. + +Megatron Bridge lets a run swap in its own checkpoint manager through +``CheckpointConfig.custom_manager_class``. This module provides one that keeps +Bridge's behavior for durable checkpoints and takes over the frequent, +non-persistent ones, which is where ML Flashpoint's memory-first saves pay off. +""" + +import logging +from typing import Any, Optional + +import torch +import torch.distributed as dist +from megatron.bridge.training.checkpointing import ( + init_checkpointing_context, + load_checkpoint, + maybe_finalize_async_save, + save_checkpoint, +) +from megatron.core import dist_checkpointing as mcore_dist_checkpointing +from megatron.core.dist_checkpointing.strategies.common import TorchCommonLoadStrategy + +from ml_flashpoint.adapter.megatron.save_utils import save_local_aware_megatron_checkpoint +from ml_flashpoint.adapter.megatron_bridge import bridge_state +from ml_flashpoint.adapter.megatron_bridge.config import MLFlashpointBridgeConfig, get_config +from ml_flashpoint.adapter.megatron_bridge.local_checkpoint_index import ( + MLFlashpointLocalCheckpointIndex, +) +from ml_flashpoint.adapter.megatron_bridge.runtime import MLFlashpointBridgeRuntime, get_runtime +from ml_flashpoint.core import mlf_logging +from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId +from ml_flashpoint.core.mlf_logging import get_logger +from ml_flashpoint.core.utils import log_execution_time + +_LOGGER = get_logger(__name__) + +SUPPORTED_CKPT_FORMAT = "torch_dist" +"""The only Megatron checkpoint format the ML Flashpoint strategies handle.""" + +LOCAL_NON_PERSISTENT_CKPT_TYPE = "local" +"""``non_persistent_ckpt_type`` value that hands non-persistent saves to this adapter.""" + + +class MLFlashpointBridgeCheckpointManager: + """Routes non-persistent checkpoints to ML Flashpoint, the rest to Megatron Bridge. + + Wire it up with:: + + checkpoint = CheckpointConfig( + 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" + ), + ) + + Megatron Bridge validates ``custom_manager_class`` against an import + allowlist, so ``ml_flashpoint`` has to be registered first. Calling + :func:`ml_flashpoint.adapter.megatron_bridge.register_with_megatron_bridge` + (or importing this package's ``enable`` helper) does that. + + Everything ML Flashpoint needs beyond ``CheckpointConfig`` comes from + :func:`ml_flashpoint.adapter.megatron_bridge.config.get_config`, because + Bridge's factory only passes the checkpoint config. + + Attributes: + checkpoint_config: The Bridge ``CheckpointConfig`` this manager was built with. + """ + + def __init__( + self, + checkpoint_config, + mlf_config: Optional[MLFlashpointBridgeConfig] = None, + runtime: Optional[MLFlashpointBridgeRuntime] = None, + ): + """Initializes the manager. + + The ML Flashpoint runtime is built lazily on first use, because Bridge + constructs the manager before ``torch.distributed`` initialization is + guaranteed to have happened on every code path. + + Args: + checkpoint_config: The Bridge ``CheckpointConfig``. + mlf_config: ML Flashpoint settings. Defaults to the registered or + environment-derived config. + runtime: An already-built runtime. Mainly for tests; production code + should let the manager build (and share) one. + """ + self.checkpoint_config = checkpoint_config + self._mlf_config = mlf_config if mlf_config is not None else get_config() + self._runtime = runtime + self._local_index: Optional[MLFlashpointLocalCheckpointIndex] = None + self._context: dict[str, Any] = init_checkpointing_context_safely(checkpoint_config) + self._enabled = self._resolve_enabled() + if self._enabled and runtime is not None: + self._install_local_index() + + def _resolve_enabled(self) -> bool: + """Determines whether ML Flashpoint should handle non-persistent saves. + + Returns: + True when ML Flashpoint is enabled and the run's checkpoint format is + one the ML Flashpoint strategies support. + """ + if not self._mlf_config.enabled: + _LOGGER.info("ML Flashpoint is disabled; delegating every checkpoint operation to Megatron Bridge.") + return False + ckpt_format = getattr(self.checkpoint_config, "ckpt_format", SUPPORTED_CKPT_FORMAT) + if ckpt_format != SUPPORTED_CKPT_FORMAT: + _LOGGER.warning( + "ML Flashpoint supports ckpt_format='%s' only, but this run uses '%s'. " + "Delegating every checkpoint operation to Megatron Bridge.", + SUPPORTED_CKPT_FORMAT, + ckpt_format, + ) + return False + return True + + @property + def enabled(self) -> bool: + """Whether ML Flashpoint handles this run's non-persistent checkpoints.""" + return self._enabled + + @property + def checkpointing_context(self) -> dict[str, Any]: + """The context Megatron Bridge caches strategies and resume hints in. + + Bridge reads ``local_checkpoint_manager`` out of this to decide whether a + resume should be attempted; the adapter installs its own index under that + key so ML Flashpoint checkpoints are discoverable. + + Reading this builds the ML Flashpoint runtime if it does not exist yet. + Megatron Bridge reads it during setup, before any save has happened, and + without the runtime there would be no index to advertise -- a restarted + job whose only checkpoint is an ML Flashpoint one would silently start + from scratch. + """ + if self._enabled and self._local_index is None: + self._ensure_runtime() + return self._context + + @property + def runtime(self) -> Optional[MLFlashpointBridgeRuntime]: + """The ML Flashpoint runtime, once it has been built.""" + return self._runtime + + def _ensure_runtime(self) -> Optional[MLFlashpointBridgeRuntime]: + """Builds the shared runtime on first use. + + Returns: + The runtime, or None if it could not be built (ML Flashpoint is then + disabled for the rest of the run). + """ + if not self._enabled: + return None + if self._runtime is not None: + return self._runtime + try: + self._runtime = get_runtime(self._mlf_config) + self._install_local_index() + except Exception: + _LOGGER.exception( + "Failed to initialize the ML Flashpoint runtime. Falling back to Megatron Bridge checkpointing " + "for the remainder of this run." + ) + self._enabled = False + self._runtime = None + return self._runtime + + def _install_local_index(self) -> None: + """Publishes the resume index into the checkpointing context.""" + if self._runtime is None or self._local_index is not None: + return + self._local_index = MLFlashpointLocalCheckpointIndex( + base_container=self._runtime.base_container, + checkpoint_loader=self._runtime.checkpoint_loader, + ) + self._context["local_checkpoint_manager"] = self._local_index + + def _version_container(self, step: int) -> CheckpointContainerId: + """Returns the container for a given step. + + Args: + step: The training step. + + Returns: + The child container ID for that step. + """ + return CheckpointContainerId.create_child( + self._runtime.base_container, + CheckpointContainerId.format_version_container(step), + ) + + @log_execution_time(logger=_LOGGER, name="MLFlashpointBridgeCheckpointManager.save", level=logging.INFO) + def save(self, ctx, callback_manager=None) -> None: + """Saves a checkpoint. + + Non-persistent checkpoints go to ML Flashpoint; everything else is + delegated to Megatron Bridge unchanged. + + A failure on the ML Flashpoint path is logged and swallowed: a + non-persistent checkpoint is an optimization for crash recovery, and + losing one must not take the training job down. + + Args: + ctx: The Bridge ``CheckpointSaveContext``. + callback_manager: The Bridge callback manager, if any. + """ + 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) + 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, + ) + + def _delegate_save(self, ctx, callback_manager) -> None: + """Runs Megatron Bridge's own save. + + Args: + ctx: The Bridge ``CheckpointSaveContext``. + callback_manager: The Bridge callback manager, if any. + + Raises: + RuntimeError: If Bridge is asked for a local non-persistent save that + ML Flashpoint was supposed to own. + """ + if ctx.non_persistent_ckpt and self._is_local_non_persistent(): + raise RuntimeError( + "Megatron Bridge was asked to write a local non-persistent checkpoint, but " + "non_persistent_ckpt_type='local' is what routes those to ML Flashpoint and " + "ML Flashpoint is unavailable. Set non_persistent_ckpt_type='global' (or drop " + "non_persistent_save_interval) to run without ML Flashpoint." + ) + save_checkpoint( + state=ctx.state, + model=ctx.model, + optimizer=ctx.optimizer, + opt_param_scheduler=ctx.opt_param_scheduler, + num_floating_point_operations_so_far=ctx.num_floating_point_operations_so_far, + checkpointing_context=self._context, + non_persistent_ckpt=ctx.non_persistent_ckpt, + train_data_iterator=ctx.train_data_iterator, + pg_collection=getattr(ctx, "pg_collection", None), + callback_manager=callback_manager, + module_name=getattr(ctx, "module_name", None), + ) + + def _is_local_non_persistent(self) -> bool: + """Whether this run routes non-persistent checkpoints to the local path.""" + return getattr(self.checkpoint_config, "non_persistent_ckpt_type", None) == LOCAL_NON_PERSISTENT_CKPT_TYPE + + @log_execution_time(logger=_LOGGER, name="MLFlashpointBridgeCheckpointManager.mlf_save", level=logging.INFO) + def _save_ml_flashpoint(self, ctx, step: int) -> None: + """Stages and schedules an ML Flashpoint save for the current step. + + Args: + ctx: The Bridge ``CheckpointSaveContext``. + step: The training step. + """ + container = self._version_container(step) + state_dict = bridge_state.build_save_state_dict( + ctx, + step=step, + num_floating_point_operations_so_far=ctx.num_floating_point_operations_so_far, + ) + + async_request = save_local_aware_megatron_checkpoint( + checkpoint=state_dict, + checkpoint_dir=str(container), + save_strategy=self._runtime.save_strategy, + async_save=self._mlf_config.async_save, + ) + if async_request is not None: + self._runtime.schedule(async_request) + # A new container exists, so any cached resume decision is stale. + if self._local_index is not None: + self._local_index.invalidate() + _LOGGER.info("Scheduled ML Flashpoint checkpoint for step %d at '%s'", step, container) + + @log_execution_time(logger=_LOGGER, name="MLFlashpointBridgeCheckpointManager.load", level=logging.INFO) + def load(self, ctx) -> tuple[int, int]: + """Loads a checkpoint, preferring the newest ML Flashpoint container. + + Args: + ctx: The Bridge ``CheckpointLoadContext``. + + Returns: + A tuple of (step, cumulative floating point operations). ``(0, 0)`` + when nothing was loaded. + """ + container = self._find_ml_flashpoint_checkpoint() + if container is not None: + loaded = self._load_ml_flashpoint(ctx, container) + if loaded is not None: + return loaded + + # Bridge's own local-checkpoint path cannot read an ML Flashpoint + # container, so make sure it is not offered one. + if self._local_index is not None: + self._local_index.disable() + return load_checkpoint( + state=ctx.state, + model=ctx.model, + optimizer=ctx.optimizer, + opt_param_scheduler=ctx.opt_param_scheduler, + strict=ctx.strict, + checkpointing_context=self._context, + skip_load_to_model_and_opt=ctx.skip_load_to_model_and_opt, + pg_collection=getattr(ctx, "pg_collection", None), + module_name=getattr(ctx, "module_name", None), + ) + + def _find_ml_flashpoint_checkpoint(self) -> Optional[CheckpointContainerId]: + """Finds the latest recoverable ML Flashpoint container, if any. + + Returns: + The container, or None. + """ + if self._ensure_runtime() is None: + return None + self._install_local_index() + return self._local_index.resolve_latest_container() + + @log_execution_time(logger=_LOGGER, name="MLFlashpointBridgeCheckpointManager.mlf_load", level=logging.INFO) + def _load_ml_flashpoint(self, ctx, container: CheckpointContainerId) -> Optional[tuple[int, int]]: + """Loads from an ML Flashpoint container. + + The read itself is allowed to fail (the caller then falls back to + Megatron Bridge), but once state has been applied to the model a failure + is fatal, because there is no consistent state to fall back from. + + Args: + ctx: The Bridge ``CheckpointLoadContext``. + container: The container to read. + + Returns: + A tuple of (step, cumulative FLOPs), or None when the read failed and + the caller should fall back. + """ + try: + sharded_state_dict = bridge_state.build_load_state_dict(ctx) + state_dict = mcore_dist_checkpointing.load( + sharded_state_dict=sharded_state_dict, + checkpoint_dir=str(container), + sharded_strategy=self._runtime.load_strategy, + common_strategy=TorchCommonLoadStrategy(), + ) + except Exception: + _LOGGER.exception( + "Failed to read the ML Flashpoint checkpoint at '%s'. Falling back to Megatron Bridge.", container + ) + return None + + step, flops = bridge_state.apply_loaded_state(ctx, state_dict) + _LOGGER.info("Recovered from ML Flashpoint checkpoint '%s' at step %d", container, step) + return step, flops + + @log_execution_time( + logger=_LOGGER, name="MLFlashpointBridgeCheckpointManager.finalize_async_saves", level=logging.DEBUG + ) + def finalize_async_saves(self, state, blocking: bool = False, terminate: bool = False) -> None: + """Finalizes pending saves on both the ML Flashpoint and Bridge queues. + + The queues are finalized independently: ML Flashpoint saves complete far + sooner than durable ones, and making them wait behind a durable save + would keep their buffers pinned long enough to exhaust the pool. + + Args: + state: The Bridge ``GlobalState``. + blocking: If True, waits for every pending save. + terminate: If True, tears down the queues afterwards. + """ + if self._runtime is not None: + try: + self._runtime.maybe_finalize(blocking=blocking) + except Exception: + _LOGGER.exception("Failed to finalize ML Flashpoint async saves.") + + maybe_finalize_async_save( + global_state=state, + ckpt_cfg=self.checkpoint_config, + blocking=blocking, + terminate=terminate, + ) + + if terminate: + self.shutdown() + + def shutdown(self) -> None: + """Releases ML Flashpoint resources held by this rank. + + Waits for in-flight saves and synchronizes across ranks before deleting + anything, so a peer is never mid-replication into a container that is + about to disappear. + """ + if self._runtime is None: + return + try: + self._runtime.maybe_finalize(blocking=True) + except Exception: + _LOGGER.exception("Failed to drain ML Flashpoint async saves during shutdown.") + + if dist.is_available() and dist.is_initialized(): + try: + dist.barrier() + except Exception: + _LOGGER.exception("Barrier before ML Flashpoint teardown failed. Continuing.") + + from ml_flashpoint.adapter.megatron_bridge.runtime import shutdown_runtime + + shutdown_runtime(remove_checkpoints=not self._mlf_config.keep_checkpoints_on_finalize) + self._runtime = None + self._local_index = None + self._context.pop("local_checkpoint_manager", None) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def init_checkpointing_context_safely(checkpoint_config) -> dict[str, Any]: + """Builds Bridge's checkpointing context without requiring NVRx. + + ``init_checkpointing_context`` insists on ``nvidia_resiliency_ext`` when + ``non_persistent_ckpt_type='local'``, but with this adapter that setting means + "ML Flashpoint owns local checkpoints", so NVRx is not needed. + + Args: + checkpoint_config: The Bridge ``CheckpointConfig``. + + Returns: + The checkpointing context dictionary. + """ + try: + return init_checkpointing_context(checkpoint_config) + except RuntimeError: + _LOGGER.info( + "Megatron Bridge could not build a local checkpointing context (nvidia_resiliency_ext is not " + "installed). ML Flashpoint provides local checkpointing instead; continuing with an empty context." + ) + return {} diff --git a/src/ml_flashpoint/adapter/megatron_bridge/config.py b/src/ml_flashpoint/adapter/megatron_bridge/config.py new file mode 100644 index 0000000..fbf1e53 --- /dev/null +++ b/src/ml_flashpoint/adapter/megatron_bridge/config.py @@ -0,0 +1,156 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration for the Megatron Bridge ML Flashpoint adapter. + +Megatron Bridge instantiates a custom checkpoint manager through +:func:`megatron.bridge.training.checkpointing.create_checkpoint_manager`, which +only passes the framework's own ``CheckpointConfig``. There is therefore no +place in the Bridge config to carry ML Flashpoint's own knobs, so they are +resolved from (in order of precedence): + +1. an explicit :func:`configure` call made before training starts, and +2. ``MLFLASHPOINT_*`` environment variables. +""" + +import dataclasses +from typing import Optional + +from ml_flashpoint.core import utils +from ml_flashpoint.core.checkpoint_saver import DEFAULT_INITIAL_BUFFER_SIZE_BYTES +from ml_flashpoint.core.mlf_logging import get_logger + +_LOGGER = get_logger(__name__) + +DEFAULT_BASE_CONTAINER = "/dev/shm/ml_flashpoint" +"""Default base container for checkpoint versions. + +``/dev/shm`` is memory backed on the training nodes, which is what makes ML +Flashpoint saves fast. Point this elsewhere only if the node exposes a faster +node-local mount. +""" + + +@dataclasses.dataclass(frozen=True) +class MLFlashpointBridgeConfig: + """ML Flashpoint settings for the Megatron Bridge adapter. + + Attributes: + enabled: Whether ML Flashpoint handles non-persistent checkpoints. When + False, the adapter delegates every operation to Megatron Bridge, which + makes it a no-op wrapper. Useful for A/B experiments that keep the + exact same launch command. + base_container: Base container (directory) that holds one child container + per checkpoint version. Must be node-local and is expected to be + memory backed. + async_save: Whether saves are scheduled asynchronously. Keeping this True + is what removes the write from the training critical path. + write_thread_count: Number of writer threads used per rank. + initial_write_buffer_size_bytes: Initial size of each write buffer. + use_optimized_save: Whether to use the optimized zero-copy tensor save. + use_cached_ckpt_structure: Whether to reuse the save plan across steps. + Only safe when the checkpoint structure is constant. + use_fully_parallel_wrapper: Whether to wrap the save/load strategies so + checkpoint data is spread evenly across ranks. + keep_checkpoints_on_finalize: Whether to keep the ML Flashpoint container + when training terminates normally. Off by default so buffers are + released back to the node. + """ + + enabled: bool = True + base_container: str = DEFAULT_BASE_CONTAINER + async_save: bool = True + write_thread_count: int = 1 + initial_write_buffer_size_bytes: int = DEFAULT_INITIAL_BUFFER_SIZE_BYTES + use_optimized_save: bool = True + use_cached_ckpt_structure: bool = False + use_fully_parallel_wrapper: bool = True + keep_checkpoints_on_finalize: bool = False + + def __post_init__(self): + if not self.base_container: + raise ValueError("base_container cannot be empty.") + if self.write_thread_count < 1: + raise ValueError(f"write_thread_count must be >= 1, got {self.write_thread_count}.") + if self.initial_write_buffer_size_bytes <= 0: + raise ValueError( + f"initial_write_buffer_size_bytes must be > 0, got {self.initial_write_buffer_size_bytes}." + ) + + @classmethod + def from_env(cls) -> "MLFlashpointBridgeConfig": + """Builds a config from ``MLFLASHPOINT_*`` environment variables. + + Every field falls back to the dataclass default when its variable is + unset. See :func:`ml_flashpoint.core.utils.get_env_var_prefix` for the + prefix applied to each name below. + + Returns: + The environment-derived configuration. + """ + defaults = cls() + return cls( + enabled=utils.get_env_val_bool("BRIDGE_ENABLED", defaults.enabled), + base_container=utils.get_env_val_str("BASE_CONTAINER", defaults.base_container), + async_save=utils.get_env_val_bool("ASYNC_SAVE", defaults.async_save), + write_thread_count=utils.get_env_val_int("WRITE_THREAD_COUNT", defaults.write_thread_count), + initial_write_buffer_size_bytes=utils.get_env_val_int( + "INITIAL_WRITE_BUFFER_SIZE_BYTES", defaults.initial_write_buffer_size_bytes + ), + use_optimized_save=utils.get_env_val_bool("USE_OPTIMIZED_SAVE", defaults.use_optimized_save), + use_cached_ckpt_structure=utils.get_env_val_bool( + "USE_CACHED_CKPT_STRUCTURE", defaults.use_cached_ckpt_structure + ), + use_fully_parallel_wrapper=utils.get_env_val_bool( + "USE_FULLY_PARALLEL_WRAPPER", defaults.use_fully_parallel_wrapper + ), + keep_checkpoints_on_finalize=utils.get_env_val_bool( + "KEEP_CHECKPOINTS_ON_FINALIZE", defaults.keep_checkpoints_on_finalize + ), + ) + + +_CONFIGURED: Optional[MLFlashpointBridgeConfig] = None + + +def configure(config: MLFlashpointBridgeConfig) -> None: + """Registers the config the adapter uses instead of reading the environment. + + Call this before Megatron Bridge builds its checkpoint manager, i.e. before + ``megatron.bridge.training.setup`` (or, for NeMo RL, before the Megatron + policy worker is initialized). + + Args: + config: The configuration to use. + """ + global _CONFIGURED + _CONFIGURED = config + _LOGGER.info("Registered ML Flashpoint Megatron Bridge config: %s", config) + + +def reset_configuration() -> None: + """Drops a previously registered config so the environment is read again.""" + global _CONFIGURED + _CONFIGURED = None + + +def get_config() -> MLFlashpointBridgeConfig: + """Returns the registered config, or one derived from the environment. + + Returns: + The effective configuration. + """ + if _CONFIGURED is not None: + return _CONFIGURED + return MLFlashpointBridgeConfig.from_env() diff --git a/src/ml_flashpoint/adapter/megatron_bridge/local_checkpoint_index.py b/src/ml_flashpoint/adapter/megatron_bridge/local_checkpoint_index.py new file mode 100644 index 0000000..e38020c --- /dev/null +++ b/src/ml_flashpoint/adapter/megatron_bridge/local_checkpoint_index.py @@ -0,0 +1,148 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Advertises ML Flashpoint checkpoints to Megatron Bridge's resume decision.""" + +import json +import os +from typing import Optional + +import torch.distributed as dist + +from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId +from ml_flashpoint.core.checkpoint_loader import MLFlashpointCheckpointLoader +from ml_flashpoint.core.mlf_logging import get_logger + +_LOGGER = get_logger(__name__) + +NO_CHECKPOINT = -1 +"""Sentinel Megatron Bridge uses for "no local checkpoint exists".""" + + +class MLFlashpointLocalCheckpointIndex: + """Duck-types the part of NVRx's ``LocalCheckpointManager`` that Bridge reads. + + Megatron Bridge decides whether to attempt a resume in + ``megatron.bridge.training.setup._should_load_checkpoint``, which looks for a + ``local_checkpoint_manager`` in the checkpoint manager's context and calls + ``find_latest()`` on it. Without this shim a run whose only checkpoint is an + ML Flashpoint one would start from scratch. + + Only ``find_latest`` and ``local_ckpt_dir`` are implemented: the adapter's own + checkpoint manager performs the load, so Bridge's local load path must never + be reached. + """ + + def __init__( + self, + base_container: CheckpointContainerId, + checkpoint_loader: MLFlashpointCheckpointLoader, + ): + """Initializes the index. + + Args: + base_container: The base container holding checkpoint versions. + checkpoint_loader: The loader used to find recoverable checkpoints. + """ + self._base_container = base_container + self._checkpoint_loader = checkpoint_loader + self._latest: Optional[CheckpointContainerId] = None + self._resolved = False + self._disabled = False + + @property + def local_ckpt_dir(self) -> str: + """The directory Bridge would use as the local checkpoint root.""" + return str(self._base_container) + + @property + def latest_container(self) -> Optional[CheckpointContainerId]: + """The resolved latest complete container, if discovery already ran.""" + return self._latest + + def disable(self) -> None: + """Makes ``find_latest`` report no checkpoint from now on. + + Called before delegating to Megatron Bridge so that Bridge does not try to + read an ML Flashpoint container through its own local-checkpoint path, + which expects an NVRx ``MCoreTensorAwareStateDict`` container. + """ + self._disabled = True + + def find_latest(self) -> int: + """Returns the step of the latest complete ML Flashpoint checkpoint. + + Discovery is collective (it gathers per-rank object inventories) and is + therefore performed at most once, with the result cached. + + Returns: + The step number, or ``NO_CHECKPOINT`` when nothing is recoverable. + """ + container = self.resolve_latest_container() + if container is None: + return NO_CHECKPOINT + step = CheckpointContainerId.parse_version_container_step(os.path.basename(str(container))) + return NO_CHECKPOINT if step is None else step + + def resolve_latest_container(self) -> Optional[CheckpointContainerId]: + """Finds (once) the latest complete checkpoint container. + + Returns: + The container, or None when nothing is recoverable. + """ + if self._disabled: + return None + if self._resolved: + return self._latest + + self._resolved = True + try: + self._latest = self._checkpoint_loader.get_latest_complete_checkpoint(self._base_container) + except Exception: + _LOGGER.exception("Failed to discover ML Flashpoint checkpoints under '%s'.", self._base_container) + self._latest = None + + if self._latest is not None: + _ensure_megatron_metadata_stub(self._latest) + return self._latest + + def invalidate(self) -> None: + """Forces the next ``find_latest`` call to rediscover.""" + self._resolved = False + self._latest = None + + +def _ensure_megatron_metadata_stub(container: CheckpointContainerId) -> None: + """Writes the ``metadata.json`` stub Megatron's loader validates. + + ML Flashpoint stores tensor data in shared-memory buffers rather than the + ``.distcp`` layout, but ``dist_checkpointing.load`` still validates a backend + marker file before dispatching to the strategy. The checks against this stub + are no-ops, and the save path writes it too; this covers a container that was + replicated from a peer node without the file. + + Args: + container: The checkpoint container to stub. + """ + try: + if dist.is_initialized() and dist.get_node_local_rank() != 0: + return + metadata_path = os.path.join(str(container), "metadata.json") + if os.path.exists(metadata_path): + return + with open(metadata_path, "w") as f: + json.dump({"sharded_backend": ""}, f) + _LOGGER.debug("Wrote Megatron metadata stub at '%s'", metadata_path) + except Exception: + _LOGGER.exception("Failed to write the Megatron metadata stub for '%s'.", container) diff --git a/src/ml_flashpoint/adapter/megatron_bridge/runtime.py b/src/ml_flashpoint/adapter/megatron_bridge/runtime.py new file mode 100644 index 0000000..a84101f --- /dev/null +++ b/src/ml_flashpoint/adapter/megatron_bridge/runtime.py @@ -0,0 +1,306 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-rank ML Flashpoint runtime shared by the Megatron Bridge integrations. + +This owns the objects that must exist exactly once per rank (buffer pool, +replication service, save/load strategies, async queue) and keeps their +construction out of the checkpoint manager itself so that NeMo RL, which does +not go through Megatron Bridge's checkpoint-manager factory, can reuse them. +""" + +import concurrent.futures +import os +import threading +from typing import Optional + +import torch.distributed as dist +from megatron.core.dist_checkpointing.strategies.async_utils import ( + AsyncCallsQueue, + AsyncRequest, +) +from megatron.core.dist_checkpointing.strategies.fully_parallel import ( + FullyParallelLoadStrategyWrapper, + FullyParallelSaveStrategyWrapper, +) +from torch import multiprocessing as torch_mp + +from ml_flashpoint.adapter.megatron.load_strategies import MLFlashpointMegatronLoadStrategy +from ml_flashpoint.adapter.megatron.save_strategies import MLFlashpointMegatronAsyncSaveStrategy +from ml_flashpoint.adapter.megatron_bridge.config import MLFlashpointBridgeConfig +from ml_flashpoint.adapter.pytorch.memory_storage_writer import MemoryStorageWriter +from ml_flashpoint.checkpoint_object_manager.checkpoint_object_manager import CheckpointObjectManager +from ml_flashpoint.core.buffer_pool import BufferPoolConfig +from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId +from ml_flashpoint.core.checkpoint_loader import DefaultMLFlashpointCheckpointLoader +from ml_flashpoint.core.checkpoint_saver import DefaultMLFlashpointCheckpointSaver +from ml_flashpoint.core.mlf_logging import get_logger +from ml_flashpoint.replication.replication_manager import ReplicationManager + +_LOGGER = get_logger(__name__) + +NUM_OF_BUFFERS_PER_OBJECT = 2 +"""Buffers reserved per writer thread, matching the NeMo adapter's pool sizing.""" + + +class MLFlashpointBridgeRuntime: + """Owns the per-rank ML Flashpoint objects used by the Bridge adapter. + + Instances are expensive (they start a buffer pool, a transfer service and a + persistent async worker), so a single instance is shared for the lifetime of + a training process. Use :func:`get_runtime` rather than constructing one + directly unless a test needs isolation. + """ + + def __init__(self, config: MLFlashpointBridgeConfig): + """Builds the runtime and initializes replication. + + Requires ``torch.distributed`` to already be initialized, because + replication address exchange is a collective. + + Args: + config: The ML Flashpoint settings to apply. + + Raises: + RuntimeError: If ``torch.distributed`` is not initialized. + """ + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "ML Flashpoint requires an initialized torch.distributed process group. " + "Build the runtime after Megatron initialization." + ) + + self._config = config + self._base_container = CheckpointContainerId(config.base_container) + self._closed = False + + pool_config = BufferPoolConfig( + pool_dir_path=os.path.join(str(self._base_container), "buffer_pool"), + rank=dist.get_rank(), + num_buffers=config.write_thread_count * NUM_OF_BUFFERS_PER_OBJECT, + buffer_size=config.initial_write_buffer_size_bytes, + ) + self._checkpoint_object_manager = CheckpointObjectManager(pool_config=pool_config) + + self._replication_manager = ReplicationManager() + self._replication_manager.initialize(checkpoint_object_manager=self._checkpoint_object_manager) + + # 'spawn' avoids inheriting the parent's CUDA context in the SyncManager + # process. A forked manager that outlives a SIGKILLed trainer keeps GPU + # memory locked and makes the in-job restart OOM. + ctx = torch_mp.get_context("spawn") + self._mp_manager_future: concurrent.futures.Future = concurrent.futures.Future() + + def start_manager(): + self._mp_manager_future.set_result(ctx.Manager()) + + threading.Thread(target=start_manager, daemon=True).start() + + self._checkpoint_saver = DefaultMLFlashpointCheckpointSaver( + global_rank_getter=dist.get_rank, + local_rank_getter=dist.get_node_local_rank, + global_barrier_func=dist.barrier, + ckpt_obj_manager=self._checkpoint_object_manager, + replication_manager=self._replication_manager, + initial_buffer_size_bytes=config.initial_write_buffer_size_bytes, + use_optimized_save=config.use_optimized_save, + ) + self._checkpoint_loader = DefaultMLFlashpointCheckpointLoader( + self._checkpoint_object_manager, + self._replication_manager, + global_rank_getter=dist.get_rank, + local_rank_getter=dist.get_node_local_rank, + broadcast_object_list_func=dist.broadcast_object_list, + all_gather_object_func=dist.all_gather_object, + world_size_getter=dist.get_world_size, + ) + + save_strategy = MLFlashpointMegatronAsyncSaveStrategy( + storage_writer=MemoryStorageWriter( + checkpoint_saver=self._checkpoint_saver, + mp_manager_future=self._mp_manager_future, + thread_count=config.write_thread_count, + ), + use_cached_ckpt_structure=config.use_cached_ckpt_structure, + ) + load_strategy = MLFlashpointMegatronLoadStrategy( + replication_manager=self._replication_manager, + checkpoint_loader=self._checkpoint_loader, + ) + if config.use_fully_parallel_wrapper: + # No parallelization group is passed, matching the NeMo adapter: the + # wrapper then distributes across the default (world) group. + save_strategy = FullyParallelSaveStrategyWrapper(save_strategy) + load_strategy = FullyParallelLoadStrategyWrapper(load_strategy) + self._save_strategy = save_strategy + self._load_strategy = load_strategy + + # Persistent so the worker process (and its buffer pool handles) is + # reused across steps instead of respawned per checkpoint. + self._async_calls_queue = AsyncCallsQueue(persistent=True) + + @property + def config(self) -> MLFlashpointBridgeConfig: + """The configuration this runtime was built with.""" + return self._config + + @property + def base_container(self) -> CheckpointContainerId: + """The base container holding one child container per checkpoint version.""" + return self._base_container + + @property + def save_strategy(self): + """The Megatron sharded save strategy backed by ML Flashpoint.""" + return self._save_strategy + + @property + def load_strategy(self): + """The Megatron sharded load strategy backed by ML Flashpoint.""" + return self._load_strategy + + @property + def checkpoint_loader(self) -> DefaultMLFlashpointCheckpointLoader: + """The loader used to discover and retrieve recoverable checkpoints.""" + return self._checkpoint_loader + + @property + def checkpoint_object_manager(self) -> CheckpointObjectManager: + """The object manager owning this rank's buffer pool.""" + return self._checkpoint_object_manager + + @property + def replication_manager(self) -> ReplicationManager: + """The replication manager used to mirror objects to peer nodes.""" + return self._replication_manager + + def schedule(self, async_request: AsyncRequest) -> int: + """Schedules an async save request on the ML Flashpoint queue. + + Args: + async_request: The request returned by the save strategy. + + Returns: + The scheduled call index. + """ + call_idx = self._async_calls_queue.schedule_async_request(async_request) + _LOGGER.debug("Scheduled ML Flashpoint async call #%d", call_idx) + return call_idx + + def num_unfinalized_calls(self) -> int: + """Returns how many scheduled saves have not been finalized yet.""" + return self._async_calls_queue.get_num_unfinalized_calls() + + def maybe_finalize(self, blocking: bool = False) -> bool: + """Finalizes completed saves. + + Args: + blocking: If True, waits for every pending save to complete. + + Returns: + True if at least one call was finalized. + """ + if self._closed or self._async_calls_queue.get_num_unfinalized_calls() == 0: + return False + finalized = self._async_calls_queue.maybe_finalize_async_calls(blocking) + if finalized: + _LOGGER.debug("Finalized ML Flashpoint async calls: %s", [f"#{idx}" for idx in finalized]) + return len(finalized) > 0 + + def shutdown(self, remove_checkpoints: bool = True) -> None: + """Tears down the runtime, releasing buffers and background workers. + + Safe to call more than once. + + Args: + remove_checkpoints: Whether to delete the base container so the + node's memory is reclaimed. + """ + if self._closed: + return + self._closed = True + + try: + self._replication_manager.shutdown() + except Exception: + _LOGGER.exception("Failed to shut down the ReplicationManager. Continuing teardown.") + + if remove_checkpoints and self._is_local_rank_zero(): + try: + self._checkpoint_object_manager.delete_container(self._base_container) + except Exception: + _LOGGER.exception("Failed to delete container '%s'. Continuing teardown.", self._base_container) + + # The buffer pool lives in the persistent worker process, so its teardown + # has to be scheduled onto that same process. + try: + self._async_calls_queue.schedule_async_request( + AsyncRequest( + async_fn=self._checkpoint_object_manager.teardown_pool, + async_fn_args=(), + finalize_fns=[], + ) + ) + except Exception: + _LOGGER.debug("Could not schedule buffer pool teardown; the queue is likely already closed.") + + self._async_calls_queue.close() + # PersistentAsyncCaller.__del__ calls close(), which calls + # torch.distributed.get_rank() and crashes if the process group is gone + # by interpreter shutdown. The queue is already closed, so make the + # second close a no-op. + caller = getattr(self._async_calls_queue, "persistent_caller", None) + if caller is not None and hasattr(caller, "close"): + caller.close = lambda: None + + @staticmethod + def _is_local_rank_zero() -> bool: + try: + return dist.get_node_local_rank() == 0 + except Exception: + return not dist.is_initialized() or dist.get_rank() == 0 + + +_RUNTIME: Optional[MLFlashpointBridgeRuntime] = None + + +def get_runtime(config: MLFlashpointBridgeConfig) -> MLFlashpointBridgeRuntime: + """Returns the process-wide runtime, building it on first use. + + Args: + config: The configuration used when the runtime has to be built. Ignored + when a runtime already exists. + + Returns: + The shared runtime instance. + """ + global _RUNTIME + if _RUNTIME is None: + _RUNTIME = MLFlashpointBridgeRuntime(config) + return _RUNTIME + + +def shutdown_runtime(remove_checkpoints: bool = True) -> None: + """Shuts down and clears the process-wide runtime, if one was built. + + Args: + remove_checkpoints: Whether to delete the base container. + """ + global _RUNTIME + if _RUNTIME is None: + return + try: + _RUNTIME.shutdown(remove_checkpoints=remove_checkpoints) + finally: + _RUNTIME = None diff --git a/src/ml_flashpoint/adapter/nemo_rl/__init__.py b/src/ml_flashpoint/adapter/nemo_rl/__init__.py new file mode 100644 index 0000000..8616623 --- /dev/null +++ b/src/ml_flashpoint/adapter/nemo_rl/__init__.py @@ -0,0 +1,25 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ML Flashpoint adapter for NeMo RL.""" + +from ml_flashpoint.adapter.nemo_rl.checkpointer import ( + MLFlashpointNeMoRLCheckpointer as MLFlashpointNeMoRLCheckpointer, +) +from ml_flashpoint.adapter.nemo_rl.integration import MODE_AUGMENT as MODE_AUGMENT +from ml_flashpoint.adapter.nemo_rl.integration import MODE_REPLACE as MODE_REPLACE +from ml_flashpoint.adapter.nemo_rl.integration import get_checkpointer as get_checkpointer +from ml_flashpoint.adapter.nemo_rl.integration import install_from_env as install_from_env +from ml_flashpoint.adapter.nemo_rl.integration import install_into_worker as install_into_worker +from ml_flashpoint.adapter.nemo_rl.integration import uninstall_from_worker as uninstall_from_worker diff --git a/src/ml_flashpoint/adapter/nemo_rl/checkpointer.py b/src/ml_flashpoint/adapter/nemo_rl/checkpointer.py new file mode 100644 index 0000000..97d8904 --- /dev/null +++ b/src/ml_flashpoint/adapter/nemo_rl/checkpointer.py @@ -0,0 +1,210 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ML Flashpoint checkpointing for NeMo RL's Megatron policy worker. + +NeMo RL builds its Megatron training state with Megatron Bridge but drives +checkpointing itself: the policy worker calls +``megatron.bridge.training.checkpointing.save_checkpoint`` directly rather than +going through ``create_checkpoint_manager``, so ``CheckpointConfig. +custom_manager_class`` is never consulted. This module therefore reuses the +Bridge checkpoint manager from the worker side, feeding it a save/load context +built from the worker's own attributes. +""" + +import dataclasses +import logging +from typing import Any, Optional + +from ml_flashpoint.adapter.megatron_bridge.checkpoint_manager import ( + MLFlashpointBridgeCheckpointManager, +) +from ml_flashpoint.adapter.megatron_bridge.config import MLFlashpointBridgeConfig, get_config +from ml_flashpoint.core.mlf_logging import get_logger +from ml_flashpoint.core.utils import log_execution_time + +_LOGGER = get_logger(__name__) + + +@dataclasses.dataclass +class _SaveContext: + """Duck-typed stand-in for Megatron Bridge's ``CheckpointSaveContext``. + + The Bridge manager only reads attributes off the context, so NeMo RL does not + need to construct Bridge's dataclass (whose field set varies across versions). + """ + + state: Any + model: list + optimizer: Any + opt_param_scheduler: Any + num_floating_point_operations_so_far: int + train_data_iterator: Any = None + non_persistent_ckpt: bool = True + pg_collection: Any = None + module_name: Optional[str] = None + + +@dataclasses.dataclass +class _LoadContext: + """Duck-typed stand-in for Megatron Bridge's ``CheckpointLoadContext``.""" + + state: Any + model: list + optimizer: Any + opt_param_scheduler: Any + strict: bool = True + skip_load_to_model_and_opt: bool = False + pg_collection: Any = None + module_name: Optional[str] = None + + +class MLFlashpointNeMoRLCheckpointer: + """Saves and restores a NeMo RL Megatron policy through ML Flashpoint. + + ML Flashpoint checkpoints are node-local and peer-replicated: they survive a + process or node failure within a run, but not the loss of the whole cluster + and not the end of the run. They are a fast recovery tier, not a replacement + for NeMo RL's durable ``checkpointing.save_period`` checkpoints, which must + stay enabled. + + Attributes: + manager: The underlying Megatron Bridge checkpoint manager. + """ + + def __init__( + self, + checkpoint_config, + mlf_config: Optional[MLFlashpointBridgeConfig] = None, + ): + """Initializes the checkpointer. + + Args: + checkpoint_config: The Bridge ``CheckpointConfig`` the worker was + configured with (``worker.mcore_state.cfg.checkpoint``). + mlf_config: ML Flashpoint settings. Defaults to the registered or + environment-derived config. + """ + self._mlf_config = mlf_config if mlf_config is not None else get_config() + self.manager = MLFlashpointBridgeCheckpointManager( + checkpoint_config=checkpoint_config, + mlf_config=self._mlf_config, + ) + + @property + def enabled(self) -> bool: + """Whether ML Flashpoint is active for this run.""" + return self.manager.enabled + + @log_execution_time(logger=_LOGGER, name="MLFlashpointNeMoRLCheckpointer.save", level=logging.INFO) + def save( + self, + state, + model, + optimizer=None, + opt_param_scheduler=None, + num_floating_point_operations_so_far: Optional[int] = None, + ) -> None: + """Writes an ML Flashpoint checkpoint for the current step. + + The step comes from ``state.train_state.step``, which NeMo RL keeps in + sync with the RL loop. + + Args: + state: The worker's ``GlobalState`` (``worker.mcore_state``). + model: The model module, or a list of module chunks. + optimizer: The optimizer, if its state should be captured. + opt_param_scheduler: The scheduler, if its state should be captured. + num_floating_point_operations_so_far: Cumulative FLOPs. Defaults to + the value already tracked on the train state. + """ + if num_floating_point_operations_so_far is None: + num_floating_point_operations_so_far = getattr(state.train_state, "floating_point_operations_so_far", 0) + ctx = _SaveContext( + state=state, + model=_as_module_list(model), + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + num_floating_point_operations_so_far=int(num_floating_point_operations_so_far), + non_persistent_ckpt=True, + ) + self.manager.save(ctx, callback_manager=None) + + @log_execution_time(logger=_LOGGER, name="MLFlashpointNeMoRLCheckpointer.load", level=logging.INFO) + def load( + self, + state, + model, + optimizer=None, + opt_param_scheduler=None, + strict: bool = True, + ) -> Optional[tuple[int, int]]: + """Restores from the newest ML Flashpoint checkpoint, if one exists. + + Unlike the Megatron Bridge manager's ``load``, this never falls back to + Bridge's durable load path: NeMo RL owns that decision and performs it + during worker setup. + + Args: + state: The worker's ``GlobalState``. + model: The model module, or a list of module chunks. + optimizer: The optimizer to restore into. + opt_param_scheduler: The scheduler to restore into. + strict: Whether to enforce strict key matching on the model load. + + Returns: + A tuple of (step, cumulative FLOPs), or None when nothing was + recovered. + """ + container = self.manager._find_ml_flashpoint_checkpoint() + if container is None: + _LOGGER.info("No recoverable ML Flashpoint checkpoint found.") + return None + ctx = _LoadContext( + state=state, + model=_as_module_list(model), + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + strict=strict, + ) + return self.manager._load_ml_flashpoint(ctx, container) + + def finalize(self, blocking: bool = True) -> None: + """Finalizes pending ML Flashpoint saves. + + Args: + blocking: If True, waits for every pending save to complete. + """ + runtime = self.manager.runtime + if runtime is None: + return + runtime.maybe_finalize(blocking=blocking) + + def shutdown(self) -> None: + """Releases every ML Flashpoint resource held by this rank.""" + self.manager.shutdown() + + +def _as_module_list(model) -> list: + """Normalizes a model argument to the list of chunks Bridge expects. + + Args: + model: A single module or a list of module chunks. + + Returns: + A list of modules. + """ + if isinstance(model, (list, tuple)): + return list(model) + return [model] diff --git a/src/ml_flashpoint/adapter/nemo_rl/integration.py b/src/ml_flashpoint/adapter/nemo_rl/integration.py new file mode 100644 index 0000000..4398c10 --- /dev/null +++ b/src/ml_flashpoint/adapter/nemo_rl/integration.py @@ -0,0 +1,249 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Attaching ML Flashpoint to NeMo RL's Megatron policy worker. + +NeMo RL calls Megatron Bridge's ``save_checkpoint`` from +``MegatronPolicyWorker.save_checkpoint`` and never consults +``CheckpointConfig.custom_manager_class``, so there is no configuration-only way +to reach the Bridge checkpoint manager. :func:`install_into_worker` wraps that +one method on a worker instance instead, which keeps the change confined to a +single, explicitly named entry point. + +Two modes are supported: + +``augment`` + Every NeMo RL checkpoint is still written durably, and an ML Flashpoint + checkpoint is written alongside it. Use this when the goal is faster + recovery without weakening durability. + +``replace`` + Only every ``durable_every_n_saves``-th checkpoint is written durably; the + rest go to ML Flashpoint alone. This is what removes checkpoint stalls from + the RL loop, and what a checkpoint-time A/B measures. +""" + +import functools +import logging +from typing import Any, Optional + +from ml_flashpoint.adapter.megatron_bridge.config import MLFlashpointBridgeConfig, get_config +from ml_flashpoint.adapter.nemo_rl.checkpointer import MLFlashpointNeMoRLCheckpointer +from ml_flashpoint.core import utils +from ml_flashpoint.core.mlf_logging import get_logger +from ml_flashpoint.core.utils import log_execution_time + +_LOGGER = get_logger(__name__) + +MODE_AUGMENT = "augment" +MODE_REPLACE = "replace" +_VALID_MODES = (MODE_AUGMENT, MODE_REPLACE) + +_INSTALLED_ATTR = "_ml_flashpoint_checkpointer" + + +class MLFlashpointWorkerHooks: + """Bookkeeping for one worker's ML Flashpoint installation. + + Attributes: + checkpointer: The ML Flashpoint checkpointer bound to the worker. + mode: Either ``augment`` or ``replace``. + durable_every_n_saves: In ``replace`` mode, how often a durable + checkpoint is still written. ``1`` means every save stays durable, + which makes ``replace`` equivalent to ``augment``. + """ + + def __init__( + self, + checkpointer: MLFlashpointNeMoRLCheckpointer, + mode: str, + durable_every_n_saves: int, + ): + self.checkpointer = checkpointer + self.mode = mode + self.durable_every_n_saves = durable_every_n_saves + self._save_count = 0 + + def should_save_durable(self) -> bool: + """Decides whether the current save must also go to durable storage. + + Returns: + True when the durable write should run. + """ + self._save_count += 1 + if self.mode == MODE_AUGMENT: + return True + return self._save_count % self.durable_every_n_saves == 0 + + +def install_into_worker( + worker, + mode: str = MODE_AUGMENT, + durable_every_n_saves: int = 1, + mlf_config: Optional[MLFlashpointBridgeConfig] = None, +) -> MLFlashpointNeMoRLCheckpointer: + """Routes a NeMo RL Megatron policy worker's saves through ML Flashpoint. + + Call this after the worker has finished initializing, i.e. once + ``worker.mcore_state``, ``worker.model`` and ``torch.distributed`` are all + live. Installing twice on the same worker is a no-op. + + Args: + worker: A NeMo RL ``MegatronPolicyWorker`` instance. + mode: ``augment`` to add an ML Flashpoint checkpoint next to every + durable one, or ``replace`` to skip most durable writes. + durable_every_n_saves: In ``replace`` mode, how often to still write + durably. Must be positive. + mlf_config: ML Flashpoint settings. Defaults to the registered or + environment-derived config. + + Returns: + The checkpointer bound to the worker. + + Raises: + ValueError: If ``mode`` or ``durable_every_n_saves`` is invalid. + AttributeError: If the worker does not expose the attributes the + adapter needs. + """ + if mode not in _VALID_MODES: + raise ValueError(f"mode must be one of {_VALID_MODES}, got '{mode}'.") + if durable_every_n_saves < 1: + raise ValueError(f"durable_every_n_saves must be a positive integer, got {durable_every_n_saves}.") + + existing = getattr(worker, _INSTALLED_ATTR, None) + if existing is not None: + _LOGGER.info("ML Flashpoint is already installed on this worker; skipping.") + return existing.checkpointer + + for attr in ("mcore_state", "model", "save_checkpoint"): + if not hasattr(worker, attr): + raise AttributeError( + f"Worker of type '{type(worker).__name__}' has no '{attr}'. This does not look like a NeMo RL " + "Megatron policy worker; ML Flashpoint cannot be installed on it." + ) + + checkpointer = MLFlashpointNeMoRLCheckpointer( + checkpoint_config=worker.mcore_state.cfg.checkpoint, + mlf_config=mlf_config, + ) + hooks = MLFlashpointWorkerHooks(checkpointer, mode=mode, durable_every_n_saves=durable_every_n_saves) + setattr(worker, _INSTALLED_ATTR, hooks) + + original_save = worker.save_checkpoint + + @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: + _LOGGER.info( + "Skipping the durable checkpoint at '%s': ML Flashpoint holds this step (mode='%s', " + "durable_every_n_saves=%d).", + weights_path, + hooks.mode, + hooks.durable_every_n_saves, + ) + return None + return original_save(weights_path, optimizer_path=optimizer_path, **kwargs) + + worker.save_checkpoint = save_checkpoint + _LOGGER.info( + "Installed ML Flashpoint on the NeMo RL Megatron policy worker (mode='%s', durable_every_n_saves=%d).", + mode, + durable_every_n_saves, + ) + return checkpointer + + +def _save_ml_flashpoint(worker, hooks: MLFlashpointWorkerHooks, include_optimizer: bool) -> None: + """Writes one ML Flashpoint checkpoint for a worker, best effort. + + Args: + worker: The NeMo RL Megatron policy worker. + hooks: The worker's installation record. + include_optimizer: Whether optimizer and scheduler state should be saved. + """ + try: + hooks.checkpointer.save( + state=worker.mcore_state, + model=worker.model, + optimizer=getattr(worker, "optimizer", None) if include_optimizer else None, + opt_param_scheduler=getattr(worker, "scheduler", None) if include_optimizer else None, + ) + except Exception: + _LOGGER.exception("ML Flashpoint save failed. Continuing; the durable checkpoint path is unaffected.") + + +def uninstall_from_worker(worker) -> None: + """Removes the ML Flashpoint wrapper and releases its resources. + + Args: + worker: A worker previously passed to :func:`install_into_worker`. + """ + hooks = getattr(worker, _INSTALLED_ATTR, None) + if hooks is None: + return + try: + hooks.checkpointer.shutdown() + finally: + # functools.wraps keeps __wrapped__ pointing at the bound original. + original = getattr(worker.save_checkpoint, "__wrapped__", None) + if original is not None: + worker.save_checkpoint = original + else: + # Fall back to the class method by dropping the instance attribute. + worker.__dict__.pop("save_checkpoint", None) + delattr(worker, _INSTALLED_ATTR) + _LOGGER.info("Removed ML Flashpoint from the NeMo RL Megatron policy worker.") + + +def install_from_env(worker) -> Optional[MLFlashpointNeMoRLCheckpointer]: + """Installs ML Flashpoint only when the environment asks for it. + + Reads ``MLFLASHPOINT_NEMO_RL_ENABLED`` (default false), + ``MLFLASHPOINT_NEMO_RL_MODE`` (default ``augment``) and + ``MLFLASHPOINT_NEMO_RL_DURABLE_EVERY_N_SAVES`` (default 1). This is what lets + a single launch command run both arms of an A/B experiment. + + Args: + worker: A NeMo RL ``MegatronPolicyWorker`` instance. + + Returns: + The checkpointer, or None when ML Flashpoint is not enabled. + """ + if not utils.get_env_val_bool("NEMO_RL_ENABLED", False): + _LOGGER.info("MLFLASHPOINT_NEMO_RL_ENABLED is not set; running without ML Flashpoint.") + return None + if not get_config().enabled: + _LOGGER.info("ML Flashpoint is disabled in its own config; running without it.") + return None + return install_into_worker( + worker, + mode=utils.get_env_val_str("NEMO_RL_MODE", MODE_AUGMENT), + durable_every_n_saves=utils.get_env_val_int("NEMO_RL_DURABLE_EVERY_N_SAVES", 1), + ) + + +def get_checkpointer(worker) -> Optional[MLFlashpointNeMoRLCheckpointer]: + """Returns the checkpointer installed on a worker, if any. + + Args: + worker: A NeMo RL ``MegatronPolicyWorker`` instance. + + Returns: + The checkpointer, or None. + """ + hooks: Optional[Any] = getattr(worker, _INSTALLED_ATTR, None) + return hooks.checkpointer if hooks is not None else None diff --git a/tests/adapter/conftest.py b/tests/adapter/conftest.py new file mode 100644 index 0000000..ce257f8 --- /dev/null +++ b/tests/adapter/conftest.py @@ -0,0 +1,139 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Stubs Megatron Bridge so the adapter can be unit tested without it. + +The Megatron Bridge adapter binds Bridge symbols at import time. Bridge itself +pulls in the full NVIDIA training stack (Transformer Engine, ModelOpt, NVRx), +which is neither installable nor runnable in a CPU-only test environment, so +this installs a minimal fake package tree when the real one is absent. + +Individual tests still patch these symbols where they are bound in the adapter +modules, so the fakes only have to satisfy ``import``. +""" + +import importlib.util +import sys +import types +from typing import Any + + +def _bridge_is_installed() -> bool: + """Whether the real Megatron Bridge can be imported. + + Returns: + True when ``megatron.bridge.training.checkpointing`` is importable. + """ + try: + return importlib.util.find_spec("megatron.bridge.training.checkpointing") is not None + except (ImportError, ValueError): + return False + + +class _FakeTrainState: + """Minimal stand-in for Megatron Bridge's ``TrainState``.""" + + def __init__(self, step: int = 0): + self.step = step + self.consumed_train_samples = 0 + self.skipped_train_samples = 0 + self.consumed_valid_samples = 0 + self.floating_point_operations_so_far = 0 + + def state_dict(self) -> dict[str, Any]: + """Returns the serializable counters.""" + return { + "step": self.step, + "consumed_train_samples": self.consumed_train_samples, + "skipped_train_samples": self.skipped_train_samples, + "consumed_valid_samples": self.consumed_valid_samples, + "floating_point_operations_so_far": self.floating_point_operations_so_far, + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Restores counters from ``state_dict``.""" + for key, value in state_dict.items(): + setattr(self, key, value) + + +def _module(name: str) -> types.ModuleType: + """Creates and registers a module. + + Args: + name: Fully qualified module name. + + Returns: + The registered module. + """ + module = types.ModuleType(name) + sys.modules[name] = module + return module + + +def _install_fake_bridge() -> None: + """Registers a fake ``megatron.bridge`` package tree in ``sys.modules``.""" + # `megatron` is a namespace package shared with the real megatron.core, so + # only reuse it -- never replace it. + megatron = sys.modules.get("megatron") + if megatron is None: + megatron = _module("megatron") + megatron.__path__ = [] + + bridge = _module("megatron.bridge") + bridge.__path__ = [] + megatron.bridge = bridge + + training = _module("megatron.bridge.training") + training.__path__ = [] + bridge.training = training + + checkpointing = _module("megatron.bridge.training.checkpointing") + checkpointing.init_checkpointing_context = lambda checkpoint_config: {} + checkpointing.load_checkpoint = lambda **kwargs: (0, 0) + checkpointing.maybe_finalize_async_save = lambda **kwargs: None + checkpointing.save_checkpoint = lambda **kwargs: None + checkpointing.generate_state_dict = lambda *args, **kwargs: {} + checkpointing.get_rng_state = lambda *args, **kwargs: None + checkpointing.set_checkpoint_version = lambda value: None + checkpointing._build_sharded_state_dict_metadata = lambda use_distributed_optimizer, cfg: {} + checkpointing._load_model_state_dict = lambda module, state_dict, strict: None + training.checkpointing = checkpointing + + state = _module("megatron.bridge.training.state") + state.TrainState = _FakeTrainState + training.state = state + + training_utils = _module("megatron.bridge.training.utils") + training_utils.__path__ = [] + training.utils = training_utils + + pg_utils = _module("megatron.bridge.training.utils.pg_utils") + pg_utils.get_pg_collection = lambda model: None + training_utils.pg_utils = pg_utils + + bridge_utils = _module("megatron.bridge.utils") + bridge_utils.__path__ = [] + bridge.utils = bridge_utils + + instantiate_utils = _module("megatron.bridge.utils.instantiate_utils") + instantiate_utils.registered_prefixes = [] + instantiate_utils.register_allowed_target_prefix = instantiate_utils.registered_prefixes.append + bridge_utils.instantiate_utils = instantiate_utils + + +MEGATRON_BRIDGE_IS_FAKE = not _bridge_is_installed() +"""Whether these tests run against the fake Bridge rather than the real one.""" + +if MEGATRON_BRIDGE_IS_FAKE: + _install_fake_bridge() diff --git a/tests/adapter/megatron_bridge/test_bridge_state.py b/tests/adapter/megatron_bridge/test_bridge_state.py new file mode 100644 index 0000000..6f11a6f --- /dev/null +++ b/tests/adapter/megatron_bridge/test_bridge_state.py @@ -0,0 +1,592 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import sys +import types +from typing import Any, Optional + +import pytest +import torch +from assertpy import assert_that + +from ml_flashpoint.adapter.megatron_bridge import bridge_state + + +@dataclasses.dataclass +class FakeCheckpointConfig: + ckpt_format: str = "torch_dist" + save_rng: bool = True + load_rng: bool = True + save_optim: bool = True + load_optim: bool = True + finetune: bool = False + + +@dataclasses.dataclass +class FakeTrainState: + step: int = 0 + consumed_train_samples: int = 0 + floating_point_operations_so_far: int = 0 + + def state_dict(self) -> dict[str, Any]: + return dataclasses.asdict(self) + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + for key, value in state_dict.items(): + setattr(self, key, value) + + +@dataclasses.dataclass +class FakeGlobalState: + cfg: Any + train_state: FakeTrainState = dataclasses.field(default_factory=FakeTrainState) + + +@dataclasses.dataclass +class FakeSaveContext: + state: FakeGlobalState + model: list + optimizer: Any = None + opt_param_scheduler: Any = None + num_floating_point_operations_so_far: int = 0 + train_data_iterator: Any = None + non_persistent_ckpt: bool = True + pg_collection: Any = None + module_name: Optional[str] = None + + +@dataclasses.dataclass +class FakeLoadContext: + state: FakeGlobalState + model: list + optimizer: Any = None + opt_param_scheduler: Any = None + strict: bool = True + skip_load_to_model_and_opt: bool = False + pg_collection: Any = None + module_name: Optional[str] = None + + +def _namespace(**kwargs): + return type("Namespace", (), kwargs)() + + +@pytest.fixture +def cfg(): + return _namespace( + checkpoint=FakeCheckpointConfig(), + optimizer=_namespace(use_distributed_optimizer=True), + rng=_namespace(data_parallel_random_init=False), + ) + + +@pytest.fixture +def pg_collection(mocker): + collection = mocker.MagicMock() + collection.dp.rank.return_value = 0 + return collection + + +@pytest.fixture +def save_ctx(cfg, pg_collection): + return FakeSaveContext( + state=FakeGlobalState(cfg=cfg, train_state=FakeTrainState(step=25)), + model=[object()], + num_floating_point_operations_so_far=987, + pg_collection=pg_collection, + ) + + +@pytest.fixture +def load_ctx(cfg, pg_collection): + return FakeLoadContext( + state=FakeGlobalState(cfg=cfg), + model=[object()], + pg_collection=pg_collection, + ) + + +@pytest.fixture(autouse=True) +def _bridge_helpers(mocker): + """Neutralizes the Bridge symbols the helpers call into.""" + mocker.patch.object(bridge_state, "get_rng_state", return_value="rng") + mocker.patch.object(bridge_state, "generate_state_dict", return_value={"model": {}}) + mocker.patch.object(bridge_state, "get_rerun_state_machine", return_value=mocker.MagicMock()) + mocker.patch.object(bridge_state, "unwrap_model", side_effect=lambda model: list(model)) + mocker.patch.object(bridge_state, "build_sharded_state_dict_metadata", return_value={}) + + +class TestUnwrapModel: + def test_prefers_megatron_core(self, mocker): + # Given + mocker.stopall() + module = types.ModuleType("megatron.core.utils") + module.unwrap_model = lambda model: ["unwrapped"] + mocker.patch.dict(sys.modules, {"megatron.core.utils": module}) + + # When + result = bridge_state.unwrap_model([object()]) + + # Then + assert_that(result).is_equal_to(["unwrapped"]) + + def test_falls_back_to_megatron_training(self, mocker): + # Given + mocker.stopall() + core_utils = types.ModuleType("megatron.core.utils") + training_utils = types.ModuleType("megatron.training.utils") + training_utils.unwrap_model = lambda model: ["legacy"] + mocker.patch.dict( + sys.modules, + {"megatron.core.utils": core_utils, "megatron.training.utils": training_utils}, + ) + + # When + result = bridge_state.unwrap_model([object()]) + + # Then + assert_that(result).is_equal_to(["legacy"]) + + def test_wraps_a_single_module_in_a_list(self, mocker): + # Given + mocker.stopall() + module = types.ModuleType("megatron.core.utils") + module.unwrap_model = lambda model: "single" + mocker.patch.dict(sys.modules, {"megatron.core.utils": module}) + + # When + result = bridge_state.unwrap_model(object()) + + # Then + assert_that(result).is_equal_to(["single"]) + + def test_raises_a_clear_error_when_unavailable(self, mocker): + # Given + mocker.stopall() + core_utils = types.ModuleType("megatron.core.utils") + mocker.patch.dict(sys.modules, {"megatron.core.utils": core_utils}) + mocker.patch.dict(sys.modules, {"megatron.training.utils": types.ModuleType("megatron.training.utils")}) + + # When/Then + with pytest.raises(RuntimeError, match="Could not resolve unwrap_model"): + bridge_state.unwrap_model([object()]) + + +class TestBuildSaveStateDict: + def test_embeds_train_state_flops_and_metadata(self, save_ctx): + # Given/When + state_dict = bridge_state.build_save_state_dict(save_ctx, step=25, num_floating_point_operations_so_far=987) + + # Then + assert_that(state_dict).contains_key(bridge_state.TRAIN_STATE_KEY) + assert_that(state_dict[bridge_state.FLOPS_KEY]).is_equal_to(987) + assert_that(state_dict).contains_key(bridge_state.CONTENT_METADATA_KEY) + + def test_passes_the_step_as_the_iteration(self, save_ctx): + # Given/When + bridge_state.build_save_state_dict(save_ctx, step=25, num_floating_point_operations_so_far=0) + + # Then + assert_that(bridge_state.generate_state_dict.call_args.kwargs["iteration"]).is_equal_to(25) + + def test_skips_rng_collection_when_disabled(self, save_ctx): + # Given + save_ctx.state.cfg.checkpoint.save_rng = False + + # When + bridge_state.build_save_state_dict(save_ctx, step=1, num_floating_point_operations_so_far=0) + + # Then + bridge_state.get_rng_state.assert_not_called() + + def test_supplies_the_process_group_to_sharding_but_not_the_checkpoint(self, save_ctx, pg_collection): + # Given/When + state_dict = bridge_state.build_save_state_dict(save_ctx, step=1, num_floating_point_operations_so_far=0) + + # Then + metadata = bridge_state.generate_state_dict.call_args.kwargs["optim_sd_kwargs"]["metadata"] + assert_that(metadata["dp_cp_group"]).is_equal_to(pg_collection.dp_cp) + assert_that(state_dict[bridge_state.CONTENT_METADATA_KEY]).does_not_contain_key("dp_cp_group") + + +class TestBuildLoadStateDict: + def test_marks_the_state_dict_as_loading(self, load_ctx): + # Given/When + bridge_state.build_load_state_dict(load_ctx) + + # Then + assert_that(bridge_state.generate_state_dict.call_args.kwargs["optim_sd_kwargs"]["is_loading"]).is_true() + + def test_omits_the_optimizer_when_load_optim_is_off(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock() + load_ctx.state.cfg.checkpoint.load_optim = False + + # When + bridge_state.build_load_state_dict(load_ctx) + + # Then + assert_that(bridge_state.generate_state_dict.call_args.args[2]).is_none() + + def test_omits_the_optimizer_when_finetuning(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock() + load_ctx.state.cfg.checkpoint.finetune = True + + # When + bridge_state.build_load_state_dict(load_ctx) + + # Then + assert_that(bridge_state.generate_state_dict.call_args.args[2]).is_none() + + def test_includes_the_optimizer_on_a_plain_resume(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock() + + # When + bridge_state.build_load_state_dict(load_ctx) + + # Then + assert_that(bridge_state.generate_state_dict.call_args.args[2]).is_same_as(load_ctx.optimizer) + + +class TestRestoreTrainState: + def test_restores_embedded_counters(self, cfg): + # Given + state = FakeGlobalState(cfg=cfg) + state_dict = {bridge_state.TRAIN_STATE_KEY: {"step": 77, "consumed_train_samples": 1024}} + + # When + bridge_state.restore_train_state(state, state_dict) + + # Then + assert_that(state.train_state.step).is_equal_to(77) + assert_that(state.train_state.consumed_train_samples).is_equal_to(1024) + + def test_falls_back_to_the_iteration_key(self, cfg): + # Given + state = FakeGlobalState(cfg=cfg) + + # When + bridge_state.restore_train_state(state, {"iteration": 12}) + + # Then + assert_that(state.train_state.step).is_equal_to(12) + + def test_restores_flops(self, cfg): + # Given + state = FakeGlobalState(cfg=cfg) + + # When + bridge_state.restore_train_state(state, {"iteration": 1, bridge_state.FLOPS_KEY: 4321}) + + # Then + assert_that(state.train_state.floating_point_operations_so_far).is_equal_to(4321) + + +class TestRestoreModel: + def test_single_chunk(self, mocker): + # Given + loader = mocker.MagicMock() + mocker.patch.object(bridge_state, "_import_optional", return_value=loader) + model = [object()] + + # When + bridge_state.restore_model(model, {"model": {"w": 1}}, strict=True) + + # Then + loader.assert_called_once_with(model[0], {"w": 1}, True) + + def test_multiple_chunks(self, mocker): + # Given + loader = mocker.MagicMock() + mocker.patch.object(bridge_state, "_import_optional", return_value=loader) + model = [object(), object()] + + # When + bridge_state.restore_model(model, {"model0": {"a": 1}, "model1": {"b": 2}}, strict=False) + + # Then + assert_that(loader.call_count).is_equal_to(2) + + def test_skips_empty_pipeline_stages(self, mocker): + # Given + loader = mocker.MagicMock() + mocker.patch.object(bridge_state, "_import_optional", return_value=loader) + model = [object(), object()] + + # When + bridge_state.restore_model(model, {"model0": {"a": 1}}, strict=True) + + # Then + assert_that(loader.call_count).is_equal_to(1) + + def test_raises_when_the_bridge_helper_is_missing(self, mocker): + # Given + mocker.patch.object(bridge_state, "_import_optional", return_value=None) + + # When/Then + with pytest.raises(RuntimeError, match="_load_model_state_dict"): + bridge_state.restore_model([object()], {"model": {}}, strict=True) + + +class TestRestoreOptimizer: + def test_restores_optimizer_and_scheduler(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock(is_stub_optimizer=False) + load_ctx.opt_param_scheduler = mocker.MagicMock() + + # When + bridge_state.restore_optimizer(load_ctx, {"optimizer": {"o": 1}, "opt_param_scheduler": {"s": 2}}) + + # Then + load_ctx.optimizer.load_state_dict.assert_called_once_with({"o": 1}) + load_ctx.opt_param_scheduler.load_state_dict.assert_called_once_with({"s": 2}) + + def test_prefers_the_legacy_lr_scheduler_key(self, load_ctx, mocker): + # Given + load_ctx.opt_param_scheduler = mocker.MagicMock() + + # When + bridge_state.restore_optimizer(load_ctx, {"lr_scheduler": {"legacy": True}}) + + # Then + load_ctx.opt_param_scheduler.load_state_dict.assert_called_once_with({"legacy": True}) + + def test_skips_stub_optimizers(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock(is_stub_optimizer=True) + + # When + bridge_state.restore_optimizer(load_ctx, {"optimizer": {"o": 1}}) + + # Then + load_ctx.optimizer.load_state_dict.assert_not_called() + + def test_skips_everything_when_finetuning(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock(is_stub_optimizer=False) + load_ctx.opt_param_scheduler = mocker.MagicMock() + load_ctx.state.cfg.checkpoint.finetune = True + + # When + bridge_state.restore_optimizer(load_ctx, {"optimizer": {}, "opt_param_scheduler": {}}) + + # Then + load_ctx.optimizer.load_state_dict.assert_not_called() + load_ctx.opt_param_scheduler.load_state_dict.assert_not_called() + + def test_skips_the_optimizer_when_skipping_model_load(self, load_ctx, mocker): + # Given + load_ctx.optimizer = mocker.MagicMock(is_stub_optimizer=False) + load_ctx.skip_load_to_model_and_opt = True + + # When + bridge_state.restore_optimizer(load_ctx, {"optimizer": {"o": 1}}) + + # Then + load_ctx.optimizer.load_state_dict.assert_not_called() + + +class TestRestoreRerunState: + def test_restores_when_present(self, mocker): + # Given + machine = mocker.MagicMock() + mocker.patch.object(bridge_state, "get_rerun_state_machine", return_value=machine) + + # When + bridge_state.restore_rerun_state({"rerun_state_machine": {"x": 1}}) + + # Then + machine.load_state_dict.assert_called_once_with({"x": 1}) + + def test_absent_key_is_a_noop(self, mocker): + # Given + machine = mocker.MagicMock() + mocker.patch.object(bridge_state, "get_rerun_state_machine", return_value=machine) + + # When + bridge_state.restore_rerun_state({}) + + # Then + machine.load_state_dict.assert_not_called() + + def test_failure_is_swallowed(self, mocker): + # Given + machine = mocker.MagicMock() + machine.load_state_dict.side_effect = RuntimeError("incompatible") + mocker.patch.object(bridge_state, "get_rerun_state_machine", return_value=machine) + + # When + bridge_state.restore_rerun_state({"rerun_state_machine": {}}) + + # Then no exception escapes. + + +class TestRestoreRngState: + def _rng_state(self): + return { + "random_rng_state": __import__("random").getstate(), + "np_rng_state": __import__("numpy").random.get_state(), + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.get_rng_state(), + "rng_tracker_states": {"tracker": "state"}, + } + + def test_restores_the_data_parallel_rank_slot(self, load_ctx, pg_collection, mocker): + # Given + load_ctx.state.cfg.rng.data_parallel_random_init = True + pg_collection.dp.rank.return_value = 1 + tracker = mocker.MagicMock() + mocker.patch.object(bridge_state.tensor_parallel, "get_cuda_rng_tracker", return_value=tracker) + mocker.patch.object( + bridge_state.tensor_parallel, "is_graph_safe_cuda_rng_tracker", return_value=False, create=True + ) + mocker.patch.object( + bridge_state.tensor_parallel, "convert_cuda_rng_state", side_effect=lambda v, **_: v, create=True + ) + mocker.patch.object(torch.cuda, "set_rng_state") + state_dict = {"rng_state": [self._rng_state(), self._rng_state()]} + + # When + bridge_state.restore_rng_state(load_ctx, state_dict, pg_collection) + + # Then + tracker.set_states.assert_called_once_with({"tracker": "state"}) + + def test_skips_when_load_rng_is_off(self, load_ctx, pg_collection, mocker): + # Given + load_ctx.state.cfg.checkpoint.load_rng = False + tracker = mocker.MagicMock() + mocker.patch.object(bridge_state.tensor_parallel, "get_cuda_rng_tracker", return_value=tracker) + + # When + bridge_state.restore_rng_state(load_ctx, {"rng_state": [self._rng_state()]}, pg_collection) + + # Then + tracker.set_states.assert_not_called() + + def test_missing_rng_state_is_a_noop(self, load_ctx, pg_collection, mocker): + # Given + tracker = mocker.MagicMock() + mocker.patch.object(bridge_state.tensor_parallel, "get_cuda_rng_tracker", return_value=tracker) + + # When + bridge_state.restore_rng_state(load_ctx, {}, pg_collection) + + # Then + tracker.set_states.assert_not_called() + + def test_failure_is_swallowed(self, load_ctx, pg_collection, mocker): + # Given a payload missing the keys the restore path needs. + mocker.patch.object(bridge_state.tensor_parallel, "get_cuda_rng_tracker", return_value=mocker.MagicMock()) + + # When + bridge_state.restore_rng_state(load_ctx, {"rng_state": [{}]}, pg_collection) + + # Then no exception escapes. + + +class TestApplyLoadedState: + def test_returns_step_and_flops(self, load_ctx, mocker): + # Given + mocker.patch.object(bridge_state, "set_checkpoint_version") + mocker.patch.object(bridge_state, "update_num_microbatches") + mocker.patch.object(bridge_state, "restore_model") + mocker.patch.object(bridge_state, "restore_optimizer") + mocker.patch.object(bridge_state, "restore_rerun_state") + mocker.patch.object(bridge_state, "restore_rng_state") + state_dict = {bridge_state.TRAIN_STATE_KEY: {"step": 88}, bridge_state.FLOPS_KEY: 2048} + + # When + step, flops = bridge_state.apply_loaded_state(load_ctx, state_dict) + + # Then + assert_that(step).is_equal_to(88) + assert_that(flops).is_equal_to(2048) + + def test_skips_the_model_when_asked_to(self, load_ctx, mocker): + # Given + mocker.patch.object(bridge_state, "set_checkpoint_version") + mocker.patch.object(bridge_state, "update_num_microbatches") + restore_model = mocker.patch.object(bridge_state, "restore_model") + mocker.patch.object(bridge_state, "restore_optimizer") + mocker.patch.object(bridge_state, "restore_rerun_state") + mocker.patch.object(bridge_state, "restore_rng_state") + load_ctx.skip_load_to_model_and_opt = True + + # When + bridge_state.apply_loaded_state(load_ctx, {bridge_state.TRAIN_STATE_KEY: {"step": 1}}) + + # Then + restore_model.assert_not_called() + + +class TestMisc: + def test_get_content_metadata(self): + # Given/When/Then + assert_that(bridge_state.get_content_metadata({bridge_state.CONTENT_METADATA_KEY: {"v": 1}})).is_equal_to( + {"v": 1} + ) + assert_that(bridge_state.get_content_metadata({})).is_none() + + def test_resolve_pg_collection_prefers_the_explicit_value(self, mocker): + # Given + explicit = mocker.MagicMock() + get_pg = mocker.patch.object(bridge_state, "get_pg_collection") + + # When + result = bridge_state.resolve_pg_collection([object()], explicit) + + # Then + assert_that(result).is_same_as(explicit) + get_pg.assert_not_called() + + def test_resolve_pg_collection_falls_back_to_the_model(self, mocker): + # Given + derived = mocker.MagicMock() + mocker.patch.object(bridge_state, "get_pg_collection", return_value=derived) + + # When + result = bridge_state.resolve_pg_collection([object()], None) + + # Then + assert_that(result).is_same_as(derived) + + def test_build_sharded_state_dict_metadata_requires_the_bridge_helper(self, mocker): + # Given + mocker.stopall() + mocker.patch.object(bridge_state, "_import_optional", return_value=None) + + # When/Then + with pytest.raises(RuntimeError, match="_build_sharded_state_dict_metadata"): + bridge_state.build_sharded_state_dict_metadata(True, object()) + + def test_clean_metadata_drops_process_groups_without_the_helper(self, mocker): + # Given a Megatron build that predates the serialization helper. + real_import = bridge_state.importlib.import_module + + def fake_import(name, *args, **kwargs): + if name == "megatron.core.dist_checkpointing.utils": + raise ImportError("not available") + return real_import(name, *args, **kwargs) + + mocker.patch.object(bridge_state.importlib, "import_module", side_effect=fake_import) + + # When + result = bridge_state.clean_metadata_for_serialization({"keep": 1}) + + # Then + assert_that(result).is_equal_to({"keep": 1}) diff --git a/tests/adapter/megatron_bridge/test_checkpoint_manager.py b/tests/adapter/megatron_bridge/test_checkpoint_manager.py new file mode 100644 index 0000000..46ef1ef --- /dev/null +++ b/tests/adapter/megatron_bridge/test_checkpoint_manager.py @@ -0,0 +1,490 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +from typing import Any, Optional + +import pytest +from assertpy import assert_that + +from ml_flashpoint.adapter.megatron_bridge import checkpoint_manager as manager_module +from ml_flashpoint.adapter.megatron_bridge.checkpoint_manager import ( + MLFlashpointBridgeCheckpointManager, + init_checkpointing_context_safely, +) +from ml_flashpoint.adapter.megatron_bridge.config import MLFlashpointBridgeConfig +from ml_flashpoint.adapter.megatron_bridge.local_checkpoint_index import ( + NO_CHECKPOINT, + MLFlashpointLocalCheckpointIndex, +) +from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId + + +@dataclasses.dataclass +class FakeCheckpointConfig: + """The subset of Megatron Bridge's ``CheckpointConfig`` the adapter reads.""" + + save: str = "/durable/checkpoints" + ckpt_format: str = "torch_dist" + non_persistent_ckpt_type: str = "local" + non_persistent_save_interval: int = 10 + save_interval: int = 100 + custom_manager_class: Optional[str] = None + + +@dataclasses.dataclass +class FakeTrainState: + step: int = 0 + floating_point_operations_so_far: int = 0 + + +@dataclasses.dataclass +class FakeGlobalState: + cfg: Any = None + train_state: FakeTrainState = dataclasses.field(default_factory=FakeTrainState) + + +@dataclasses.dataclass +class FakeSaveContext: + state: FakeGlobalState + model: list + optimizer: Any = None + opt_param_scheduler: Any = None + num_floating_point_operations_so_far: int = 0 + train_data_iterator: Any = None + non_persistent_ckpt: bool = False + pg_collection: Any = None + module_name: Optional[str] = None + + +@dataclasses.dataclass +class FakeLoadContext: + state: FakeGlobalState + model: list + optimizer: Any = None + opt_param_scheduler: Any = None + strict: bool = True + skip_load_to_model_and_opt: bool = False + pg_collection: Any = None + module_name: Optional[str] = None + + +@pytest.fixture +def checkpoint_config() -> FakeCheckpointConfig: + return FakeCheckpointConfig() + + +@pytest.fixture +def mlf_config(tmp_path) -> MLFlashpointBridgeConfig: + return MLFlashpointBridgeConfig(base_container=str(tmp_path / "mlf")) + + +@pytest.fixture +def runtime(mocker, mlf_config): + """A runtime double whose base container matches the config under test.""" + fake = mocker.MagicMock() + fake.base_container = CheckpointContainerId(mlf_config.base_container) + fake.config = mlf_config + return fake + + +@pytest.fixture +def global_state(checkpoint_config) -> FakeGlobalState: + cfg = mocker_namespace(checkpoint=checkpoint_config) + return FakeGlobalState(cfg=cfg, train_state=FakeTrainState(step=40)) + + +def mocker_namespace(**kwargs): + """Builds a lightweight attribute container for nested config objects.""" + return type("Namespace", (), kwargs)() + + +@pytest.fixture +def manager(checkpoint_config, mlf_config, runtime, mocker) -> MLFlashpointBridgeCheckpointManager: + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + return MLFlashpointBridgeCheckpointManager( + checkpoint_config=checkpoint_config, + mlf_config=mlf_config, + runtime=runtime, + ) + + +class TestEnablement: + def test_enabled_by_default(self, manager): + # Given/When/Then + assert_that(manager.enabled).is_true() + + def test_disabled_when_config_disabled(self, checkpoint_config, runtime, mocker, tmp_path): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + config = MLFlashpointBridgeConfig(enabled=False, base_container=str(tmp_path / "mlf")) + + # When + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=config, runtime=runtime) + + # Then + assert_that(manager.enabled).is_false() + + @pytest.mark.parametrize("ckpt_format", ["torch", "fsdp_dtensor"]) + def test_disabled_for_unsupported_checkpoint_format(self, ckpt_format, mlf_config, runtime, mocker): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + checkpoint_config = FakeCheckpointConfig(ckpt_format=ckpt_format) + + # When + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=mlf_config, runtime=runtime) + + # Then + assert_that(manager.enabled).is_false() + + def test_installs_local_index_into_context(self, manager): + # Given/When + context = manager.checkpointing_context + + # Then + assert_that(context).contains_key("local_checkpoint_manager") + assert_that(context["local_checkpoint_manager"]).is_instance_of(MLFlashpointLocalCheckpointIndex) + + def test_disabled_manager_does_not_install_local_index(self, checkpoint_config, runtime, mocker, tmp_path): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + config = MLFlashpointBridgeConfig(enabled=False, base_container=str(tmp_path / "mlf")) + + # When + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=config, runtime=runtime) + + # Then + assert_that(manager.checkpointing_context).does_not_contain_key("local_checkpoint_manager") + + +class TestSaveRouting: + def test_persistent_save_is_delegated(self, manager, global_state, mocker): + # Given + delegate = mocker.patch.object(manager_module, "save_checkpoint") + mlf_save = mocker.patch.object(manager_module, "save_local_aware_megatron_checkpoint") + ctx = FakeSaveContext(state=global_state, model=[object()], non_persistent_ckpt=False) + + # When + manager.save(ctx, callback_manager=None) + + # Then + delegate.assert_called_once() + mlf_save.assert_not_called() + + def test_non_persistent_save_goes_to_ml_flashpoint(self, manager, global_state, mocker): + # Given + delegate = mocker.patch.object(manager_module, "save_checkpoint") + mocker.patch.object(manager_module.bridge_state, "build_save_state_dict", return_value={"model": {}}) + mlf_save = mocker.patch.object( + manager_module, "save_local_aware_megatron_checkpoint", return_value="async-request" + ) + ctx = FakeSaveContext(state=global_state, model=[object()], non_persistent_ckpt=True) + + # When + manager.save(ctx, callback_manager=None) + + # Then + delegate.assert_not_called() + mlf_save.assert_called_once() + assert_that(mlf_save.call_args.kwargs["checkpoint_dir"]).ends_with("step-40_ckpt") + manager.runtime.schedule.assert_called_once_with("async-request") + + def test_synchronous_save_schedules_nothing(self, checkpoint_config, runtime, global_state, mocker, tmp_path): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + config = MLFlashpointBridgeConfig(async_save=False, base_container=str(tmp_path / "mlf")) + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=config, runtime=runtime) + mocker.patch.object(manager_module.bridge_state, "build_save_state_dict", return_value={}) + mocker.patch.object(manager_module, "save_local_aware_megatron_checkpoint", return_value=None) + ctx = FakeSaveContext(state=global_state, model=[object()], non_persistent_ckpt=True) + + # When + manager.save(ctx, callback_manager=None) + + # Then + runtime.schedule.assert_not_called() + + def test_ml_flashpoint_failure_does_not_propagate(self, manager, global_state, mocker): + # Given + mocker.patch.object( + manager_module.bridge_state, "build_save_state_dict", side_effect=RuntimeError("out of buffers") + ) + ctx = FakeSaveContext(state=global_state, model=[object()], non_persistent_ckpt=True) + + # When + manager.save(ctx, callback_manager=None) + + # Then no exception escapes; the training loop keeps going. + + def test_disabled_manager_delegates_non_persistent_save( + self, checkpoint_config, runtime, global_state, mocker, tmp_path + ): + # Given a run that opted out of ML Flashpoint but kept the local cadence. + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + config = MLFlashpointBridgeConfig(enabled=False, base_container=str(tmp_path / "mlf")) + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=config, runtime=runtime) + ctx = FakeSaveContext(state=global_state, model=[object()], non_persistent_ckpt=True) + + # When/Then: Bridge cannot service a local checkpoint ML Flashpoint owns. + with pytest.raises(RuntimeError, match="non_persistent_ckpt_type='local'"): + manager.save(ctx, callback_manager=None) + + def test_disabled_manager_delegates_global_non_persistent_save(self, runtime, mocker, tmp_path): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + delegate = mocker.patch.object(manager_module, "save_checkpoint") + checkpoint_config = FakeCheckpointConfig(non_persistent_ckpt_type="global") + config = MLFlashpointBridgeConfig(enabled=False, base_container=str(tmp_path / "mlf")) + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=config, runtime=runtime) + state = FakeGlobalState(cfg=mocker_namespace(checkpoint=checkpoint_config), train_state=FakeTrainState(1)) + ctx = FakeSaveContext(state=state, model=[object()], non_persistent_ckpt=True) + + # When + manager.save(ctx, callback_manager=None) + + # Then + delegate.assert_called_once() + + def test_save_invalidates_cached_resume_decision(self, manager, global_state, mocker): + # Given + mocker.patch.object(manager_module.bridge_state, "build_save_state_dict", return_value={}) + mocker.patch.object(manager_module, "save_local_aware_megatron_checkpoint", return_value=None) + index = manager.checkpointing_context["local_checkpoint_manager"] + invalidate = mocker.spy(index, "invalidate") + ctx = FakeSaveContext(state=global_state, model=[object()], non_persistent_ckpt=True) + + # When + manager.save(ctx, callback_manager=None) + + # Then + assert_that(invalidate.call_count).is_equal_to(1) + + +class TestLoadRouting: + def test_prefers_ml_flashpoint_checkpoint(self, manager, global_state, mocker): + # Given + container = CheckpointContainerId(str(manager.runtime.base_container) + "/step-40_ckpt") + mocker.patch.object(manager, "_find_ml_flashpoint_checkpoint", return_value=container) + mocker.patch.object(manager, "_load_ml_flashpoint", return_value=(40, 1234)) + delegate = mocker.patch.object(manager_module, "load_checkpoint") + ctx = FakeLoadContext(state=global_state, model=[object()]) + + # When + result = manager.load(ctx) + + # Then + assert_that(result).is_equal_to((40, 1234)) + delegate.assert_not_called() + + def test_falls_back_when_no_ml_flashpoint_checkpoint(self, manager, global_state, mocker): + # Given + mocker.patch.object(manager, "_find_ml_flashpoint_checkpoint", return_value=None) + delegate = mocker.patch.object(manager_module, "load_checkpoint", return_value=(7, 99)) + ctx = FakeLoadContext(state=global_state, model=[object()]) + + # When + result = manager.load(ctx) + + # Then + assert_that(result).is_equal_to((7, 99)) + delegate.assert_called_once() + + def test_falls_back_when_ml_flashpoint_read_fails(self, manager, global_state, mocker): + # Given + container = CheckpointContainerId(str(manager.runtime.base_container) + "/step-40_ckpt") + mocker.patch.object(manager, "_find_ml_flashpoint_checkpoint", return_value=container) + mocker.patch.object(manager, "_load_ml_flashpoint", return_value=None) + delegate = mocker.patch.object(manager_module, "load_checkpoint", return_value=(0, 0)) + ctx = FakeLoadContext(state=global_state, model=[object()]) + + # When + result = manager.load(ctx) + + # Then + assert_that(result).is_equal_to((0, 0)) + delegate.assert_called_once() + + def test_fallback_hides_the_local_index_from_bridge(self, manager, global_state, mocker): + # Given: Bridge's own local path cannot read an ML Flashpoint container. + mocker.patch.object(manager, "_find_ml_flashpoint_checkpoint", return_value=None) + mocker.patch.object(manager_module, "load_checkpoint", return_value=(0, 0)) + index = manager.checkpointing_context["local_checkpoint_manager"] + ctx = FakeLoadContext(state=global_state, model=[object()]) + + # When + manager.load(ctx) + + # Then + assert_that(index.find_latest()).is_equal_to(NO_CHECKPOINT) + + def test_read_failure_returns_none_for_fallback(self, manager, global_state, mocker): + # Given + container = CheckpointContainerId(str(manager.runtime.base_container) + "/step-40_ckpt") + mocker.patch.object( + manager_module.bridge_state, "build_load_state_dict", side_effect=RuntimeError("missing objects") + ) + ctx = FakeLoadContext(state=global_state, model=[object()]) + + # When + result = manager._load_ml_flashpoint(ctx, container) + + # Then + assert_that(result).is_none() + + def test_successful_read_applies_state(self, manager, global_state, mocker): + # Given + container = CheckpointContainerId(str(manager.runtime.base_container) + "/step-40_ckpt") + mocker.patch.object(manager_module.bridge_state, "build_load_state_dict", return_value={}) + mocker.patch.object(manager_module.mcore_dist_checkpointing, "load", return_value={"model": {}}) + apply_state = mocker.patch.object(manager_module.bridge_state, "apply_loaded_state", return_value=(40, 5)) + ctx = FakeLoadContext(state=global_state, model=[object()]) + + # When + result = manager._load_ml_flashpoint(ctx, container) + + # Then + assert_that(result).is_equal_to((40, 5)) + apply_state.assert_called_once() + + +class TestFinalizeAndShutdown: + def test_finalizes_both_queues(self, manager, global_state, mocker): + # Given + bridge_finalize = mocker.patch.object(manager_module, "maybe_finalize_async_save") + + # When + manager.finalize_async_saves(global_state, blocking=False) + + # Then + manager.runtime.maybe_finalize.assert_called_once_with(blocking=False) + bridge_finalize.assert_called_once() + + def test_ml_flashpoint_finalize_failure_still_finalizes_bridge(self, manager, global_state, mocker): + # Given + manager.runtime.maybe_finalize.side_effect = RuntimeError("worker died") + bridge_finalize = mocker.patch.object(manager_module, "maybe_finalize_async_save") + + # When + manager.finalize_async_saves(global_state, blocking=True) + + # Then + bridge_finalize.assert_called_once() + + def test_terminate_shuts_the_runtime_down(self, manager, global_state, mocker): + # Given + mocker.patch.object(manager_module, "maybe_finalize_async_save") + shutdown = mocker.patch("ml_flashpoint.adapter.megatron_bridge.runtime.shutdown_runtime") + mocker.patch.object(manager_module.dist, "is_initialized", return_value=False) + + # When + manager.finalize_async_saves(global_state, blocking=True, terminate=True) + + # Then + shutdown.assert_called_once_with(remove_checkpoints=True) + assert_that(manager.runtime).is_none() + + def test_shutdown_keeps_checkpoints_when_configured(self, checkpoint_config, runtime, mocker, tmp_path): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + mocker.patch.object(manager_module.dist, "is_initialized", return_value=False) + shutdown = mocker.patch("ml_flashpoint.adapter.megatron_bridge.runtime.shutdown_runtime") + config = MLFlashpointBridgeConfig(base_container=str(tmp_path / "mlf"), keep_checkpoints_on_finalize=True) + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=config, runtime=runtime) + + # When + manager.shutdown() + + # Then + shutdown.assert_called_once_with(remove_checkpoints=False) + + def test_shutdown_without_runtime_is_a_noop(self, checkpoint_config, mlf_config, mocker): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + shutdown = mocker.patch("ml_flashpoint.adapter.megatron_bridge.runtime.shutdown_runtime") + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=mlf_config) + + # When + manager.shutdown() + + # Then + shutdown.assert_not_called() + + def test_shutdown_clears_the_local_index_from_context(self, manager, mocker): + # Given + mocker.patch.object(manager_module.dist, "is_initialized", return_value=False) + mocker.patch("ml_flashpoint.adapter.megatron_bridge.runtime.shutdown_runtime") + assert_that(manager.checkpointing_context).contains_key("local_checkpoint_manager") + + # When + manager.shutdown() + + # Then + assert_that(manager._context).does_not_contain_key("local_checkpoint_manager") + + +class TestRuntimeBootstrap: + def test_runtime_is_built_lazily(self, checkpoint_config, mlf_config, mocker): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + built = mocker.MagicMock() + built.base_container = CheckpointContainerId(mlf_config.base_container) + get_runtime = mocker.patch.object(manager_module, "get_runtime", return_value=built) + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=mlf_config) + + # When + assert_that(manager.runtime).is_none() + result = manager._ensure_runtime() + + # Then + assert_that(result).is_same_as(built) + get_runtime.assert_called_once_with(mlf_config) + + def test_runtime_failure_disables_ml_flashpoint(self, checkpoint_config, mlf_config, mocker): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={}) + mocker.patch.object(manager_module, "get_runtime", side_effect=RuntimeError("no process group")) + manager = MLFlashpointBridgeCheckpointManager(checkpoint_config, mlf_config=mlf_config) + + # When + result = manager._ensure_runtime() + + # Then + assert_that(result).is_none() + assert_that(manager.enabled).is_false() + + +class TestInitCheckpointingContextSafely: + def test_passes_through(self, mocker, checkpoint_config): + # Given + mocker.patch.object(manager_module, "init_checkpointing_context", return_value={"a": 1}) + + # When + result = init_checkpointing_context_safely(checkpoint_config) + + # Then + assert_that(result).is_equal_to({"a": 1}) + + def test_missing_nvrx_yields_empty_context(self, mocker, checkpoint_config): + # Given + mocker.patch.object( + manager_module, + "init_checkpointing_context", + side_effect=RuntimeError("nvidia_resiliency_ext is required"), + ) + + # When + result = init_checkpointing_context_safely(checkpoint_config) + + # Then + assert_that(result).is_equal_to({}) diff --git a/tests/adapter/megatron_bridge/test_config.py b/tests/adapter/megatron_bridge/test_config.py new file mode 100644 index 0000000..448646e --- /dev/null +++ b/tests/adapter/megatron_bridge/test_config.py @@ -0,0 +1,175 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from assertpy import assert_that + +from ml_flashpoint.adapter.megatron_bridge import config as config_module +from ml_flashpoint.adapter.megatron_bridge.config import ( + DEFAULT_BASE_CONTAINER, + MLFlashpointBridgeConfig, + configure, + get_config, + reset_configuration, +) +from ml_flashpoint.core.checkpoint_saver import DEFAULT_INITIAL_BUFFER_SIZE_BYTES + + +@pytest.fixture(autouse=True) +def _clean_registration(): + reset_configuration() + yield + reset_configuration() + + +class TestMLFlashpointBridgeConfig: + def test_defaults(self): + # Given/When + config = MLFlashpointBridgeConfig() + + # Then + assert_that(config.enabled).is_true() + assert_that(config.base_container).is_equal_to(DEFAULT_BASE_CONTAINER) + assert_that(config.async_save).is_true() + assert_that(config.write_thread_count).is_equal_to(1) + assert_that(config.initial_write_buffer_size_bytes).is_equal_to(DEFAULT_INITIAL_BUFFER_SIZE_BYTES) + assert_that(config.use_optimized_save).is_true() + assert_that(config.use_cached_ckpt_structure).is_false() + assert_that(config.use_fully_parallel_wrapper).is_true() + assert_that(config.keep_checkpoints_on_finalize).is_false() + + def test_is_frozen(self): + # Given + config = MLFlashpointBridgeConfig() + + # When/Then + with pytest.raises(Exception): + config.enabled = False + + @pytest.mark.parametrize( + "kwargs", + [ + {"base_container": ""}, + {"write_thread_count": 0}, + {"write_thread_count": -3}, + {"initial_write_buffer_size_bytes": 0}, + {"initial_write_buffer_size_bytes": -1}, + ], + ) + def test_rejects_invalid_values(self, kwargs): + # Given/When/Then + with pytest.raises(ValueError): + MLFlashpointBridgeConfig(**kwargs) + + +class TestFromEnv: + def test_uses_defaults_when_env_is_empty(self, monkeypatch): + # Given + for name in list(os_environ_keys()): + monkeypatch.delenv(name, raising=False) + + # When + config = MLFlashpointBridgeConfig.from_env() + + # Then + assert_that(config).is_equal_to(MLFlashpointBridgeConfig()) + + def test_reads_every_field(self, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_BRIDGE_ENABLED", "false") + monkeypatch.setenv("MLFLASHPOINT_BASE_CONTAINER", "/mnt/local/mlf") + monkeypatch.setenv("MLFLASHPOINT_ASYNC_SAVE", "false") + monkeypatch.setenv("MLFLASHPOINT_WRITE_THREAD_COUNT", "4") + monkeypatch.setenv("MLFLASHPOINT_INITIAL_WRITE_BUFFER_SIZE_BYTES", "2048") + monkeypatch.setenv("MLFLASHPOINT_USE_OPTIMIZED_SAVE", "false") + monkeypatch.setenv("MLFLASHPOINT_USE_CACHED_CKPT_STRUCTURE", "true") + monkeypatch.setenv("MLFLASHPOINT_USE_FULLY_PARALLEL_WRAPPER", "false") + monkeypatch.setenv("MLFLASHPOINT_KEEP_CHECKPOINTS_ON_FINALIZE", "true") + + # When + config = MLFlashpointBridgeConfig.from_env() + + # Then + assert_that(config.enabled).is_false() + assert_that(config.base_container).is_equal_to("/mnt/local/mlf") + assert_that(config.async_save).is_false() + assert_that(config.write_thread_count).is_equal_to(4) + assert_that(config.initial_write_buffer_size_bytes).is_equal_to(2048) + assert_that(config.use_optimized_save).is_false() + assert_that(config.use_cached_ckpt_structure).is_true() + assert_that(config.use_fully_parallel_wrapper).is_false() + assert_that(config.keep_checkpoints_on_finalize).is_true() + + def test_non_integer_value_falls_back_to_default(self, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_WRITE_THREAD_COUNT", "not-a-number") + + # When + config = MLFlashpointBridgeConfig.from_env() + + # Then + assert_that(config.write_thread_count).is_equal_to(1) + + def test_invalid_env_value_still_validated(self, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_WRITE_THREAD_COUNT", "0") + + # When/Then + with pytest.raises(ValueError): + MLFlashpointBridgeConfig.from_env() + + +class TestConfigure: + def test_get_config_reads_env_when_unregistered(self, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_BASE_CONTAINER", "/from/env") + + # When + config = get_config() + + # Then + assert_that(config.base_container).is_equal_to("/from/env") + + def test_registered_config_wins_over_env(self, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_BASE_CONTAINER", "/from/env") + configure(MLFlashpointBridgeConfig(base_container="/registered")) + + # When + config = get_config() + + # Then + assert_that(config.base_container).is_equal_to("/registered") + + def test_reset_restores_env_lookup(self, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_BASE_CONTAINER", "/from/env") + configure(MLFlashpointBridgeConfig(base_container="/registered")) + + # When + reset_configuration() + + # Then + assert_that(get_config().base_container).is_equal_to("/from/env") + + def test_module_state_starts_clean(self): + # Given/When/Then + assert_that(config_module._CONFIGURED).is_none() + + +def os_environ_keys(): + """Returns every ``MLFLASHPOINT_`` variable currently set.""" + import os + + return [name for name in os.environ if name.startswith("MLFLASHPOINT_")] diff --git a/tests/adapter/megatron_bridge/test_enable.py b/tests/adapter/megatron_bridge/test_enable.py new file mode 100644 index 0000000..4e55ae4 --- /dev/null +++ b/tests/adapter/megatron_bridge/test_enable.py @@ -0,0 +1,137 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import sys +import types +from typing import Optional + +import pytest +from assertpy import assert_that + +import ml_flashpoint.adapter.megatron_bridge as adapter +from ml_flashpoint.adapter.megatron_bridge import CUSTOM_MANAGER_CLASS, MLFlashpointBridgeConfig +from ml_flashpoint.adapter.megatron_bridge.config import reset_configuration + + +@dataclasses.dataclass +class FakeCheckpointConfig: + save: str = "/durable/checkpoints" + save_interval: int = 100 + non_persistent_ckpt_type: Optional[str] = None + non_persistent_save_interval: Optional[int] = None + custom_manager_class: Optional[str] = None + + +@pytest.fixture(autouse=True) +def _clean_registration(): + reset_configuration() + yield + reset_configuration() + + +@pytest.fixture +def instantiate_utils(mocker): + """A stand-in for Megatron Bridge's target allowlist module.""" + module = types.ModuleType("megatron.bridge.utils.instantiate_utils") + module.prefixes = [] + module.register_allowed_target_prefix = module.prefixes.append + mocker.patch.dict(sys.modules, {"megatron.bridge.utils.instantiate_utils": module}) + return module + + +class TestRegisterWithMegatronBridge: + def test_registers_the_package_prefix(self, instantiate_utils): + # Given/When + adapter.register_with_megatron_bridge() + + # Then + assert_that(instantiate_utils.prefixes).contains("ml_flashpoint") + + def test_is_idempotent(self, instantiate_utils): + # Given/When + adapter.register_with_megatron_bridge() + adapter.register_with_megatron_bridge() + + # Then: registering twice is harmless; Bridge de-duplicates prefixes. + assert_that(set(instantiate_utils.prefixes)).is_equal_to({"ml_flashpoint"}) + + def test_missing_helper_is_tolerated(self, mocker): + # Given a Bridge build without the allowlist helper. + module = types.ModuleType("megatron.bridge.utils.instantiate_utils") + mocker.patch.dict(sys.modules, {"megatron.bridge.utils.instantiate_utils": module}) + + # When/Then: no exception, just a warning. + adapter.register_with_megatron_bridge() + + +class TestEnable: + def test_points_the_config_at_the_adapter(self, instantiate_utils): + # Given + config = FakeCheckpointConfig() + + # When + adapter.enable(config, non_persistent_save_interval=20) + + # Then + assert_that(config.custom_manager_class).is_equal_to(CUSTOM_MANAGER_CLASS) + assert_that(config.non_persistent_ckpt_type).is_equal_to("local") + assert_that(config.non_persistent_save_interval).is_equal_to(20) + + def test_leaves_the_durable_cadence_alone(self, instantiate_utils): + # Given + config = FakeCheckpointConfig(save="/durable/checkpoints", save_interval=100) + + # When + adapter.enable(config, non_persistent_save_interval=20) + + # Then + assert_that(config.save).is_equal_to("/durable/checkpoints") + assert_that(config.save_interval).is_equal_to(100) + + def test_registers_the_allowlist_prefix(self, instantiate_utils): + # Given + config = FakeCheckpointConfig() + + # When + adapter.enable(config, non_persistent_save_interval=5) + + # Then + assert_that(instantiate_utils.prefixes).contains("ml_flashpoint") + + def test_registers_the_supplied_ml_flashpoint_config(self, instantiate_utils, tmp_path): + # Given + mlf_config = MLFlashpointBridgeConfig(base_container=str(tmp_path / "mlf")) + + # When + adapter.enable(FakeCheckpointConfig(), non_persistent_save_interval=5, mlf_config=mlf_config) + + # Then + assert_that(adapter.get_config()).is_equal_to(mlf_config) + + @pytest.mark.parametrize("interval", [0, -1]) + def test_rejects_a_non_positive_interval(self, instantiate_utils, interval): + # Given/When/Then + with pytest.raises(ValueError, match="non_persistent_save_interval"): + adapter.enable(FakeCheckpointConfig(), non_persistent_save_interval=interval) + + def test_custom_manager_class_resolves_to_the_manager(self): + # Given + module_path, class_name = CUSTOM_MANAGER_CLASS.rsplit(".", 1) + + # When + module = __import__(module_path, fromlist=[class_name]) + + # Then + assert_that(getattr(module, class_name)).is_same_as(adapter.MLFlashpointBridgeCheckpointManager) diff --git a/tests/adapter/megatron_bridge/test_local_checkpoint_index.py b/tests/adapter/megatron_bridge/test_local_checkpoint_index.py new file mode 100644 index 0000000..fe59ba1 --- /dev/null +++ b/tests/adapter/megatron_bridge/test_local_checkpoint_index.py @@ -0,0 +1,230 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os + +import pytest +from assertpy import assert_that + +from ml_flashpoint.adapter.megatron_bridge import local_checkpoint_index as index_module +from ml_flashpoint.adapter.megatron_bridge.local_checkpoint_index import ( + NO_CHECKPOINT, + MLFlashpointLocalCheckpointIndex, +) +from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId + + +@pytest.fixture +def base_container(tmp_path) -> CheckpointContainerId: + return CheckpointContainerId(str(tmp_path / "mlf")) + + +@pytest.fixture(autouse=True) +def _single_rank(mocker): + """Runs the index as if this were the only, node-local-zero rank.""" + mocker.patch.object(index_module.dist, "is_initialized", return_value=True) + mocker.patch.object(index_module.dist, "get_node_local_rank", return_value=0) + + +def _make_container(base_container: CheckpointContainerId, step: int) -> CheckpointContainerId: + container = CheckpointContainerId.create_child(base_container, CheckpointContainerId.format_version_container(step)) + os.makedirs(str(container), exist_ok=True) + return container + + +class TestFindLatest: + def test_returns_sentinel_when_nothing_recoverable(self, base_container, mocker): + # Given + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = None + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + result = index.find_latest() + + # Then + assert_that(result).is_equal_to(NO_CHECKPOINT) + + def test_returns_step_of_latest_container(self, base_container, mocker): + # Given + container = _make_container(base_container, 120) + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + result = index.find_latest() + + # Then + assert_that(result).is_equal_to(120) + + def test_returns_sentinel_for_unparseable_container_name(self, base_container, mocker): + # Given + container = CheckpointContainerId(str(base_container) + "/not-a-version-dir") + os.makedirs(str(container), exist_ok=True) + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + result = index.find_latest() + + # Then + assert_that(result).is_equal_to(NO_CHECKPOINT) + + def test_discovery_runs_once(self, base_container, mocker): + # Given + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = _make_container(base_container, 5) + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + index.find_latest() + index.find_latest() + index.resolve_latest_container() + + # Then + assert_that(loader.get_latest_complete_checkpoint.call_count).is_equal_to(1) + + def test_discovery_failure_is_swallowed(self, base_container, mocker): + # Given + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.side_effect = RuntimeError("peer unreachable") + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + result = index.find_latest() + + # Then + assert_that(result).is_equal_to(NO_CHECKPOINT) + + def test_invalidate_forces_rediscovery(self, base_container, mocker): + # Given + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = None + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + index.find_latest() + + # When + index.invalidate() + index.find_latest() + + # Then + assert_that(loader.get_latest_complete_checkpoint.call_count).is_equal_to(2) + + def test_disable_reports_no_checkpoint(self, base_container, mocker): + # Given + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = _make_container(base_container, 7) + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + index.disable() + + # Then + assert_that(index.find_latest()).is_equal_to(NO_CHECKPOINT) + assert_that(index.resolve_latest_container()).is_none() + loader.get_latest_complete_checkpoint.assert_not_called() + + +class TestMetadataStub: + def test_writes_stub_for_discovered_container(self, base_container, mocker): + # Given + container = _make_container(base_container, 42) + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + index.resolve_latest_container() + + # Then + stub_path = os.path.join(str(container), "metadata.json") + assert_that(os.path.exists(stub_path)).is_true() + with open(stub_path) as handle: + assert_that(json.load(handle)).is_equal_to({"sharded_backend": ""}) + + def test_does_not_overwrite_existing_stub(self, base_container, mocker): + # Given + container = _make_container(base_container, 42) + stub_path = os.path.join(str(container), "metadata.json") + with open(stub_path, "w") as handle: + json.dump({"sharded_backend": "already-here"}, handle) + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + index.resolve_latest_container() + + # Then + with open(stub_path) as handle: + assert_that(json.load(handle)).is_equal_to({"sharded_backend": "already-here"}) + + def test_non_local_rank_zero_does_not_write(self, base_container, mocker): + # Given + mocker.patch.object(index_module.dist, "get_node_local_rank", return_value=1) + container = _make_container(base_container, 42) + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + index.resolve_latest_container() + + # Then + assert_that(os.path.exists(os.path.join(str(container), "metadata.json"))).is_false() + + def test_write_failure_does_not_propagate(self, base_container, mocker): + # Given a container path that was never created on disk. + container = CheckpointContainerId(str(base_container) + "/step-1_ckpt") + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + resolved = index.resolve_latest_container() + + # Then + assert_that(resolved).is_equal_to(container) + + +class TestProperties: + def test_local_ckpt_dir_is_base_container(self, base_container, mocker): + # Given + index = MLFlashpointLocalCheckpointIndex(base_container, mocker.MagicMock()) + + # When/Then + assert_that(index.local_ckpt_dir).is_equal_to(str(base_container)) + + def test_latest_container_is_none_before_discovery(self, base_container, mocker): + # Given + index = MLFlashpointLocalCheckpointIndex(base_container, mocker.MagicMock()) + + # When/Then + assert_that(index.latest_container).is_none() + + def test_latest_container_after_discovery(self, base_container, mocker): + # Given + container = _make_container(base_container, 3) + loader = mocker.MagicMock() + loader.get_latest_complete_checkpoint.return_value = container + index = MLFlashpointLocalCheckpointIndex(base_container, loader) + + # When + index.resolve_latest_container() + + # Then + assert_that(index.latest_container).is_equal_to(container) diff --git a/tests/adapter/megatron_bridge/test_runtime.py b/tests/adapter/megatron_bridge/test_runtime.py new file mode 100644 index 0000000..49cb927 --- /dev/null +++ b/tests/adapter/megatron_bridge/test_runtime.py @@ -0,0 +1,300 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from assertpy import assert_that + +from ml_flashpoint.adapter.megatron_bridge import runtime as runtime_module +from ml_flashpoint.adapter.megatron_bridge.config import MLFlashpointBridgeConfig +from ml_flashpoint.adapter.megatron_bridge.runtime import ( + NUM_OF_BUFFERS_PER_OBJECT, + MLFlashpointBridgeRuntime, + get_runtime, + shutdown_runtime, +) + + +@pytest.fixture +def config(tmp_path) -> MLFlashpointBridgeConfig: + return MLFlashpointBridgeConfig(base_container=str(tmp_path / "mlf"), write_thread_count=2) + + +@pytest.fixture +def distributed(mocker): + """Makes the runtime believe it is rank 0 of an initialized process group.""" + mocker.patch.object(runtime_module.dist, "is_available", return_value=True) + mocker.patch.object(runtime_module.dist, "is_initialized", return_value=True) + mocker.patch.object(runtime_module.dist, "get_rank", return_value=0) + mocker.patch.object(runtime_module.dist, "get_node_local_rank", return_value=0) + mocker.patch.object(runtime_module.dist, "get_world_size", return_value=1) + mocker.patch.object(runtime_module.dist, "barrier") + + +@pytest.fixture +def collaborators(mocker): + """Replaces everything the runtime constructs with doubles.""" + return { + "object_manager": mocker.patch.object(runtime_module, "CheckpointObjectManager"), + "replication_manager": mocker.patch.object(runtime_module, "ReplicationManager"), + "saver": mocker.patch.object(runtime_module, "DefaultMLFlashpointCheckpointSaver"), + "loader": mocker.patch.object(runtime_module, "DefaultMLFlashpointCheckpointLoader"), + "storage_writer": mocker.patch.object(runtime_module, "MemoryStorageWriter"), + "save_strategy": mocker.patch.object(runtime_module, "MLFlashpointMegatronAsyncSaveStrategy"), + "load_strategy": mocker.patch.object(runtime_module, "MLFlashpointMegatronLoadStrategy"), + "parallel_save": mocker.patch.object(runtime_module, "FullyParallelSaveStrategyWrapper"), + "parallel_load": mocker.patch.object(runtime_module, "FullyParallelLoadStrategyWrapper"), + "queue": mocker.patch.object(runtime_module, "AsyncCallsQueue"), + "mp": mocker.patch.object(runtime_module, "torch_mp"), + } + + +@pytest.fixture(autouse=True) +def _clean_module_runtime(): + runtime_module._RUNTIME = None + yield + runtime_module._RUNTIME = None + + +class TestConstruction: + def test_requires_an_initialized_process_group(self, config, mocker, collaborators): + # Given + mocker.patch.object(runtime_module.dist, "is_initialized", return_value=False) + + # When/Then + with pytest.raises(RuntimeError, match="torch.distributed"): + MLFlashpointBridgeRuntime(config) + + def test_sizes_the_buffer_pool_from_the_thread_count(self, config, distributed, collaborators, mocker): + # Given + pool_config = mocker.patch.object(runtime_module, "BufferPoolConfig") + + # When + MLFlashpointBridgeRuntime(config) + + # Then + assert_that(pool_config.call_args.kwargs["num_buffers"]).is_equal_to( + config.write_thread_count * NUM_OF_BUFFERS_PER_OBJECT + ) + + def test_initializes_replication(self, config, distributed, collaborators): + # Given/When + MLFlashpointBridgeRuntime(config) + + # Then + collaborators["replication_manager"].return_value.initialize.assert_called_once() + + def test_wraps_strategies_for_fully_parallel_saving(self, config, distributed, collaborators): + # Given/When + MLFlashpointBridgeRuntime(config) + + # Then + collaborators["parallel_save"].assert_called_once() + collaborators["parallel_load"].assert_called_once() + + def test_skips_the_parallel_wrapper_when_disabled(self, tmp_path, distributed, collaborators): + # Given + config = MLFlashpointBridgeConfig(base_container=str(tmp_path / "mlf"), use_fully_parallel_wrapper=False) + + # When + runtime = MLFlashpointBridgeRuntime(config) + + # Then + collaborators["parallel_save"].assert_not_called() + assert_that(runtime.save_strategy).is_same_as(collaborators["save_strategy"].return_value) + + def test_uses_a_persistent_async_queue(self, config, distributed, collaborators): + # Given/When + MLFlashpointBridgeRuntime(config) + + # Then + collaborators["queue"].assert_called_once_with(persistent=True) + + def test_exposes_the_base_container(self, config, distributed, collaborators): + # Given/When + runtime = MLFlashpointBridgeRuntime(config) + + # Then + assert_that(str(runtime.base_container)).is_equal_to(config.base_container) + assert_that(runtime.config).is_same_as(config) + + +class TestAsyncQueue: + def test_schedule_forwards_to_the_queue(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + queue = collaborators["queue"].return_value + queue.schedule_async_request.return_value = 7 + + # When + result = runtime.schedule("request") + + # Then + assert_that(result).is_equal_to(7) + queue.schedule_async_request.assert_called_once_with("request") + + def test_maybe_finalize_short_circuits_when_idle(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + queue = collaborators["queue"].return_value + queue.get_num_unfinalized_calls.return_value = 0 + + # When + result = runtime.maybe_finalize() + + # Then + assert_that(result).is_false() + queue.maybe_finalize_async_calls.assert_not_called() + + def test_maybe_finalize_reports_completions(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + queue = collaborators["queue"].return_value + queue.get_num_unfinalized_calls.return_value = 2 + queue.maybe_finalize_async_calls.return_value = [0, 1] + + # When + result = runtime.maybe_finalize(blocking=True) + + # Then + assert_that(result).is_true() + queue.maybe_finalize_async_calls.assert_called_once_with(True) + + def test_num_unfinalized_calls_is_forwarded(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + collaborators["queue"].return_value.get_num_unfinalized_calls.return_value = 3 + + # When/Then + assert_that(runtime.num_unfinalized_calls()).is_equal_to(3) + + +class TestShutdown: + def test_releases_replication_buffers_and_queue(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + + # When + runtime.shutdown() + + # Then + collaborators["replication_manager"].return_value.shutdown.assert_called_once() + collaborators["object_manager"].return_value.delete_container.assert_called_once() + collaborators["queue"].return_value.close.assert_called_once() + + def test_can_keep_the_container(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + + # When + runtime.shutdown(remove_checkpoints=False) + + # Then + collaborators["object_manager"].return_value.delete_container.assert_not_called() + + def test_is_idempotent(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + + # When + runtime.shutdown() + runtime.shutdown() + + # Then + assert_that(collaborators["queue"].return_value.close.call_count).is_equal_to(1) + + def test_replication_failure_does_not_stop_teardown(self, config, distributed, collaborators): + # Given + collaborators["replication_manager"].return_value.shutdown.side_effect = RuntimeError("socket closed") + runtime = MLFlashpointBridgeRuntime(config) + + # When + runtime.shutdown() + + # Then + collaborators["queue"].return_value.close.assert_called_once() + + def test_container_deletion_failure_does_not_stop_teardown(self, config, distributed, collaborators): + # Given + collaborators["object_manager"].return_value.delete_container.side_effect = OSError("busy") + runtime = MLFlashpointBridgeRuntime(config) + + # When + runtime.shutdown() + + # Then + collaborators["queue"].return_value.close.assert_called_once() + + def test_finalize_after_shutdown_is_a_noop(self, config, distributed, collaborators): + # Given + runtime = MLFlashpointBridgeRuntime(config) + runtime.shutdown() + collaborators["queue"].return_value.get_num_unfinalized_calls.return_value = 5 + + # When + result = runtime.maybe_finalize() + + # Then + assert_that(result).is_false() + + def test_only_local_rank_zero_deletes_the_container(self, config, distributed, collaborators, mocker): + # Given + mocker.patch.object(runtime_module.dist, "get_node_local_rank", return_value=1) + runtime = MLFlashpointBridgeRuntime(config) + + # When + runtime.shutdown() + + # Then + collaborators["object_manager"].return_value.delete_container.assert_not_called() + + +class TestModuleRuntime: + def test_get_runtime_builds_once(self, config, distributed, collaborators, mocker): + # Given + build = mocker.spy(runtime_module, "MLFlashpointBridgeRuntime") + + # When + first = get_runtime(config) + second = get_runtime(config) + + # Then + assert_that(second).is_same_as(first) + assert_that(build.call_count).is_equal_to(1) + + def test_shutdown_runtime_clears_the_module_state(self, config, distributed, collaborators): + # Given + get_runtime(config) + + # When + shutdown_runtime() + + # Then + assert_that(runtime_module._RUNTIME).is_none() + + def test_shutdown_runtime_without_a_runtime_is_a_noop(self): + # Given/When + shutdown_runtime() + + # Then + assert_that(runtime_module._RUNTIME).is_none() + + def test_shutdown_runtime_clears_state_even_on_failure(self, config, distributed, collaborators, mocker): + # Given + runtime = get_runtime(config) + mocker.patch.object(runtime, "shutdown", side_effect=RuntimeError("teardown failed")) + + # When/Then + with pytest.raises(RuntimeError): + shutdown_runtime() + assert_that(runtime_module._RUNTIME).is_none() diff --git a/tests/adapter/nemo_rl/test_checkpointer.py b/tests/adapter/nemo_rl/test_checkpointer.py new file mode 100644 index 0000000..093b3d4 --- /dev/null +++ b/tests/adapter/nemo_rl/test_checkpointer.py @@ -0,0 +1,201 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from assertpy import assert_that + +from ml_flashpoint.adapter.nemo_rl import checkpointer as checkpointer_module +from ml_flashpoint.adapter.nemo_rl.checkpointer import ( + MLFlashpointNeMoRLCheckpointer, + _as_module_list, +) + + +@pytest.fixture +def manager(mocker): + """Replaces the Bridge manager so no ML Flashpoint runtime is built.""" + instance = mocker.MagicMock() + mocker.patch.object(checkpointer_module, "MLFlashpointBridgeCheckpointManager", return_value=instance) + return instance + + +@pytest.fixture +def checkpoint_config(): + return type("CheckpointConfig", (), {"ckpt_format": "torch_dist"})() + + +@pytest.fixture +def state(): + train_state = type("TrainState", (), {"step": 12, "floating_point_operations_so_far": 555})() + return type("GlobalState", (), {"train_state": train_state})() + + +@pytest.fixture +def checkpointer(checkpoint_config, manager) -> MLFlashpointNeMoRLCheckpointer: + return MLFlashpointNeMoRLCheckpointer(checkpoint_config) + + +class TestSave: + def test_marks_the_checkpoint_non_persistent(self, checkpointer, manager, state): + # Given + model = object() + + # When + checkpointer.save(state=state, model=model) + + # Then + ctx = manager.save.call_args.args[0] + assert_that(ctx.non_persistent_ckpt).is_true() + assert_that(ctx.state).is_same_as(state) + assert_that(ctx.model).is_equal_to([model]) + + def test_defaults_flops_from_the_train_state(self, checkpointer, manager, state): + # Given/When + checkpointer.save(state=state, model=object()) + + # Then + assert_that(manager.save.call_args.args[0].num_floating_point_operations_so_far).is_equal_to(555) + + def test_explicit_flops_win(self, checkpointer, manager, state): + # Given/When + checkpointer.save(state=state, model=object(), num_floating_point_operations_so_far=1) + + # Then + assert_that(manager.save.call_args.args[0].num_floating_point_operations_so_far).is_equal_to(1) + + def test_passes_optimizer_and_scheduler_through(self, checkpointer, manager, state, mocker): + # Given + optimizer = mocker.MagicMock() + scheduler = mocker.MagicMock() + + # When + checkpointer.save(state=state, model=object(), optimizer=optimizer, opt_param_scheduler=scheduler) + + # Then + ctx = manager.save.call_args.args[0] + assert_that(ctx.optimizer).is_same_as(optimizer) + assert_that(ctx.opt_param_scheduler).is_same_as(scheduler) + + def test_no_callback_manager_is_passed(self, checkpointer, manager, state): + # Given/When + checkpointer.save(state=state, model=object()) + + # Then + assert_that(manager.save.call_args.kwargs["callback_manager"]).is_none() + + +class TestLoad: + def test_returns_none_when_nothing_recoverable(self, checkpointer, manager, state): + # Given + manager._find_ml_flashpoint_checkpoint.return_value = None + + # When + result = checkpointer.load(state=state, model=object()) + + # Then + assert_that(result).is_none() + manager._load_ml_flashpoint.assert_not_called() + + def test_loads_the_discovered_container(self, checkpointer, manager, state): + # Given + manager._find_ml_flashpoint_checkpoint.return_value = "container" + manager._load_ml_flashpoint.return_value = (12, 555) + + # When + result = checkpointer.load(state=state, model=object()) + + # Then + assert_that(result).is_equal_to((12, 555)) + assert_that(manager._load_ml_flashpoint.call_args.args[1]).is_equal_to("container") + + def test_never_falls_back_to_the_durable_path(self, checkpointer, manager, state): + # Given: NeMo RL owns the durable resume decision. + manager._find_ml_flashpoint_checkpoint.return_value = "container" + manager._load_ml_flashpoint.return_value = None + + # When + result = checkpointer.load(state=state, model=object()) + + # Then + assert_that(result).is_none() + manager.load.assert_not_called() + + def test_forwards_strictness(self, checkpointer, manager, state): + # Given + manager._find_ml_flashpoint_checkpoint.return_value = "container" + + # When + checkpointer.load(state=state, model=object(), strict=False) + + # Then + assert_that(manager._load_ml_flashpoint.call_args.args[0].strict).is_false() + + +class TestLifecycle: + def test_enabled_follows_the_manager(self, checkpointer, manager): + # Given + manager.enabled = False + + # When/Then + assert_that(checkpointer.enabled).is_false() + + def test_finalize_without_a_runtime_is_a_noop(self, checkpointer, manager): + # Given + manager.runtime = None + + # When + checkpointer.finalize() + + # Then no exception escapes. + + def test_finalize_drains_the_runtime(self, checkpointer, manager, mocker): + # Given + runtime = mocker.MagicMock() + manager.runtime = runtime + + # When + checkpointer.finalize(blocking=True) + + # Then + runtime.maybe_finalize.assert_called_once_with(blocking=True) + + def test_shutdown_delegates_to_the_manager(self, checkpointer, manager): + # Given/When + checkpointer.shutdown() + + # Then + manager.shutdown.assert_called_once() + + +class TestAsModuleList: + def test_wraps_a_single_module(self): + # Given + module = object() + + # When/Then + assert_that(_as_module_list(module)).is_equal_to([module]) + + def test_passes_a_list_through(self): + # Given + modules = [object(), object()] + + # When/Then + assert_that(_as_module_list(modules)).is_equal_to(modules) + + def test_normalizes_a_tuple(self): + # Given + modules = (object(),) + + # When/Then + assert_that(_as_module_list(modules)).is_equal_to(list(modules)) diff --git a/tests/adapter/nemo_rl/test_integration.py b/tests/adapter/nemo_rl/test_integration.py new file mode 100644 index 0000000..262221f --- /dev/null +++ b/tests/adapter/nemo_rl/test_integration.py @@ -0,0 +1,273 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from assertpy import assert_that + +from ml_flashpoint.adapter.nemo_rl import integration as integration_module +from ml_flashpoint.adapter.nemo_rl.integration import ( + MODE_AUGMENT, + MODE_REPLACE, + get_checkpointer, + install_from_env, + install_into_worker, + uninstall_from_worker, +) + + +class FakeWorker: + """Stands in for NeMo RL's ``MegatronPolicyWorker``.""" + + def __init__(self, checkpoint_config): + self.mcore_state = type("State", (), {"cfg": type("Cfg", (), {"checkpoint": checkpoint_config})()})() + self.model = object() + self.optimizer = object() + self.scheduler = object() + self.durable_saves: list[tuple[str, str]] = [] + + def save_checkpoint(self, weights_path, optimizer_path=None, **kwargs): + self.durable_saves.append((weights_path, optimizer_path)) + return "durable-result" + + +@pytest.fixture +def checkpoint_config(): + return type("CheckpointConfig", (), {"ckpt_format": "torch_dist", "non_persistent_ckpt_type": "local"})() + + +@pytest.fixture +def worker(checkpoint_config) -> FakeWorker: + return FakeWorker(checkpoint_config) + + +@pytest.fixture +def checkpointer(mocker): + """Replaces the real checkpointer so no ML Flashpoint runtime is built.""" + instance = mocker.MagicMock() + mocker.patch.object(integration_module, "MLFlashpointNeMoRLCheckpointer", return_value=instance) + return instance + + +class TestInstall: + def test_returns_checkpointer_and_records_it(self, worker, checkpointer): + # Given/When + result = install_into_worker(worker) + + # Then + assert_that(result).is_same_as(checkpointer) + assert_that(get_checkpointer(worker)).is_same_as(checkpointer) + + def test_installing_twice_is_a_noop(self, worker, checkpointer): + # Given + first = install_into_worker(worker) + + # When + second = install_into_worker(worker) + + # Then + assert_that(second).is_same_as(first) + + @pytest.mark.parametrize("mode", ["", "both", "AUGMENT"]) + def test_rejects_unknown_mode(self, worker, checkpointer, mode): + # Given/When/Then + with pytest.raises(ValueError, match="mode must be one of"): + install_into_worker(worker, mode=mode) + + @pytest.mark.parametrize("interval", [0, -1]) + def test_rejects_non_positive_durable_interval(self, worker, checkpointer, interval): + # Given/When/Then + with pytest.raises(ValueError, match="durable_every_n_saves"): + install_into_worker(worker, mode=MODE_REPLACE, durable_every_n_saves=interval) + + def test_rejects_object_that_is_not_a_megatron_worker(self, checkpointer): + # Given + not_a_worker = object() + + # When/Then + with pytest.raises(AttributeError, match="mcore_state"): + install_into_worker(not_a_worker) + + +class TestAugmentMode: + def test_every_save_stays_durable(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_AUGMENT) + + # When + for step in range(3): + worker.save_checkpoint(f"/durable/step_{step}", optimizer_path=f"/durable/step_{step}/optim") + + # Then + assert_that(worker.durable_saves).is_length(3) + assert_that(checkpointer.save.call_count).is_equal_to(3) + + def test_returns_the_wrapped_result(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_AUGMENT) + + # When + result = worker.save_checkpoint("/durable/step_0") + + # Then + assert_that(result).is_equal_to("durable-result") + + def test_optimizer_is_skipped_when_no_optimizer_path(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_AUGMENT) + + # When + worker.save_checkpoint("/durable/step_0", optimizer_path=None) + + # Then + assert_that(checkpointer.save.call_args.kwargs["optimizer"]).is_none() + assert_that(checkpointer.save.call_args.kwargs["opt_param_scheduler"]).is_none() + + def test_optimizer_is_included_when_optimizer_path_given(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_AUGMENT) + + # When + worker.save_checkpoint("/durable/step_0", optimizer_path="/durable/step_0/optim") + + # Then + assert_that(checkpointer.save.call_args.kwargs["optimizer"]).is_same_as(worker.optimizer) + assert_that(checkpointer.save.call_args.kwargs["opt_param_scheduler"]).is_same_as(worker.scheduler) + + def test_ml_flashpoint_failure_does_not_break_the_durable_save(self, worker, checkpointer): + # Given + checkpointer.save.side_effect = RuntimeError("buffer pool exhausted") + install_into_worker(worker, mode=MODE_AUGMENT) + + # When + result = worker.save_checkpoint("/durable/step_0") + + # Then + assert_that(result).is_equal_to("durable-result") + assert_that(worker.durable_saves).is_length(1) + + +class TestReplaceMode: + def test_keeps_every_nth_durable_save(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_REPLACE, durable_every_n_saves=3) + + # When + for step in range(6): + worker.save_checkpoint(f"/durable/step_{step}") + + # Then + assert_that(checkpointer.save.call_count).is_equal_to(6) + assert_that([path for path, _ in worker.durable_saves]).is_equal_to(["/durable/step_2", "/durable/step_5"]) + + def test_skipped_save_returns_none(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_REPLACE, durable_every_n_saves=2) + + # When + first = worker.save_checkpoint("/durable/step_0") + second = worker.save_checkpoint("/durable/step_1") + + # Then + assert_that(first).is_none() + assert_that(second).is_equal_to("durable-result") + + def test_interval_of_one_matches_augment(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_REPLACE, durable_every_n_saves=1) + + # When + worker.save_checkpoint("/durable/step_0") + worker.save_checkpoint("/durable/step_1") + + # Then + assert_that(worker.durable_saves).is_length(2) + + +class TestUninstall: + def test_restores_the_original_method(self, worker, checkpointer): + # Given + install_into_worker(worker, mode=MODE_REPLACE, durable_every_n_saves=100) + + # When + uninstall_from_worker(worker) + worker.save_checkpoint("/durable/step_0") + + # Then + assert_that(worker.durable_saves).is_length(1) + assert_that(get_checkpointer(worker)).is_none() + + def test_shuts_the_checkpointer_down(self, worker, checkpointer): + # Given + install_into_worker(worker) + + # When + uninstall_from_worker(worker) + + # Then + checkpointer.shutdown.assert_called_once() + + def test_uninstalling_a_clean_worker_is_a_noop(self, worker, checkpointer): + # Given/When + uninstall_from_worker(worker) + + # Then + assert_that(get_checkpointer(worker)).is_none() + + +class TestInstallFromEnv: + def test_disabled_by_default(self, worker, checkpointer, monkeypatch): + # Given + monkeypatch.delenv("MLFLASHPOINT_NEMO_RL_ENABLED", raising=False) + + # When + result = install_from_env(worker) + + # Then + assert_that(result).is_none() + assert_that(get_checkpointer(worker)).is_none() + + def test_installs_when_enabled(self, worker, checkpointer, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_NEMO_RL_ENABLED", "true") + + # When + result = install_from_env(worker) + + # Then + assert_that(result).is_same_as(checkpointer) + + def test_reads_mode_and_interval(self, worker, checkpointer, monkeypatch, mocker): + # Given + monkeypatch.setenv("MLFLASHPOINT_NEMO_RL_ENABLED", "true") + monkeypatch.setenv("MLFLASHPOINT_NEMO_RL_MODE", MODE_REPLACE) + monkeypatch.setenv("MLFLASHPOINT_NEMO_RL_DURABLE_EVERY_N_SAVES", "4") + install = mocker.spy(integration_module, "install_into_worker") + + # When + install_from_env(worker) + + # Then + assert_that(install.call_args.kwargs["mode"]).is_equal_to(MODE_REPLACE) + assert_that(install.call_args.kwargs["durable_every_n_saves"]).is_equal_to(4) + + def test_respects_the_global_disable_switch(self, worker, checkpointer, monkeypatch): + # Given + monkeypatch.setenv("MLFLASHPOINT_NEMO_RL_ENABLED", "true") + monkeypatch.setenv("MLFLASHPOINT_BRIDGE_ENABLED", "false") + + # When + result = install_from_env(worker) + + # Then + assert_that(result).is_none() diff --git a/tests/scripts/test_checkpoint_timing_scripts.py b/tests/scripts/test_checkpoint_timing_scripts.py new file mode 100644 index 0000000..cde81fc --- /dev/null +++ b/tests/scripts/test_checkpoint_timing_scripts.py @@ -0,0 +1,296 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import json +from pathlib import Path + +import pytest +from assertpy import assert_that + +_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "scripts" / "benchmarks" + + +def _load(name: str): + """Imports a benchmark script by path, since scripts/ is not a package.""" + spec = importlib.util.spec_from_file_location(name, _SCRIPTS_DIR / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def parser_module(): + return _load("parse_checkpoint_timings") + + +@pytest.fixture(scope="module") +def compare_module(): + return _load("compare_checkpoint_timings") + + +BASELINE_LOG = """ + iteration 10/100 | elapsed time per iteration (ms): 812.3 + save-checkpoint ................................: (18450.20, 18512.90) + iteration 20/100 + save-checkpoint ................................: (17980.00, 18100.10) +""" + +FLASHPOINT_LOG = """ + save-checkpoint-non-persistent .................: (612.40, 640.10) +MLFlashpointBridgeCheckpointManager.save took 0.5981s + save-checkpoint ................................: (18300.00, 18402.10) + load-checkpoint ................................: (2100.00, 2150.00) +nemo_rl.save_checkpoint took 12.3456s +""" + + +class TestParseLines: + def test_reads_the_slowest_rank_in_seconds(self, parser_module): + # Given/When + samples = parser_module.parse_lines(BASELINE_LOG.splitlines()) + + # Then + assert_that(samples["save-checkpoint"]).is_length(2) + assert_that(samples["save-checkpoint"][0]).is_close_to(18.5129, 1e-6) + assert_that(samples["save-checkpoint"][1]).is_close_to(18.1001, 1e-6) + + def test_reads_the_max_only_log_option(self, parser_module): + """A run configured with timing_log_option='max' logs a bare value, not a (min, max) pair.""" + # Given + lines = [" save-checkpoint ................................: 2345.67"] + + # When + samples = parser_module.parse_lines(lines) + + # Then + assert_that(samples["save-checkpoint"]).is_length(1) + assert_that(samples["save-checkpoint"][0]).is_close_to(2.34567, 1e-6) + + def test_both_log_options_produce_comparable_samples(self, parser_module): + """The same underlying duration parses identically whichever shape it was logged in.""" + # Given + minmax = [" save-checkpoint ....: (100.00, 2345.67)"] + max_only = [" save-checkpoint ....: 2345.67"] + + # When + expected_samples = parser_module.parse_lines(minmax) + actual_samples = parser_module.parse_lines(max_only) + + # Then + assert_that(actual_samples).is_equal_to(expected_samples) + + def test_recognizes_non_persistent_and_load_timers(self, parser_module): + # Given/When + samples = parser_module.parse_lines(FLASHPOINT_LOG.splitlines()) + + # Then + assert_that(samples).contains_key("save-checkpoint-non-persistent", "load-checkpoint") + + def test_recognizes_ml_flashpoint_timers(self, parser_module): + # Given/When + samples = parser_module.parse_lines(FLASHPOINT_LOG.splitlines()) + + # Then + assert_that(samples["MLFlashpointBridgeCheckpointManager.save"]).is_equal_to([0.5981]) + assert_that(samples["nemo_rl.save_checkpoint"]).is_equal_to([12.3456]) + + def test_ignores_unrelated_timing_lines(self, parser_module): + # Given + lines = ["some_unrelated_function took 3.5s", "forward-backward ...: (10.0, 11.0)"] + + # When + samples = parser_module.parse_lines(lines) + + # Then + assert_that(samples).is_empty() + + def test_empty_input_yields_no_samples(self, parser_module): + # Given/When/Then + assert_that(parser_module.parse_lines([])).is_empty() + + +class TestSummarize: + def test_reports_the_expected_statistics(self, parser_module): + # Given + samples = {"save-checkpoint": [1.0, 3.0, 2.0]} + + # When + summary = parser_module.summarize(samples)["save-checkpoint"] + + # Then + assert_that(summary["count"]).is_equal_to(3) + assert_that(summary["mean_s"]).is_equal_to(2.0) + assert_that(summary["median_s"]).is_equal_to(2.0) + assert_that(summary["min_s"]).is_equal_to(1.0) + assert_that(summary["max_s"]).is_equal_to(3.0) + + def test_single_sample_has_zero_stdev(self, parser_module): + # Given/When + summary = parser_module.summarize({"save-checkpoint": [4.0]})["save-checkpoint"] + + # Then + assert_that(summary["stdev_s"]).is_equal_to(0.0) + + +class TestParserCli: + def test_writes_a_labelled_report(self, parser_module, tmp_path): + # Given + log = tmp_path / "run.log" + log.write_text(BASELINE_LOG) + out = tmp_path / "report.json" + + # When + exit_code = parser_module.main(["--label", "baseline", str(log), "--output", str(out)]) + + # Then + assert_that(exit_code).is_equal_to(0) + report = json.loads(out.read_text()) + assert_that(report["label"]).is_equal_to("baseline") + assert_that(report["timers"]).contains_key("save-checkpoint") + + def test_reports_nothing_found_without_failing(self, parser_module, tmp_path, capsys): + # Given + log = tmp_path / "empty.log" + log.write_text("nothing interesting here\n") + out = tmp_path / "report.json" + + # When + exit_code = parser_module.main(["--label", "baseline", str(log), "--output", str(out)]) + + # Then + assert_that(exit_code).is_equal_to(0) + assert_that(capsys.readouterr().err).contains("No checkpoint timings found") + + +class TestCompare: + def _report(self, label, timers): + return {"label": label, "timers": timers, "raw_samples": {}} + + def test_orders_headline_timers_first(self, compare_module): + # Given + stats = {"count": 1, "mean_s": 1.0, "max_s": 1.0} + baseline = self._report("baseline", {"zzz-other": stats, "save-checkpoint": stats}) + candidate = self._report("flashpoint", {"zzz-other": stats, "save-checkpoint": stats}) + + # When + rows = compare_module.compare(baseline, candidate) + + # Then + assert_that(rows[0]["timer"]).is_equal_to("save-checkpoint") + + def test_computes_the_speedup(self, compare_module): + # Given + baseline = self._report("baseline", {"save-checkpoint": {"count": 2, "mean_s": 18.0, "max_s": 19.0}}) + candidate = self._report("flashpoint", {"save-checkpoint": {"count": 2, "mean_s": 0.6, "max_s": 0.7}}) + + # When + rows = compare_module.compare(baseline, candidate) + + # Then + assert_that(rows[0]["mean_delta"]).contains("30.00x") + assert_that(rows[0]["mean_delta"]).contains("-17.400s") + + def test_missing_timer_on_one_side_is_not_an_error(self, compare_module): + # Given + baseline = self._report("baseline", {}) + candidate = self._report( + "flashpoint", {"save-checkpoint-non-persistent": {"count": 1, "mean_s": 0.6, "max_s": 0.6}} + ) + + # When + rows = compare_module.compare(baseline, candidate) + + # Then + assert_that(rows[0]["mean_delta"]).is_equal_to("n/a") + assert_that(rows[0]["baseline_count"]).is_equal_to(0) + + def test_zero_baseline_is_reported_rather_than_dividing(self, compare_module): + # Given + baseline = self._report("baseline", {"save-checkpoint": {"count": 1, "mean_s": 0.0, "max_s": 0.0}}) + candidate = self._report("flashpoint", {"save-checkpoint": {"count": 1, "mean_s": 1.0, "max_s": 1.0}}) + + # When + rows = compare_module.compare(baseline, candidate) + + # Then + assert_that(rows[0]["mean_delta"]).contains("baseline is 0") + + def test_render_includes_both_labels(self, compare_module): + # Given + rows = [ + { + "timer": "save-checkpoint", + "baseline_mean_s": 18.0, + "candidate_mean_s": 0.6, + "baseline_count": 2, + "candidate_count": 2, + "mean_delta": "-17.400s", + } + ] + + # When + rendered = compare_module.render("baseline", "flashpoint", rows) + + # Then + assert_that(rendered).contains("flashpoint vs baseline") + assert_that(rendered).contains("save-checkpoint") + + +class TestCompareCli: + def _write(self, path, label, timers): + path.write_text(json.dumps({"label": label, "timers": timers, "raw_samples": {}})) + + def test_prints_a_table(self, compare_module, tmp_path, capsys): + # Given + stats = {"count": 1, "mean_s": 2.0, "max_s": 2.0} + base = tmp_path / "base.json" + cand = tmp_path / "cand.json" + self._write(base, "baseline", {"save-checkpoint": stats}) + self._write(cand, "flashpoint", {"save-checkpoint": stats}) + + # When + exit_code = compare_module.main(["--baseline", str(base), "--candidate", str(cand)]) + + # Then + assert_that(exit_code).is_equal_to(0) + assert_that(capsys.readouterr().out).contains("save-checkpoint") + + def test_emits_json_on_request(self, compare_module, tmp_path, capsys): + # Given + stats = {"count": 1, "mean_s": 2.0, "max_s": 2.0} + base = tmp_path / "base.json" + cand = tmp_path / "cand.json" + self._write(base, "baseline", {"save-checkpoint": stats}) + self._write(cand, "flashpoint", {"save-checkpoint": stats}) + + # When + compare_module.main(["--baseline", str(base), "--candidate", str(cand), "--json"]) + + # Then + payload = json.loads(capsys.readouterr().out) + assert_that(payload["candidate_label"]).is_equal_to("flashpoint") + + def test_empty_reports_exit_non_zero(self, compare_module, tmp_path): + # Given + base = tmp_path / "base.json" + cand = tmp_path / "cand.json" + self._write(base, "baseline", {}) + self._write(cand, "flashpoint", {}) + + # When + exit_code = compare_module.main(["--baseline", str(base), "--candidate", str(cand)]) + + # Then + assert_that(exit_code).is_equal_to(1)