Skip to content

[PyTorch] Tests for loading previously-generated checkpoints - #1899

Merged
ksivaman merged 4 commits into
NVIDIA:mainfrom
timmoon10:checkpoint-test
Jul 9, 2025
Merged

[PyTorch] Tests for loading previously-generated checkpoints#1899
ksivaman merged 4 commits into
NVIDIA:mainfrom
timmoon10:checkpoint-test

Conversation

@timmoon10

@timmoon10timmoon10 commented Jun 26, 2025

Copy link
Copy Markdown
Member

Description

Some downstream users experienced pain when we merged #1033 because it affected the checkpoint format, namely by adding _extra_state to the state dict for the LayerNorm and RMSNorm modules. This PR adds tests for loading previously-generated checkpoints with strict=True. Even if we decide to change the checkpoint format in the future, these tests will make sure that it is a deliberate decision.

Note that these tests don't check correctness since that's tested in a few other places:

TransformerEngine/tests/pytorch/test_numerics.py

Line 846 in cc0cb35

deftest_gpt_checkpointing(dtype, bs, model):

TransformerEngine/tests/pytorch/test_fusible_ops.py

Line 1971 in cc0cb35

classTestCheckpointing:

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring
  • Testing

Changes

  • Add tests for loading previously-generated checkpoints

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@timmoon10
timmoon10 requested a review from ksivamanJune 26, 2025 22:53
@timmoon10timmoon10 added the testing Improvements to tests or testing infrastructure label Jun 26, 2025
@timmoon10

timmoon10 commented Jun 27, 2025

Copy link
Copy Markdown
MemberAuthor

/te-ci pytorch

1 similar comment
@timmoon10

timmoon10 commented Jun 30, 2025

Copy link
Copy Markdown
MemberAuthor

/te-ci pytorch

@timmoon10
Add tests for loading previously-generated checkpoints
9005ae1
Signed-off-by: Tim Moon <tmoon@nvidia.com>
ksivaman
ksivaman reviewed Jul 7, 2025
View reviewed changes
Comment threadqa/L0_pytorch_unittest/test.sh Outdated
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fused_attn.xml $TE_PATH/tests/pytorch/fused_attn/test_fused_attn.py || test_fail "test_fused_attn.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/fused_attn/test_kv_cache.py || test_fail "test_kv_cache.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py"
TE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py"

@ksivamanksivamanJul 7, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the CI purposes, this doesn't actually test older checkpoints, right?

@timmoon10timmoon10Jul 7, 2025•
edited
Loading

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm planning on keeping the files in our internal repos and modifying the CI pipeline to install them at this path. Unfortunately, this means non-NVIDIA users won't be able to run the tests themselves.

My original approach was to include the checkpoint files with Git LFS, but that could add some burden for users who don't have Git configured properly. We've had users struggle even with Git submodules, so I don't think it would be worth it to introduce LFS for such a basic test.

ksivaman
ksivaman reviewed Jul 7, 2025
View reviewed changes
Comment threadqa/L0_pytorch_unittest/test.sh Outdated
ksivaman
ksivaman reviewed Jul 7, 2025
View reviewed changes
Comment threadtests/pytorch/test_checkpoint.py Outdated
timmoon10and others added 3 commits July 7, 2025 11:32
@timmoon10@ksivaman
Use NVTE_ prefix for envvar
014c188
Review suggestion from @ksivaman
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
@timmoon10
Merge branch 'main' into checkpoint-test
4ebc5c9
@ksivaman
Merge branch 'main' into checkpoint-test
9f6e5ef
@ksivaman

ksivaman commented Jul 8, 2025

Copy link
Copy Markdown
Member

/te-ci pytorch

ksivaman
ksivaman approved these changes Jul 9, 2025
View reviewed changes

@ksivamanksivaman left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ksivaman
ksivaman merged commit 637facc into NVIDIA:mainJul 9, 2025
@yaox12

yaox12 commented Jul 10, 2025

Copy link
Copy Markdown
Member

@timmoon10@ksivaman This PR caused a CI failure. Can you take a look?

@timmoon10

timmoon10 commented Jul 10, 2025

Copy link
Copy Markdown
MemberAuthor

@yaox12#1934 changed the pickling logic for MXFP8 tensors, which broke checkpoint backward compatibility. #1943 should fix it.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ksivamanksivamanksivaman approved these changes

Assignees

No one assigned

Labels

testingImprovements to tests or testing infrastructure

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@timmoon10@ksivaman@yaox12