Skip to content

[sync][training] fix: honor checkpoint sharding validation policy - #6434

Draft
yaoyu-33 wants to merge 1 commit into
mainfrom
codex/sync-save-sharding-validation-20261013
Draft

yaoyu-33 wants to merge 1 commit into
mainfrom
codex/sync-save-sharding-validation-20261013

Conversation

@yaoyu-33

Copy link
Copy Markdown
Contributor

Setting CheckpointConfig.ckpt_load_validate_sharding_integrity=False currently still sends validate_access_integrity=True to Core's global checkpoint loader and save path. The fully-parallel save wrapper also retains its default validation policy. This leaves the inherited checkpoint configuration ineffective in supported recipe training.

Forward the policy to the global loader and to both save layers when the installed Core wrapper supports the save-side option. Preserve the development pin's load-only semantics, cached constant-structure behavior, and existing MIMO exceptions. Constructor capability is checked once at import.

Bases: Bridge b1bb0a97ec5aad323eb2dd666e3f69aaebc1daff; upstream Core main a98aed553e25b4e86c59f50be54ac0ea4227b442. Bridge's pinned Core: e30d863a6d82a3c45f2d71355dc73fbf5610967b. Development pin: 4a4de865795f123eda422ebf0a4c13d2418cf3cb.

Source changes: NVIDIA/Megatron-LM#4422 and NVIDIA/Megatron-LM#5554. Related Bridge PR 6192 handles content-hash verification, a separate contract.

Validation on unchanged production base with test-only additions, then on this fix:

uv run python -m pytest tests/unit_tests/training/test_checkpointing.py -k 'save_checkpoint_global or global_checkpoint_load_honors_sharding_validation' -q -p no:cacheprovider

Before: exit 1, 11 failed and 53 passed; failure anchor assert True is False. After: exit 0, 64 passed. Cases include both policy values, first/cached saves, constant structure, MIMO, and a strict legacy save-wrapper constructor.

uv run python -m pytest tests/unit_tests/training/test_checkpointing.py -k 'TestSaveCheckpoint or TestCheckpointPathOverride or global_checkpoint_load_honors_sharding_validation' -q -p no:cacheprovider
uv run pre-commit run --all-files

Adjacent: exit 0, 87 passed and one skipped. All-files hooks: exit 0, no modified files. Native CPU unit fixtures mock checkpoint I/O and distributed calls. No multi-rank resume, full development-pin runtime, or performance result is claimed.

uv run mypy --strict src/megatron/bridge/training/checkpointing.py tests/unit_tests/training/test_checkpointing.py exits 2 with the same mypy 2.2.0 internal error on unchanged and patched sources; no type-check pass is claimed.

@yaoyu-33 yaoyu-33 added bug Something isn't working area:training Training loop, callbacks, and runtime integration area:ckpt Checkpoint conversion, loading, export, and save paths labels Oct 13, 2026
@copy-pr-bot

copy-pr-bot Bot commented Oct 13, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
@yaoyu-33
yaoyu-33 force-pushed the codex/sync-save-sharding-validation-20261013 branch from 6a40c9e to 8246192 Compare October 13, 2026 07:28
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 8246192

This branch was successfully deployed

1 active deployment
test — 8246192e Deployed Oct 13, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #22813
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:ckpt Checkpoint conversion, loading, export, and save paths area:training Training loop, callbacks, and runtime integration bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant