Test MFSDP v2 attention activation offload

[wujingyue]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add integration coverage for MFSDP v2 using HybridModel with three MLA/MLP pairs (+-+-+-) and core_attn and attn_proj activation offload. The test compares six training steps against an identically initialized no-offload baseline, comparing the per-step microbatch losses.

The eager tests use flat ZeRO-3 with one optimizer instance. They use unfused attention (matching the pipeline-parallel offload tests), BF16 compute/main gradients, FP32 main weights, two microbatches, gradient clipping at 1.0, and max inflight offloads 1.

A pytest fixture handles process-group setup and offload/process-group teardown, including after assertion failures.

The loss-comparison test does not run a profiler. A separate transfer test requires positive offloaded bytes for both attention groups and positive DtoH/HtoD transfer counts when enabled, versus zero counts when disabled during the final steady-state iteration. The transfer test builds the model and optimizer, warms up once, then profiles one _train() call. Losses stay on the GPU, so the test introduces no loss-copy traffic into the trace.

This PR adds only the integration test. It needs no NVTE_CPU_OFFLOAD_V1 override. This exercises MCore's existing fine-grained offloader, not a migration to TE's newer offload implementation. It does not add sharding support or modify production offloading. MoE overlap, paged stash, and full-scale DeepSeek training are outside this test's scope.

Validation

• Two RTX A6000 GPUs: all three flat ZeRO-3 cases passed on the updated PR, with installed TE 2.18 (no TE overlay) and NVTE_CPU_OFFLOAD_V1 unset.
• On each rank, the final profiled step had 0 DtoH / 0 HtoD events with offloading disabled and 24 DtoH / 24 HtoD events with it enabled.
• Black, isort, Ruff, mypy, shell syntax, and diff whitespace checks passed.
• tools/autoformat.sh stops at the repository's obsolete Pylint bad-builtins option. Pylint passes with only unrecognized-option disabled.

env -u NVTE_CPU_OFFLOAD_V1 OMP_NUM_THREADS=4 \
uv run python -m torch.distributed.run --standalone --nproc-per-node=4 \
  -m pytest --experimental -v -o addopts= \
  tests/unit_tests/distributed/mfsdp_v2/test_activation_offload.py

Contribution process

• I have added relevant unit tests.
• I have run tools/autoformat.sh (environment limitation noted above).

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.