[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) triggers a CUDA illegal memory access during FP8 forward when a TE module lives on cuda:1 but the current device is cuda:0. The same model works on a single GPU.

TE resolves three things from the current CUDA device instead of from the tensors it is given: the kernel-launch context (including the runtime-compiled kernel cache), recipe-state allocation, and the delayed-scaling amax reduction. All three are fixed here in the Python layer.

Changes

• Execution device. prepare_forward pins the current device to the input's device for the duration of a module forward and end_forward restores it; the pin is also released if preparation raises. When the devices already match this is one integer compare. The fusible-ops OperationFuser.__call__ runs under the input's device the same way. Backward needs nothing: autograd already runs each node on its gradient's device.
• Recipe-state placement. RecipeState.create now receives the module's own device, derived from its parameters/buffers (get_module_device), instead of defaulting to the current device. The fusible ops construct their recipe state after registering the weight for the same reason. Weights created under quantized_model_init are quantized on their own device.
• Amax finalization. The registered amax tensors of one autocast may live on several devices, so a single torch.cat is not valid. Single-device buffers (checked in O(1) via a per-key device set maintained at registration) take the existing path unchanged. Multi-device buffers are grouped by device, preserving registration order, and updated on each device. With a distributed reduction, the groups are gathered to the first-registered device in registration order (one copy per device), reduced with the same single all-reduce as before, and scattered back, so the collective count, order and size are unchanged.

Testing

tests/pytorch/test_multi_device_fp8.py (needs 2 GPUs, never calls torch.cuda.set_device, asserts device placement rather than "no exception" because peer access masks the crash):

• module on a non-current device, with and without quantized_model_init
• cleanup when prepare_forward raises
• ops.BasicLinear on a non-current device
• two modules on different devices in one autocast
• bitwise equality of a 1-GPU and a 2-GPU split model over 4 iterations
• multi-device buffer with a mocked distributed reduction: the collective sees the whole buffer in registration order on the first-registered device, and results match the local path bit-for-bit

All 7 fail on unmodified main and pass with this PR.

Also run locally on main @ 224f6ec (H200 sm_90, 2.10 torch, CUDA 12.8):

• test_sanity.py: 8812 passed. test_fusible_ops.py: 1629 passed. test_custom_recipe.py, test_recipe.py: passed.
• test_numerics.py -k "fp8 or checkpoint or recompute", test_cuda_graphs.py, test_torch_compile.py: failure sets byte-identical to unmodified main (pre-existing attention mismatches from cuDNN 9.10 in this environment).
• Real NCCL check: 2 ranks × 2 local GPUs each, modules interleaved across the two GPUs, reduce_amax=True, 10 forward/backward iterations; both ranks produce the same output/state hash.
• pre-commit and pylint clean.

Out of scope / follow-ups

• Standalone quantizer(tensor) / tensor.dequantize() calls on a tensor that is not on the current device still launch on the current device. The thorough fix is a device guard in the C++ bindings; this PR covers the module, ops and construction paths.
• Process-global caches keyed without a device (get_dummy_wgrad, the ALiBi cache, skip_fp8_weight_update_tensor).
• Recipe-state placement on checkpoint restore for parameterless modules (DPA) falls back to the current device.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge.

Summary

• Pins module and fusible-operation execution to the input or gradient device.
• Allocates recipe state on module-owned devices and restores checkpoint state accordingly.
• Finalizes global amax state per device while preserving distributed registration order.
• Adds focused multi-device FP8 forward, backward, state-placement, and reduction tests.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  I[Input tensor device] --> G[Temporary CUDA device guard]
  G --> F[FP8 forward or fused operation]
  F --> R[Device-local recipe state]
  R --> A[Register amax entries]
  A --> D{Distributed reduction?}
  D -->|No| L[Finalize per owning device]
  D -->|Yes| C[Gather in registration order]
  C --> X[Single all-reduce]
  X --> S[Scatter to owning devices]
  L --> E[Restore ambient CUDA device]
  S --> E

Loading

_{Reviews (5) · Last reviewed commit: "[PyTorch] Fix FP8 illegal memory access ..."}

[SuperGoodGame]

Rebased onto current main (224f6ec) and squashed into one commit; the diff is about half the previous size (397+/84- vs 820+/78-). Dropped the cross-file device bookkeeping, the hand-rolled fast paths, the fuser backward wrapper and the standalone NCCL script, and fixed one gap the first version missed: weights created under quantized_model_init on a non-current device were still quantized on the current device, for te.Linear as well as the ops API. The description is updated accordingly.

Re-verified on current main: the reproducer still hits the illegal memory access without the patch. Locally, test_sanity, test_fusible_ops and the recipe/numerics/cuda-graph/torch.compile suites match unmodified main, and the real-NCCL 2 ranks x 2 GPUs check passes.

@timmoon10 @ksivaman could you take a look when you have a chance and trigger /te-ci pytorch L0? The design questions from #3124 (whether single-process multi-device is a topology TE wants to support, and gather-to-one-device vs per-device collectives) are still open. Happy to switch to the smaller "raise a clear error" variant if you'd prefer that.