[PyTorch] Fix FP8 illegal memory access in single-process multi-GPU execution

[SuperGoodGame]

Summary

Fixes #3124.

Single-process multi-GPU execution (for example, accelerate.dispatch_model or plain device_map placement) can trigger a CUDA illegal memory access during FP8 training when a TE module is on cuda:1 but the ambient current device is cuda:0. The same model works on a single GPU.

This PR fixes the three independent device assumptions in TE's Python FP8 path.

Root causes and fixes

A. Execution-device mismatch

FP8 forward setup and runtime-compiled kernel lookup use the ambient CUDA device. prepare_forward now temporarily pins the current device to the input device when they differ. The guard is released by end_forward, and is also released if preparation itself raises. The guard stack keeps nested forwards and activation recomputation balanced.

B. Recipe-state placement

RecipeState.create previously defaulted to the ambient CUDA device. Persistent state such as scale and amax_history is now allocated from module-owned parameters/buffers. This applies to both the legacy TE module API and the pytorch.ops API, including quantized-weight initialization. Checkpoint restore placement is fixed for module-owned recipe states whose execution device can be derived from parameters or buffers; parameterless modules such as DPA remain a follow-up.

The standalone operation fuser also runs forward and backward work under the input/gradient device context.

C. Global amax finalization

Registered amax tensors can belong to different CUDA devices, so a single torch.cat and fused update is not always valid.

• Without distributed reduction, entries are grouped and finalized locally on each owning device. No D2D copies or collectives are added.
• With distributed reduction, entries are gathered in registration order to the first registered device, reduced with the existing single all_reduce, then scattered back before per-device updates.
• global_amax_devices is rebuilt when registered tensors are replaced so the device index cannot become stale.

Ablation

The following was observed with a module on cuda:1 and ambient device cuda:0, enabling one fix at a time:

A	B	C	Result
-	-	-	IMA in the FP8 cast/quantize kernel
✓	-	-	IMA moves to amax workspace placement
-	✓	-	IMA remains in the cast kernel
✓	✓	-	Forward completes; failure moves to amax finalization
✓	✓	✓	Forward/backward and delayed-scaling state complete successfully

Testing

tests/pytorch/test_multi_device_fp8.py now contains seven focused tests:

1. A module on a non-current device, including forward/backward state and output placement.
2. Cleanup of the current-device guard when prepare_forward raises.
3. Device-aware recipe state and forward/backward execution through pytorch.ops.BasicLinear.
4. Two modules on different devices in one autocast context.
5. Bitwise comparison of a single-device and split-device model.
6. Ordered multi-device gather/reduce/scatter bookkeeping with a mocked collective.
7. Sentinel-based verification that the distributed collective receives entries in logical registration order and applies a simulated MAX reduction.

Additional validation performed locally:

• test_multi_device_fp8.py: 6 passed.
• Relevant fusible-operation tests: 272 passed.
• test_custom_recipe.py: 25 passed.
• Real NCCL harness: tests/pytorch/distributed/run_multi_device_fp8.py runs 2 processes × 2 GPUs per process for 10 forward/backward iterations with one all-reduce per direction; both ranks produced the same final hash.
• Pre-commit hooks: passed.

Run the real-NCCL check with:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=2 tests/pytorch/distributed/run_multi_device_fp8.py

Performance

No stable performance regression was observed in a CUDA-event microbenchmark against the PR's original HEAD. The normal TE module path only adds cleanup around exceptional paths; the ops path adds a conditional current-device check per fuser call. The measured variation was within run-to-run noise.

Known environment limitation

The current test environment uses cuDNN 9.10.2. Some strict CUDA-graph attention comparisons require cuDNN 9.15.1 or newer and fail with small eager versus graph numerical differences. The representative failure reproduces on the unmodified base commit as well.

Related process-global caches whose keys do not include a device are outside this PR and should be handled separately.

[greptile-apps]

Greptile Summary

The PR makes FP8 execution device-aware for single-process multi-GPU models and repairs cleanup after forward preparation fails.

• Pins module and operation execution to the input or gradient device.
• Allocates persistent recipe state on module-owned devices.
• Finalizes global amax state per device while preserving distributed registration order.
• Adds focused local and distributed multi-GPU validation.

Confidence Score: 5/5

The PR appears safe to merge.

The previously reported device-guard leak is fixed because prepare_forward now releases the entered guard on every preparation exception, while successful preparations remain paired with end_forward; no blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Adds device-aware recipe-state placement and balanced CUDA-device cleanup for successful and exceptional forward preparation.
transformer_engine/pytorch/quantization.py	Tracks registered amax devices and performs local or distributed delayed-scale finalization without combining tensors directly across devices.
transformer_engine/pytorch/ops/fuser.py	Runs fused forward and backward work under the corresponding input or gradient CUDA device.
transformer_engine/pytorch/ops/op.py	Allocates operation recipe state from operation-owned CUDA parameters or buffers and refreshes global device bookkeeping after state replacement.
tests/pytorch/test_multi_device_fp8.py	Covers off-current-device execution, preparation failure cleanup, recipe-state placement, multi-device finalization, and distributed registration ordering.
tests/pytorch/distributed/run_multi_device_fp8.py	Provides a real NCCL validation harness for two ranks with two local GPUs per process.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[prepare_forward receives input] --> B{Input device differs from current device?}
    B -- Yes --> C[Enter CUDA device guard]
    B -- No --> D[Push empty guard entry]
    C --> E[Push guard entry]
    D --> F[Prepare FP8 state]
    E --> F
    F --> G{Preparation succeeds?}
    G -- No --> H[Pop NVTX if needed]
    H --> I[Release device guard]
    G -- Yes --> J[Run module forward]
    J --> K[end_forward in finally]
    K --> L[Restore recompute state]
    L --> M[Pop NVTX range]
    M --> I

Loading

_{Reviews (4): Last reviewed commit: "[PyTorch] Validate mixed-device amax ord..." | Re-trigger Greptile}