diff --git a/tests/pytorch/test_multi_device_fp8.py b/tests/pytorch/test_multi_device_fp8.py
new file mode 100644
index 0000000000..8c8273839e
--- /dev/null
+++ b/tests/pytorch/test_multi_device_fp8.py
@@ -0,0 +1,192 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Single-process multi-device FP8 tests (issue #3124).
+
+These tests never call ``torch.cuda.set_device`` and never enable peer access
+before the first forward: the current device stays ``cuda:0`` while the modules
+under test live on other devices, as with ``accelerate.dispatch_model`` and plain
+``device_map`` placement. Placement is asserted explicitly because, once peer
+access is enabled, the unfixed code silently runs on the wrong device instead of
+raising.
+"""
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+import transformer_engine.pytorch.ops as te_ops
+from transformer_engine.common.recipe import DelayedScaling, Format
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+
+pytestmark = pytest.mark.skipif(
+    torch.cuda.device_count() < 2, reason="needs >= 2 CUDA devices in one process"
+)
+
+DT = torch.bfloat16
+RECIPE = DelayedScaling(fp8_format=Format.HYBRID)
+
+
+@pytest.fixture(autouse=True)
+def _fresh_fp8_state():
+    FP8GlobalStateManager.reset()
+    yield
+    FP8GlobalStateManager.reset()
+
+
+def _assert_state_on(module, device):
+    """Recipe state lives on ``device`` and the delayed-scaling update ran there.
+
+    The fused update rolls the history so the newest amax lands in the last row
+    (row 0 is zeroed for the next iteration) and moves the scale off its init
+    value of 1.
+    """
+    for key in ("scaling_fwd", "scaling_bwd"):
+        state = module.fp8_meta[key]
+        assert state.scale.device == device
+        assert state.amax_history.device == device
+        assert not torch.all(state.amax_history[-1] == 0), f"{key} amax never updated"
+        assert not torch.all(state.scale == 1.0), f"{key} scale never updated"
+
+
+@pytest.mark.parametrize("quantized_weight", (False, True))
+def test_module_off_current_device(quantized_weight):
+    """One module on cuda:1 while the current device is cuda:0 (#3124 case A)."""
+    assert torch.cuda.current_device() == 0
+    with te.quantized_model_init(enabled=quantized_weight, recipe=RECIPE):
+        module = te.Linear(512, 1024, bias=False, params_dtype=DT, device="cuda:1")
+    inp = torch.randn(128, 512, device="cuda:1", dtype=DT, requires_grad=True)
+    with te.autocast(enabled=True, recipe=RECIPE):
+        out = module(inp)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    assert out.device == torch.device("cuda:1")
+    assert inp.grad.device == torch.device("cuda:1")
+    assert torch.cuda.current_device() == 0
+    _assert_state_on(module, torch.device("cuda:1"))
+
+
+def test_prepare_forward_exception_restores_current_device(monkeypatch):
+    """A failed prepare must not leak the temporary CUDA device guard."""
+    module = te.Linear(16, 16, bias=False, params_dtype=DT, device="cuda:1")
+    inp = torch.randn(2, 16, device="cuda:1", dtype=DT)
+
+    def fail_init(*args, **kwargs):
+        raise RuntimeError("injected prepare_forward failure")
+
+    monkeypatch.setattr(module, "init_fp8_metadata", fail_init)
+    with pytest.raises(RuntimeError, match="injected prepare_forward failure"):
+        module.prepare_forward(inp)
+    assert torch.cuda.current_device() == 0
+
+
+def test_basic_operation_off_current_device():
+    """The fusible-ops API must allocate state and execute on the op's device."""
+    with te.quantized_model_init(enabled=True, recipe=RECIPE):
+        op = te_ops.basic.BasicLinear(512, 1024, device="cuda:1", dtype=DT)
+    inp = torch.randn(128, 512, device="cuda:1", dtype=DT, requires_grad=True)
+    with te.autocast(enabled=True, recipe=RECIPE):
+        out = op(inp)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    assert out.device == torch.device("cuda:1")
+    assert inp.grad.device == torch.device("cuda:1")
+    assert torch.cuda.current_device() == 0
+    for mode in ("forward", "backward"):
+        state = op._fp8_metas[mode][FP8GlobalStateManager.get_meta_tensor_key(mode == "forward")]
+        assert state.scale.device == torch.device("cuda:1")
+        assert state.amax_history.device == torch.device("cuda:1")
+
+
+def test_two_modules_on_different_devices_one_autocast():
+    """Two modules on different devices inside one autocast (#3124 case C)."""
+    a = te.Linear(512, 1024, bias=False, params_dtype=DT, device="cuda:0")
+    b = te.Linear(1024, 512, bias=False, params_dtype=DT, device="cuda:1")
+    for _ in range(3):
+        inp = torch.randn(128, 512, device="cuda:0", dtype=DT, requires_grad=True)
+        with te.autocast(enabled=True, recipe=RECIPE):
+            hidden = a(inp)
+            out = b(hidden.to("cuda:1"))
+        out.sum().backward()
+    torch.cuda.synchronize()
+
+    assert hidden.device == torch.device("cuda:0")
+    assert out.device == torch.device("cuda:1")
+    _assert_state_on(a, torch.device("cuda:0"))
+    _assert_state_on(b, torch.device("cuda:1"))
+
+
+def _run_chain(devices, iters=4):
+    """Run a chain of Linears placed on ``devices``; return outputs and FP8 state."""
+    modules = []
+    for i, (device, (in_f, out_f)) in enumerate(
+        zip(devices, ((512, 1024), (1024, 512), (512, 256)))
+    ):
+        torch.manual_seed(1234 + i)
+        modules.append(te.Linear(in_f, out_f, bias=False, params_dtype=DT, device=device))
+
+    outs, states = [], []
+    for it in range(iters):
+        torch.manual_seed(100 + it)
+        x = torch.randn(128, 512, device="cuda:0", dtype=DT, requires_grad=True)
+        with te.autocast(enabled=True, recipe=RECIPE):
+            for module, device in zip(modules, devices):
+                x = module(x.to(device))
+        x.sum().backward()
+        torch.cuda.synchronize()
+        outs.append(x.detach().float().cpu())
+    for module in modules:
+        for key in ("scaling_fwd", "scaling_bwd"):
+            states.append(module.fp8_meta[key].scale.cpu())
+            states.append(module.fp8_meta[key].amax_history.cpu())
+    return outs, states
+
+
+def test_split_model_matches_single_device_bitwise():
+    """Same weights and inputs: the 2-GPU split must match the 1-GPU run exactly."""
+    ref = _run_chain(("cuda:0", "cuda:0", "cuda:0"))
+    split = _run_chain(("cuda:0", "cuda:1", "cuda:0"))
+    for r, s in zip(ref[0] + ref[1], split[0] + split[1]):
+        torch.testing.assert_close(r, s, rtol=0, atol=0)
+
+
+def test_multi_device_amax_reduction(monkeypatch):
+    """Multi-device buffer with a distributed amax reduction.
+
+    A fake two-rank world routes the interleaved (cuda:0, cuda:1, cuda:0) buffer
+    through the gather -> collective -> scatter path. The collective must see the
+    whole buffer in registration order on the first-registered device, and with the
+    collective mocked as the identity the results must be bit-identical to the
+    fully local path.
+    """
+    fake_dist = {"on": True}
+    calls = []
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: fake_dist["on"])
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda *a, **k: 2)
+
+    def fake_reduce(tensor, _group):
+        qstate = FP8GlobalStateManager.quantization_state
+        entries = [
+            entries
+            for entries in qstate.global_amax_buffer.values()
+            if sum(entry.numel() for entry in entries) == tensor.numel()
+        ]
+        assert len(entries) == 1
+        expected = torch.cat([entry.to(tensor.device) for entry in entries[0]])
+        torch.testing.assert_close(tensor, expected, rtol=0, atol=0)
+        calls.append(tensor.device)
+
+    monkeypatch.setattr(
+        FP8GlobalStateManager, "reduce_tensor_across_group_op_max", staticmethod(fake_reduce)
+    )
+    devices = ("cuda:0", "cuda:1", "cuda:0")
+    gathered = _run_chain(devices)
+    fake_dist["on"] = False
+    local = _run_chain(devices)
+
+    for g, l in zip(gathered[0] + gathered[1], local[0] + local[1]):
+        torch.testing.assert_close(g, l, rtol=0, atol=0)
+    # One collective per direction per iteration, on the first-registered device.
+    assert calls == [torch.device("cuda:0")] * 8
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index a3131f7436..bf4b59ac3d 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -62,6 +62,7 @@
     is_non_tn_fp8_gemm_supported,
     torch_get_autocast_gpu_dtype,
     get_device_compute_capability,
+    get_module_device,
     get_nvtx_range_context,
     nvtx_range_push,
     nvtx_range_pop,
@@ -918,6 +919,9 @@ def __init__(self, name: Optional[str] = None) -> None:
         self.wgrad_store = None
         self._output_quantizer_role: Optional[QuantizerRole] = None
         self._grad_input_quantizer_role: Optional[QuantizerRole] = None
+        # Current CUDA devices saved by prepare_forward and restored by end_forward.
+        # A stack so nested and activation-recompute forwards stay balanced.
+        self._forward_prev_devices: List[int] = []
 
         if not TEDebugState.debug_enabled:
             TEDebugState.initialize()
@@ -1151,6 +1155,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
             mode=("forward" if fwd else "backward"),
             num_quantizers=num_fp8_tensors,
             roles=roles,
+            device=get_module_device(self),
         )
 
         # Reached the rebuild path because ``fp8_meta_tensors_initialized``
@@ -1592,53 +1597,78 @@ def prepare_forward(
         allow_different_data_and_param_types: bool = False,
     ) -> torch.Tensor:
         """Checks and prepares for FWD execution."""
-        self.fast_setattr(
-            "allow_different_data_and_param_types", allow_different_data_and_param_types
-        )
-        self.fast_setattr("forwarded_at_least_once", True)
-
-        # Activation recomputation is used and this is the second forward phase.
-        if self.fp8 and in_fp8_activation_recompute_phase():
-            delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta)
-            FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta)
+        # TE kernels resolve their launch device from the current CUDA device, not
+        # from ``inp`` (e.g. the runtime-compiled kernel cache is keyed on
+        # cuda::current_device()). Pin the current device to the input's for the
+        # duration of the forward so that a module living off the current device
+        # (single-process multi-GPU, e.g. accelerate.dispatch_model) does not raise
+        # an illegal memory access. Restored in ``end_forward``, or here if
+        # preparation fails.
+        prev_device = torch.cuda.current_device()
+        if inp.is_cuda and inp.device.index != prev_device:
+            torch.cuda.set_device(inp.device.index)
         else:
-            if not inp.is_cuda:
-                raise RuntimeError(
-                    f"TransformerEngine needs CUDA. Got input on device: {inp.device}"
-                )
+            prev_device = -1
+        self._forward_prev_devices.append(prev_device)
+        try:
+            self.fast_setattr(
+                "allow_different_data_and_param_types", allow_different_data_and_param_types
+            )
+            self.fast_setattr("forwarded_at_least_once", True)
 
-            if self.tp_size > 1:
-                if not self.tp_group_initialized:
+            # Activation recomputation is used and this is the second forward phase.
+            if self.fp8 and in_fp8_activation_recompute_phase():
+                delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta)
+                FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta)
+            else:
+                if not inp.is_cuda:
                     raise RuntimeError(
-                        "Tensor parallel group not initialized. Call "
-                        "set_tensor_parallel_group() before forward pass when tp_size > 1."
+                        f"TransformerEngine needs CUDA. Got input on device: {inp.device}"
                     )
 
-            self.set_activation_dtype(inp)
-            self.init_fp8_metadata(num_gemms=num_gemms)
-            self._check_weight_tensor_recipe_correspondence()
-
-            delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta)
-            if delayed_scaling_recipe:
-                if self.sequence_parallel:
-                    assert (
-                        self.fp8_meta["recipe"].custom() or self.fp8_meta["recipe"].reduce_amax
-                    ), (
-                        "Amax reduction across tensor parallel group is "
-                        "necessary when using sequence parallelism with FP8."
-                    )
+                if self.tp_size > 1:
+                    if not self.tp_group_initialized:
+                        raise RuntimeError(
+                            "Tensor parallel group not initialized. Call "
+                            "set_tensor_parallel_group() before forward pass when tp_size > 1."
+                        )
 
-                if not FP8GlobalStateManager.fp8_graph_capturing():
-                    FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
+                self.set_activation_dtype(inp)
+                self.init_fp8_metadata(num_gemms=num_gemms)
+                self._check_weight_tensor_recipe_correspondence()
+
+                delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta)
+                if delayed_scaling_recipe:
+                    if self.sequence_parallel:
+                        assert (
+                            self.fp8_meta["recipe"].custom() or self.fp8_meta["recipe"].reduce_amax
+                        ), (
+                            "Amax reduction across tensor parallel group is "
+                            "necessary when using sequence parallelism with FP8."
+                        )
+
+                    if not FP8GlobalStateManager.fp8_graph_capturing():
+                        FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
+
+                    # Activation recomputation is used and this is the first forward phase.
+                    if is_fp8_activation_recompute_enabled():
+                        FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(
+                            self.fp8_meta
+                        )
 
-                # Activation recomputation is used and this is the first forward phase.
-                if is_fp8_activation_recompute_enabled():
-                    FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)
+            nvtx_range_push(self.__class__.__name__ + " forward")
+            if not allow_non_contiguous and not inp.is_contiguous():
+                inp = inp.contiguous()
+            return inp
+        except BaseException:
+            self._restore_forward_device()
+            raise
 
-        nvtx_range_push(self.__class__.__name__ + " forward")
-        if not allow_non_contiguous and not inp.is_contiguous():
-            inp = inp.contiguous()
-        return inp
+    def _restore_forward_device(self) -> None:
+        """Restore the current CUDA device saved by ``prepare_forward``."""
+        prev_device = self._forward_prev_devices.pop()
+        if prev_device >= 0:
+            torch.cuda.set_device(prev_device)
 
     def end_forward(self):
         """
@@ -1649,6 +1679,7 @@ def end_forward(self):
         if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase():
             FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta)
         nvtx_range_pop()
+        self._restore_forward_device()
 
     @contextmanager
     def prepare_forward_ctx(
@@ -1861,8 +1892,9 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None:
                     )
                     quantizer.amax_reduction_group = amax_reduction_group
                     quantizer.with_amax_reduction = True
-                # Quantize parameter
-                param = quantizer(param)
+                # Quantize parameter on its own device (see prepare_forward)
+                with torch.cuda.device(param.get_device()):
+                    param = quantizer(param)
 
             # Redo parameter wrap in case we broke it above
             # NOTE: Currently this can only be broken when primary weights are in Fp8 but
diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py
index cb429055a4..fabb6781d1 100644
--- a/transformer_engine/pytorch/ops/basic/basic_linear.py
+++ b/transformer_engine/pytorch/ops/basic/basic_linear.py
@@ -143,11 +143,6 @@ def __init__(
             out_features=out_features,
         )
 
-        # Initialize recipe state if needed for natively quantized weight
-        self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters()
-        if self._with_quantized_weight:
-            self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe())
-
         # Initialize parameters if needed
         weight = torch.empty(
             self.local_out_features,
@@ -160,6 +155,13 @@ def __init__(
         self.register_parameter("weight", weight)
         self._rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]]
         self._rng_state_tracker_function = rng_state_tracker_function
+
+        # Initialize recipe state if needed for natively quantized weight
+        # Note: After registering the weight so the state is allocated on its device.
+        self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters()
+        if self._with_quantized_weight:
+            self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe())
+
         if weight.device.type != "meta":
             self.reset_parameters()
 
@@ -333,7 +335,8 @@ def reset_parameters(self) -> None:
                 columnwise=torch.is_grad_enabled(),
             )
             quantizer.internal = False
-            with torch.no_grad():
+            # Quantize on the weight's device: TE kernels launch on the current device
+            with torch.cuda.device(weight.get_device()), torch.no_grad():
                 weight = quantizer(weight)
 
         # Save updated parameter
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index 9551650045..be8d138847 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -239,11 +239,6 @@ def __init__(
         if dtype not in (torch.float32, torch.float16, torch.bfloat16):
             raise ValueError(f"Supported dtypes are float32, float16, bfloat16 (got {dtype})")
 
-        # Initialize recipe state if needed for natively quantized weight
-        self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters()
-        if self._with_quantized_weight:
-            self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe())
-
         # RNG state tracker
         self._rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]]
         self._rng_state_tracker_function = rng_state_tracker_function
@@ -278,6 +273,12 @@ def __init__(
                 bias_tensor = torch.nn.Parameter(bias_tensor)
             self.register_parameter(f"bias{group_idx}", bias_tensor)
 
+        # Initialize recipe state if needed for natively quantized weight
+        # Note: After registering the weights so the state is allocated on their device.
+        self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters()
+        if self._with_quantized_weight:
+            self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe())
+
         # Initialize weights if needed
         if device.type != "meta":
             self.reset_parameters()
@@ -479,8 +480,9 @@ def reset_parameters(self) -> None:
                 )
                 quantizer.internal = False
 
-            # Quantize weights
-            weights = self._quantize_weights(weights, quantizers)
+            # Quantize on the weights' device: TE kernels launch on the current device
+            with torch.cuda.device(weights[0].get_device()):
+                weights = self._quantize_weights(weights, quantizers)
 
         # Register weights
         for group_idx, weight in enumerate(weights):
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 2500002700..0bb8655000 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -713,6 +713,17 @@ def __call__(
         *extra_inputs: torch.Tensor,
         basic_op_kwargs: Optional[list[dict[str, Any]]] = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        # TE kernels resolve their launch device from the current CUDA device, so
+        # run under the input's device (single-process multi-GPU placement).
+        with torch.cuda.device(input.get_device()):
+            return self._call_impl(input, *extra_inputs, basic_op_kwargs=basic_op_kwargs)
+
+    def _call_impl(
+        self,
+        input: torch.Tensor,  # pylint: disable=redefined-builtin
+        *extra_inputs: torch.Tensor,
+        basic_op_kwargs: Optional[list[dict[str, Any]]] = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
 
         # Verify extra input count
         if len(extra_inputs) != self.num_extra_inputs:
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index d057d46816..c71b44e1f2 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -22,6 +22,7 @@
     autocast,
 )
 from ..tensor import Quantizer
+from ..utils import get_module_device
 
 
 @dataclasses.dataclass
@@ -370,6 +371,7 @@ def reset_recipe_state(
                     mode=mode,
                     num_quantizers=num_quantizers,
                     roles=roles,
+                    device=get_module_device(self),
                 )
                 fp8_meta_key = FP8GlobalStateManager.get_meta_tensor_key(
                     forward=(mode == "forward"),
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 34ae20b498..aa8644a58d 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -405,6 +405,9 @@ class FP8GlobalState:
     global_amax_buffer: Dict[str, list] = field(default_factory=dict)
     global_amax_history_buffer: Dict[str, list] = field(default_factory=dict)
     global_scale_buffer: Dict[str, list] = field(default_factory=dict)
+    # Devices of each key's registered amax tensors. Modules placed off the
+    # current device (single-process multi-GPU) make a buffer span devices.
+    global_amax_devices: Dict[str, set] = field(default_factory=dict)
     fp8_tensors_recompute_buffer: list = field(default_factory=list)
     autocast_arguments: Dict[Any, Tuple[Recipe, Optional[dist_group_type]]] = field(
         default_factory=dict
@@ -555,12 +558,14 @@ def add_fp8_tensors_to_global_buffer(
                     fp8_meta[fp8_meta_tensor_key].amax_history
                 ]
                 qstate.global_scale_buffer[key] = [fp8_meta[fp8_meta_tensor_key].scale]
+                qstate.global_amax_devices[key] = {fp8_meta[fp8_meta_tensor_key].scale.device}
             else:
                 qstate.global_amax_buffer[key].append(fp8_meta[fp8_meta_tensor_key].amax_history[0])
                 qstate.global_amax_history_buffer[key].append(
                     fp8_meta[fp8_meta_tensor_key].amax_history
                 )
                 qstate.global_scale_buffer[key].append(fp8_meta[fp8_meta_tensor_key].scale)
+                qstate.global_amax_devices[key].add(fp8_meta[fp8_meta_tensor_key].scale.device)
             fp8_meta[index_in_buffer].append(len(qstate.global_amax_buffer[key]) - 1)
             fp8_meta[index_in_buffer].append(key)
 
@@ -652,6 +657,40 @@ def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_ty
                 async_op=False,
             )
 
+    @classmethod
+    def _update_amax_histories_and_scales(
+        cls,
+        contiguous_amax: torch.Tensor,
+        amax_buffer: List[torch.Tensor],
+        amax_histories: List[torch.Tensor],
+        scales: List[torch.Tensor],
+        recipe: Recipe,
+        forward: bool,
+    ) -> None:
+        """Update amax histories and scales from reduced amaxes that all live on one device."""
+        unfused_update = (
+            bool(int(os.getenv("NVTE_UNFUSED_FP8_UPDATE", "0")))
+            or callable(recipe.amax_compute_algo)
+            or callable(recipe.scaling_factor_compute_algo)
+        )
+        # The fused kernel launches on the current device, not the tensors' device.
+        with torch.cuda.device(contiguous_amax.device):
+            if not unfused_update:
+                tex.fused_amax_and_scale_update_after_reduction(
+                    contiguous_amax,
+                    amax_histories,
+                    scales,
+                    recipe.amax_compute_algo,
+                    get_fp8_te_dtype(recipe, forward),
+                    recipe.margin,
+                )
+            else:
+                split_and_copy(contiguous_amax, amax_buffer, [x.numel() for x in amax_buffer])
+                for amax_history, scale in zip(amax_histories, scales):
+                    _amax_and_scale_update(
+                        amax_history, scale, get_fp8_max(recipe, forward), recipe
+                    )
+
     @classmethod
     def reduce_and_update_fp8_tensors(
         cls,
@@ -671,44 +710,63 @@ def reduce_and_update_fp8_tensors(
             if len(amax_buffer) == 0:
                 continue
 
-            # Retrieve autocast specific args and concat amaxes.
+            # Retrieve autocast specific args.
             recipe, group = qstate.autocast_arguments[autocast_key]
-            contiguous_amax = torch.cat(amax_buffer)
-
-            # Reduction.
-            if (
+            amax_histories = qstate.global_amax_history_buffer[buffer_key]
+            scales = qstate.global_scale_buffer[buffer_key]
+            need_reduce = (
                 recipe.reduce_amax
                 and torch.distributed.is_initialized()
                 and torch.distributed.get_world_size(group=group) > 1
-            ):
-                cls.reduce_tensor_across_group_op_max(contiguous_amax, group)
-
-            # Amax and scale update.
-            unfused_update = (
-                bool(int(os.getenv("NVTE_UNFUSED_FP8_UPDATE", "0")))
-                or callable(recipe.amax_compute_algo)
-                or callable(recipe.scaling_factor_compute_algo)
             )
 
-            if not unfused_update:
-                tex.fused_amax_and_scale_update_after_reduction(
-                    contiguous_amax,
-                    qstate.global_amax_history_buffer[buffer_key],
-                    qstate.global_scale_buffer[buffer_key],
-                    recipe.amax_compute_algo,
-                    get_fp8_te_dtype(recipe, forward),
-                    recipe.margin,
+            if len(qstate.global_amax_devices[buffer_key]) == 1:
+                contiguous_amax = torch.cat(amax_buffer)
+                if need_reduce:
+                    cls.reduce_tensor_across_group_op_max(contiguous_amax, group)
+                cls._update_amax_histories_and_scales(
+                    contiguous_amax, amax_buffer, amax_histories, scales, recipe, forward
                 )
-            else:
-                split_and_copy(contiguous_amax, amax_buffer, [x.numel() for x in amax_buffer])
+                continue
 
-                for amax_history, scale in zip(
-                    qstate.global_amax_history_buffer[buffer_key],
-                    qstate.global_scale_buffer[buffer_key],
-                ):
-                    _amax_and_scale_update(
-                        amax_history, scale, get_fp8_max(recipe, forward), recipe
-                    )
+            # Modules registered under one autocast live on several CUDA devices
+            # (single-process multi-GPU, e.g. accelerate.dispatch_model). Group the
+            # entries by device, keeping registration order within each group, and
+            # finalize each group on its own device.
+            by_device: Dict[torch.device, List[int]] = {}
+            for i, amax in enumerate(amax_buffer):
+                by_device.setdefault(amax.device, []).append(i)
+            group_amaxes = {
+                dev: torch.cat([amax_buffer[i] for i in idxs]) for dev, idxs in by_device.items()
+            }
+            if need_reduce:
+                # Keep the single collective over the whole buffer in registration
+                # order, so entry i still refers to the same quantizer on every
+                # rank: gather each device's entries to the first-registered device
+                # with one copy per device, reduce, and scatter back the same way.
+                reduce_device = amax_buffer[0].device
+                numels = [amax.numel() for amax in amax_buffer]
+                parts: List[torch.Tensor] = [None] * len(amax_buffer)
+                for dev, idxs in by_device.items():
+                    chunks = group_amaxes[dev].to(reduce_device).split([numels[i] for i in idxs])
+                    for i, chunk in zip(idxs, chunks):
+                        parts[i] = chunk
+                contiguous_amax = torch.cat(parts)
+                cls.reduce_tensor_across_group_op_max(contiguous_amax, group)
+                chunks = contiguous_amax.split(numels)
+                group_amaxes = {
+                    dev: torch.cat([chunks[i] for i in idxs]).to(dev)
+                    for dev, idxs in by_device.items()
+                }
+            for dev, idxs in by_device.items():
+                cls._update_amax_histories_and_scales(
+                    group_amaxes[dev],
+                    [amax_buffer[i] for i in idxs],
+                    [amax_histories[i] for i in idxs],
+                    [scales[i] for i in idxs],
+                    recipe,
+                    forward,
+                )
 
     @staticmethod
     def get_unique_autocast_key(
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index 513632e095..4ea9d8eb2c 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -822,6 +822,19 @@ def canonicalize_dtype(dtype: Optional[torch.dtype]) -> torch.dtype:
     return dtype
 
 
+def get_module_device(module: torch.nn.Module) -> torch.device:
+    """CUDA device of a module's parameters or buffers, else the current CUDA device.
+
+    Per-module state (e.g. FP8 scales and amax histories) must live on the module's
+    device, which is not necessarily the current device in single-process multi-GPU
+    execution (e.g. accelerate.dispatch_model).
+    """
+    for tensor in (*module.parameters(), *module.buffers()):
+        if tensor.device.type == "cuda":
+            return tensor.device
+    return torch.device("cuda", torch.cuda.current_device())
+
+
 def devices_match(device1: torch.device, device2: torch.device) -> bool:
     """Whether two devices are the same"""
     device1 = torch.device(device1)