ci: re-enable GLM-4.7-Flash ckpt save/load e2e test - #1592
Open
guapisolo wants to merge 6 commits into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request re-enables the GLM-4.7 flash checkpoint end-to-end test in the CUDA CI suite by removing the disabled parameter and its associated FIXME comment. There are no review comments, and I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
guapisolo
force-pushed
the
ci/enable-glm47-flash-ckpt
branch
from
July 12, 2026 22:50
65bf379 to
03f2dd6
Compare
Disabled in #1149 as "Disabled due to bugs" with a FIXME pointing at Megatron. Since then the Megatron pin picked up the dist-ckpt load fix for bare BytesIO _extra_state (radixark/Megatron-LM#53, fixes #1293), and the sibling qwen3_4B ckpt test disabled in the same sweep proved stale on re-run (#1271). The training config itself (TP2/PP2/CP2/EP4 + MTP + EAGLE) already runs green in test_glm47_flash/test_r3_mtp; the save -> load -> async_save -> load roundtrip is the only untested part. Re-enable and let the run-ci-ckpt CI run on 8xH100 validate it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The re-enabled test fails deterministically at first load:
dist_checkpointing merge raises "Cannot merge two lists with different
lengths (736 and 768)" at ('optimizer', 1, 'param_state', ...). Under
dp_reshardable the expert optimizer's param_state is a positional
per-param list whose length differs across PP ranks (PP0 has 23 MoE
layers = 736 entries, PP1 has 23 + MTP = 768), and common.pt is written
from rank 0's view, so the other-shaped rank can never merge on load.
--dist-ckpt-optim-fully-reshardable stores optimizer state model-space
(name-keyed), which is PP-heterogeneity-safe; load follows the ckpt
metadata automatically. Optimizer DP=1 in this config, so the
save-side gather costs nothing.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Round 2 moved the failure from load to save: fully_reshardable's default NCCL variant all-gathers optimizer state into GPU buffers and died allocating 12.94 GiB with 8.47 GiB free on the 80 GB H100 (the state normally lives CPU-side under --optimizer-cpu-offload). --distrib-optim-fully-reshardable-mem-efficient switches the gather to Gloo, keeps it in host memory, and returns it only on DP rank 0 -- matching this test's CPU-offloaded optimizer; the DP=1 layout makes the lost save parallelism irrelevant. Requires gloo process groups, which are on by default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Run the checkpoint roundtrip in the 8-GPU H200 suite and enable its tight-host-memory configuration by default for the target CI environment.
guapisolo
force-pushed
the
ci/enable-glm47-flash-ckpt
branch
from
July 14, 2026 11:01
b08d4f0 to
660421a
Compare
The GLM checkpoint CI allocates excessive FP32 gather buffers in fully-reshardable mode. Use dp_reshardable and disable the ineffective tight-host-memory override by default.
HybridDeviceOptimizer forces master parameters and Adam states to FP32 while accepting lower-precision dtype flags. Remove those flags from shared E2E recipes and fail validation instead of silently ignoring them.
guapisolo
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
July 14, 2026 20:40
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
1 participant
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.
Summary
Re-enable GLM-4.7-Flash checkpoint CI without unsupported optimizer-memory settings.
Symptom & Reproduction
tests/e2e/ckpt/test_glm47_flash_ckpt.pywas disabled, leaving the save → load → async-save → load roundtrip without CI coverage.python tests/ci/run_suite.py --hw cuda --suite stage-c-8-gpu-h200 --labels run-ci-ckpt --list-onlyomitted the disabled test before this PR.fully_reshardabletriggered host OOM through 103–116 GiB of additional FP32 buffers per actor.Root Cause
register_cuda_cikept the checkpoint roundtrip disabled after its original Megatron blocker was fixed.dp_reshardablekeysparam_statepositionally, so uneven PP + MTP produced unmergeable optimizer lists.fully_reshardableuses name-keyed state but gathers every optimizer shard through FP32 buffers.HybridDeviceOptimizerignores the three low-precision state dtypes, soTIGHT_HOST_MEMORYcannot shrink those buffers.Fix
Re-enable the test in
stage-c-8-gpu-h200and remove thefully_reshardablesave path, eliminating the profiled host-OOM allocation. Remove the ineffective FP16 optimizer-state overrides from CPU-offload E2E recipes and reject any non-FP32 optimizer-state dtype after Megatron argument normalization.This fixes the observed host OOM only. The branch returns to
dp_reshardable; the uneven-PP optimizer-list load mismatch remains unverified until the H200 checkpoint roundtrip completes.Verification
test_glm47_flash_ckpt.pyis registered instage-c-8-gpu-h200underrun-ci-ckpt; the remote checkpoint job has not completed yet.test_optimizer_cpu_offload_rejects_lower_precision_state_dtypespasses and verifies the fail-fast boundary.torch.float16, then Miles rejected the CPU-offload combination.--optimizer-cpu-offloadwithout the three low-precision dtype flags.Review Focus
miles/backends/megatron_utils/arguments.py::validate_argsfor the CPU-offload dtype boundary.dp_reshardableload behavior.