feat(adapter): add Megatron Bridge and NeMo RL checkpoint adapters - #110
Draft
g-husam wants to merge 1 commit into
Draft
feat(adapter): add Megatron Bridge and NeMo RL checkpoint adapters#110g-husam wants to merge 1 commit into
g-husam wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.