Opt in trusted FA CI checkpoints to pickle loading

[sudhakarsingh27]

Description

The L3 FA-version tests create and reload trusted delayed-scaling FP8 checkpoints, but their launcher did not opt into legacy pickle extra-state loading. This caused the FP16 and BF16 checkpoint cases to fail across the FA version matrix.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Infra/Build change

Changes

• Enable NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 for every L3 test_attention.py execution path.
• Keep the opt-in scoped to trusted test-generated checkpoints; runtime defaults and CP-only tests are unchanged.

Validation

Scenario	Result
Shell syntax	bash -n qa/L3_pytorch_FA_versions_test/test.sh passed
QA launcher audit	Every direct test_attention.py launcher has the trusted-checkpoint opt-in
Full GPU FA matrix	Pending PR checks

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 (not applicable: CI-only fix)
• My changes generate no new warnings
• I have added tests that prove my fix is effective (covered by the existing FA-version matrix)
• New and existing unit tests pass locally with my changes (requires GPU CI)

[sudhakarsingh27]

/te-ci pytorch L3

[greptile-apps]

Greptile Summary

This PR fixes L3 Flash Attention version tests that were failing because the test_attention.py launcher did not opt into legacy pickle extra-state loading when reloading trusted delayed-scaling FP8 checkpoints. The fix adds NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 inline to all three test_attention.py invocations in test.sh.

• All three test_attention.py execution paths (parallel ≥5 GPU, sequential <5 GPU, and non-CP FA iteration) now include the env var opt-in, scoped per-command so it does not bleed into other processes.
• test_attention_with_cp.py invocations are intentionally unchanged, since CP tests do not reload the FP8 checkpoint requiring pickle support.

Confidence Score: 5/5

Safe to merge — a minimal, targeted CI fix that adds a missing env var opt-in to three test launcher commands.

The change touches only one CI shell script, adding NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 inline to each test_attention.py invocation. All three execution branches are covered, the env var is correctly scoped per-command (no global export), and test_attention_with_cp.py is deliberately left unchanged. There is no production code, no logic change, and no missing case.

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 inline env var to all three test_attention.py invocations (parallel path, sequential path, and non-CP FA path); CP-only runs are correctly left unchanged.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[FA version loop] --> B{fa_version == CP_FA_VERSION?}
    B -- yes --> C{NUM_GPUS >= 5?}
    C -- yes --> D["CUDA_VISIBLE_DEVICES=0<br/>NVTE_TORCH_COMPILE=0<br/>NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 ✅<br/>test_attention.py & (GPU 0)"]
    C -- yes --> E["CUDA_VISIBLE_DEVICES=CP_GPUS<br/>NVTE_TORCH_COMPILE=0<br/>test_attention_with_cp.py &<br/>(no pickle opt-in, correct)"]
    C -- no --> F["NVTE_TORCH_COMPILE=0<br/>NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 ✅<br/>test_attention.py (sequential)"]
    C -- no --> G["NVTE_TORCH_COMPILE=0<br/>test_attention_with_cp.py (sequential)<br/>(no pickle opt-in, correct)"]
    B -- no --> H["NVTE_TORCH_COMPILE=0<br/>NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 ✅<br/>test_attention.py"]
    D --> I[wait / test_fail check]
    E --> I
    F --> J[test_fail check]
    G --> J
    H --> K[test_fail check]

Loading

_{Reviews (1): Last reviewed commit: "Opt in trusted FA CI checkpoints to pick..." | Re-trigger Greptile}

[sudhakarsingh27]

/te-ci pytorch L3