Skip to content

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

Draft
g-husam wants to merge 1 commit into
mainfrom
feature/megatron-bridge-adapter
Draft

feat(adapter): add Megatron Bridge and NeMo RL checkpoint adapters#110
g-husam wants to merge 1 commit into
mainfrom
feature/megatron-bridge-adapter

Conversation

@g-husam

@g-husam g-husam commented Sep 7, 2026

Copy link
Copy Markdown
Collaborator

Megatron Bridge exposes a CheckpointManager protocol that a run can swap in via CheckpointConfig.custom_manager_class. MLFlashpointBridgeCheckpointManager implements it: non-persistent checkpoints (non_persistent_save_interval) are written by ML Flashpoint into node-local memory, while durable checkpoints (save_interval) keep going through Megatron Bridge untouched. Bridge takes the non-persistent branch only on steps that are not also durable-checkpoint steps, so the two cadences never collide.

Resume works by publishing an MLFlashpointLocalCheckpointIndex under checkpointing_context["local_checkpoint_manager"], which is what Bridge's setup._should_load_checkpoint consults. Load prefers the newest recoverable ML Flashpoint container and falls back to Bridge's own path when there is none or when the in-memory read fails.

NeMo RL needs a different entry point: it builds its Megatron state with Megatron Bridge but calls Bridge's functional save_checkpoint directly from MegatronPolicyWorker, so custom_manager_class is never consulted anywhere in that repo. install_into_worker wraps that one method instead, in either "augment" mode (durable writes unchanged, ML Flashpoint alongside) or "replace" mode (only every Nth checkpoint stays durable).

Also adds a checkpoint-timing benchmark harness that parses Megatron's save-checkpoint timers and ML Flashpoint's own execution timings out of training logs and diffs two runs, plus a runbook for measuring the difference on a GKE training cluster, and source-verified research notes on all three codebases for future work.

Megatron Bridge exposes a CheckpointManager protocol that a run can swap in
via CheckpointConfig.custom_manager_class. MLFlashpointBridgeCheckpointManager
implements it: non-persistent checkpoints (non_persistent_save_interval) are
written by ML Flashpoint into node-local memory, while durable checkpoints
(save_interval) keep going through Megatron Bridge untouched. Bridge takes the
non-persistent branch only on steps that are not also durable-checkpoint steps,
so the two cadences never collide.

Resume works by publishing an MLFlashpointLocalCheckpointIndex under
checkpointing_context["local_checkpoint_manager"], which is what Bridge's
setup._should_load_checkpoint consults. Load prefers the newest recoverable
ML Flashpoint container and falls back to Bridge's own path when there is none
or when the in-memory read fails.

NeMo RL needs a different entry point: it builds its Megatron state with
Megatron Bridge but calls Bridge's functional save_checkpoint directly from
MegatronPolicyWorker, so custom_manager_class is never consulted anywhere in
that repo. install_into_worker wraps that one method instead, in either
"augment" mode (durable writes unchanged, ML Flashpoint alongside) or
"replace" mode (only every Nth checkpoint stays durable).

Also adds a checkpoint-timing benchmark harness that parses Megatron's
save-checkpoint timers and ML Flashpoint's own execution timings out of
training logs and diffs two runs, plus a runbook for measuring the difference
on a GKE training cluster, and source-verified research notes on all three
codebases for future work.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant