[PyTorch] Allow multi-version Flash Attention tests to use checkpoints with pickles

[timmoon10]

Description

#3123 added a guard to avoid using pickle in checkpoints. However, the legacy FP8 delayed scaling format still requires pickle. This causes failures in the multi-version Flash Attention tests. The fix is to allow checkpoints with pickles.

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
• Tests

Changes

• Allow multi-version Flash Attention tests to use checkpoints with pickles.

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]

/te-ci pytorch L3

[greptile-apps]

Greptile Summary

The multi-version Flash Attention test script was failing because FP8 delayed-scaling checkpoints use pickle, but PR #3123 introduced a guard that refuses to load pickled extra state by default. This fix enables the opt-in env var for the entire test suite so all FA version iterations can load those legacy checkpoints.

• Adds export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 before the FA version loop, unblocking checkpoint loading for delayed-scaling FP8 across all tested FA versions.
• Mirrors the pattern already used in qa/L0_pytorch_unittest/test.sh, where the same env var is set inline for the attention tests that exercise checkpointing.

Confidence Score: 5/5

Safe to merge — single-line addition to a QA test script with a clear comment explaining the reason.

The change is confined to a test harness script and simply enables an already-designed opt-in mechanism for loading legacy pickled checkpoints. The same env var is already used in the L0 unit test script for the same purpose, and all checkpoints in this test environment originate from the TE test suite itself.

No files require special attention.

Important Files Changed

Filename	Overview
qa/L3_pytorch_FA_versions_test/test.sh	Adds NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 as a global export before the test loop so that FP8 delayed-scaling checkpoints (which still use pickle) can be loaded across all Flash Attention version iterations.

Sequence Diagram

sequenceDiagram
    participant Script as test.sh
    participant Pytest as pytest (FA version loop)
    participant TE as TransformerEngine
    participant ES as _extra_state.py

    Script->>Script: "export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1"
    loop For each FA version (2.8.3, 3.0.0b1, 4.0.0b11)
        Script->>Pytest: run test_attention.py
        Pytest->>TE: load FP8 checkpoint (set_extra_state)
        TE->>ES: should_load_extra_state_pickle(data)
        ES->>ES: _classify_extra_state_pickle → UNSAFE_LOAD
        ES->>ES: "unsafe_pickle_extra_state_enabled() → True (env var=1)"
        ES-->>TE: return True (allow pickle)
        TE-->>Pytest: checkpoint loaded successfully
    end

Loading

_{Reviews (1): Last reviewed commit: "Allow Flash Attention tests to use check..." | Re-trigger Greptile}