Fix deepgemm on multiple devices

src/transformers/integrations/deepgemm.py

@@ -16,7 +16,7 @@
 
 Provides:
 - `deepgemm_bf16_experts_forward`: BF16 M-grouped experts forward.
-- `deepgemm_fp8_fp4_linear`: end-to-end FP8/FP4 linear (BF16 in, BF16 out).
+- `deepgemm_fp8_fp4_linear`: end-to-end FP8/FP4 linear (output dtype follows the input).
 - `deepgemm_fp8_fp4_experts_forward`: FP8 (or FP4 on SM100+) M-grouped experts forward.
 - `deepgemm_fp8_fp4_megamoe_experts_forward`: FP8xFP4 Mega MoE forward (SM100+).
 
@@ -314,45 +314,6 @@ def _assert_sm100_scales_are_ue8m0(scale: torch.Tensor) -> None:
     )
 
 
-_DEEPGEMM_VISITED_DEVICES: set[int] = set()
-
-
-def _assert_single_device(device: torch.device, context: str) -> None:
-    """Reject DeepGEMM calls that span multiple CUDA devices in the same process
-    (e.g. ``device_map="auto"`` across N GPUs). DeepGEMM loads each kernel via
-    ``cuKernelGetFunction``, which binds the resulting ``CUfunction`` handle to
-    the CUDA context that was current at load time. Driving the same cached
-    handle from a different device's context launches it against the wrong
-    module/context and produces garbage. Distributed setups (torchrun + TP/EP)
-    don't trip this because each process owns exactly one device's context.
-
-    The fix is a build-time choice on the DeepGEMM side: compiling with
-    ``DG_JIT_USE_RUNTIME_API=1`` swaps the loader for the runtime API
-    (context-free ``cudaKernel_t``) and lifts the restriction — but it has to
-    be baked into the wheel, setting the env var at Python runtime won't change
-    the loader the cached ``.so`` already uses. Until the kernels-community build
-    we ship picks that up, we reject single-process multi-device by default.
-
-    Raised as :class:`ImportError` from the per-linear path so :func:`fp8_linear`
-    falls back to Triton (which loads through the runtime API and has no such
-    binding); raised as :class:`RuntimeError` from the experts path where there's
-    no fallback — the user explicitly chose ``experts_implementation="deepgemm"``
-    and must switch to ``"grouped_mm"`` / ``"eager"`` or run distributed.
-    """
-    idx = device.index if device.index is not None else torch.cuda.current_device()
-    _DEEPGEMM_VISITED_DEVICES.add(idx)
-    if len(_DEEPGEMM_VISITED_DEVICES) <= 1:
-        return
-    msg = (
-        "DeepGEMM caches each kernel's `CUfunction` against the CUDA context it was first "
-        "loaded under, so driving it from a different device in the same process produces "
-        "garbage. Run distributed (TP/EP) so each process owns one device, "
-    )
-    if context == "linear":
-        raise ImportError(msg + "or fall back to the Triton kernel (handled automatically).")
-    raise RuntimeError(msg + "or pick `experts_implementation='grouped_mm'`.")
-
-
 def _ceil_to_ue8m0(sf: torch.Tensor) -> torch.Tensor:
     """Round each fp32 SF up to the nearest power of 2 (zero mantissa).
 
@@ -588,16 +549,13 @@ def deepgemm_fp8_fp4_linear(
     weight_scale_inv: torch.Tensor,
     bias: torch.Tensor | None = None,
     block_size: tuple[int, int] | None = None,
-    output_dtype: torch.dtype = torch.bfloat16,
     activation_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """End-to-end DeepGEMM linear: per-token activation quant + FP8/FP4 matmul.
 
     Static (per-tensor) activation quantization is rejected — DeepGEMM needs
     per-row SFs. Callers should route static activations through the Triton fallback.
     """
-    _assert_single_device(input.device, context="linear")
-
     if activation_scale is not None:
         raise NotImplementedError("DeepGEMM linear does not support static activation quantization.")
     if input.dtype not in (torch.bfloat16, torch.float16):
@@ -608,7 +566,7 @@ def deepgemm_fp8_fp4_linear(
 
     input_2d = input.view(-1, input.shape[-1])
     qinput_2d, scale_2d = deepgemm.per_token_cast_to_fp8(input_2d, **cast_kwargs)
-    output = torch.empty(qinput_2d.shape[0], weight.shape[0], device=input.device, dtype=output_dtype)
+    output = torch.empty(qinput_2d.shape[0], weight.shape[0], device=input.device, dtype=input.dtype)
 
     # Pass `(1, 1, gran_k)` for int-SF paths so the kernel uses the right K granularity
     # (the default `(1, 1, 128)` mismatches FP4's gran_k=32). Float-SF leaves it None.
@@ -696,7 +654,15 @@ def deepgemm_fp8_fp4_experts_forward(
     top_k_index: torch.Tensor,
     top_k_weights: torch.Tensor,
 ) -> torch.Tensor:
-    _assert_single_device(hidden_states.device, context="experts")
+    if self._deepgemm_disabled:
+        # Set at load when the model spans >1 CUDA device in this process, where DeepGEMM's
+        # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`).
+        raise RuntimeError(
+            "DeepGEMM experts selected on a model spanning multiple CUDA devices in one process; "
+            "its kernels are bound to a single CUDA context and corrupt across devices. Use "
+            "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)."
+        )
+
     _assert_sm100_scales_are_ue8m0(self.down_proj_scale_inv)
 
     if self.activation_scheme == "static":

src/transformers/integrations/finegrained_fp8.py

@@ -183,7 +183,6 @@ def finegrained_fp8_linear(
     block_size: list[int] | None = None,
     bias: torch.Tensor | None = None,
     activation_scale: torch.Tensor | None = None,
-    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Triton FP8/FP4 linear: fused act-quant + matmul, then optional bias add.
 
@@ -197,7 +196,7 @@ def finegrained_fp8_linear(
         weight,
         weight_scale_inv,
         block_size,
-        output_dtype,
+        input.dtype,
         activation_scale=activation_scale,
     )
     if bias is not None:
@@ -212,10 +211,12 @@ def fp8_linear(
     block_size: list[int] | None = None,
     bias: torch.Tensor | None = None,
     activation_scale: torch.Tensor | None = None,
-    output_dtype: torch.dtype | None = None,
+    allow_deepgemm: bool = True,
 ) -> torch.Tensor:
     """End-to-end FP8/FP4 linear used by `FP8Linear` and the eager `FP8Experts` loop.
 
+    The output dtype always follows ``input``.
+
     Dispatch order — both backends handle FP8 and FP4 weights with fp32 or UE8M0 scales:
       1. DeepGEMM (`deepgemm_fp8_fp4_linear`) — 3-6× faster on the shapes it supports.
          Preferred for FP4, UE8M0 SFs, and 128×128 block FP8.
@@ -233,26 +234,25 @@ def fp8_linear(
         bias: optional bias added to the matmul output.
         activation_scale: pass a per-tensor scalar to use static activation quant; leave `None`
             for dynamic (per-token) quant.
-        output_dtype: desired output dtype.
+        allow_deepgemm: set ``False`` to force the Triton fallback for this call. Used when the
+            model spans multiple CUDA devices in one process — DeepGEMM's cached kernels are bound
+            to a single CUDA context and produce garbage across devices (see the multi-device guard
+            in ``quantizer_finegrained_fp8.py``).
     """
     # DeepGEMM is CUDA-only, dynamic-only, SM90+ only, FP4/FP8-block-128-only.
     # ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces the Triton fallback for this single
     # dispatcher (the experts ``"deepgemm"`` impl is unaffected — use ``set_experts_implementation``
     # for that). Used by the FP8 MoE batched_mm / grouped_mm paths to avoid a still-unexplained
     # DeepGEMM-vs-Triton interaction that degrades end-to-end generation on B200 (per-row kernel
     # outputs still measure bit-perfect, but final tokens drift; not reproducible with the
-    # DeepGEMM linear off). Also temporarily skipped under ``torch.compile`` — DeepGEMM's
-    # per-token cast calls ``pack_ue8m0_to_int`` which has data-dependent bit-twiddling that
-    # dynamo can't guard. TODO: remove the ``is_torchdynamo_compiling`` gate once the upstream
-    # ``pack_ue8m0_to_int`` is rewritten to be FakeTensor-friendly; the Triton fallback is
-    # dynamo-friendly today via its ``@triton_op`` registration.
+    # DeepGEMM linear off).
     deepgemm_preferred = (
-        activation_scale is None
+        allow_deepgemm
+        and activation_scale is None
         and weight.device.type == "cuda"
         and torch.cuda.get_device_properties().major >= 9
         and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128))
         and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1"
-        and not is_torchdynamo_compiling()
     )
 
     if deepgemm_preferred:
@@ -262,7 +262,6 @@ def fp8_linear(
                 weight,
                 weight_scale_inv,
                 block_size=block_size,
-                output_dtype=output_dtype,
                 activation_scale=activation_scale,
                 bias=bias,
             )
@@ -274,10 +273,14 @@ def fp8_linear(
                 "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for FP8 linear entirely."
             )
 
-    return finegrained_fp8_linear(input, weight, weight_scale_inv, block_size, bias, activation_scale, output_dtype)
+    return finegrained_fp8_linear(input, weight, weight_scale_inv, block_size, bias, activation_scale)
 
 
 class FP8Linear(nn.Linear):
+    # Set True at load when the model spans >1 CUDA device in one process; DeepGEMM's
+    # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`).
+    _deepgemm_disabled = False
+
     def __init__(
         self,
         in_features: int,
@@ -329,8 +332,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
             scale_inv,
             block_size=self.block_size,
             activation_scale=self.activation_scale,
-            output_dtype=input.dtype,
             bias=self.bias,
+            allow_deepgemm=not self._deepgemm_disabled,
         )
 
 
@@ -574,6 +577,10 @@ def fp8_grouped_mm_experts_forward(
 
 
 class FP8Experts(nn.Module):
+    # Set True at load when the model spans >1 CUDA device in one process; DeepGEMM's
+    # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`).
+    _deepgemm_disabled = False
+
     # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan.
     # The plan dicts store `{module-path-pattern: parallel-layer-kind}`; this maps an
     # old kind to a new kind, and the quantizer rewrites every plan VALUE that matches.
@@ -729,7 +736,7 @@ def linear(
             weight_scale_inv,
             self.block_size,
             activation_scale=activation_scale,
-            output_dtype=input.dtype,
+            allow_deepgemm=not self._deepgemm_disabled,
         )
 
 

src/transformers/quantizers/quantizer_finegrained_fp8.py

@@ -153,6 +153,34 @@ def _process_model_after_weight_loading(self, model, **kwargs):
                 module = model.get_submodule(module_name)
                 scale = getattr(module, attr)
                 setattr(module, attr, torch.nn.Parameter(scale.data.to(ue8m0), requires_grad=False))
+
+        # DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` handle
+        # to the CUDA context live at load time, so one process driving it across >1 device launches
+        # against the wrong context and produces garbage. (The build-time fix is compiling DeepGEMM
+        # with `DG_JIT_USE_RUNTIME_API=1` for a context-free `cudaKernel_t` loader; until the wheel
+        # we ship picks that up, we avoid single-process multi-device.) When this model's FP8 weights
+        # span multiple CUDA devices (single-process `device_map="auto"`), flag every FP8 module so
+        # its linear and experts paths skip DeepGEMM entirely and run through Triton/grouped_mm. A
+        # model that fits on one device keeps DeepGEMM even when other GPUs are visible; TP/EP put one
+        # device per process, so this is a no-op there.
+        from ..integrations.finegrained_fp8 import FP8Experts, FP8Linear
+
+        fp8_modules = [m for m in model.modules() if isinstance(m, (FP8Linear, FP8Experts))]
+        cuda_devices = set()
+        for m in fp8_modules:
+            param = next(m.parameters(), None)
+            if param is not None and param.device.type == "cuda":
+                cuda_devices.add(param.device.index)
+        if len(cuda_devices) > 1:
+            for m in fp8_modules:
+                m._deepgemm_disabled = True
+            logger.warning_once(
+                "This FP8 model spans multiple CUDA devices in one process; routing its FP8 linear "
+                "and experts layers through Triton/grouped_mm instead of DeepGEMM (DeepGEMM's cached "
+                "kernels are bound to a single CUDA context and corrupt across devices). Run "
+                "tensor/expert parallel (one device per process) to use the faster DeepGEMM path."
+            )
+
         return model
 
     def update_tp_plan(self, config):