fallback embedding bwd on xpu only

jianyizh · jianyizh · commit 6bd1e220cbcb · 2025-04-18T03:19:56.000Z
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
@@ -1236,6 +1236,8 @@ def embedding_dense_backward(
     padding_idx: int,
     scale_grad_by_freq: bool,
 ):
+    if grad_output.is_xpu:
+        return NotImplemented
     computation_dtype, result_dtype = utils.elementwise_dtypes(
         grad_output, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
     )
@@ -1474,7 +1476,7 @@ def _addmm_activation(
 ):
     out = addmm(self, mat1, mat2, beta, alpha)
     if use_gelu:
-        if self.is_cuda:
+        if self.is_cuda or self.is_xpu:
             return aten.gelu(out, approximate="tanh")
         else:
             return aten.gelu(out)
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -2814,6 +2814,7 @@ def is_aligned(x):
 
 # index_reduce requires fallback when use_scatter_fallback(...) returns True
 make_fallback(aten.index_reduce)
+make_fallback(aten.embedding_dense_backward)
 
 
 # Register with type_promotion_kind None.
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
@@ -7101,6 +7101,18 @@ def _check_for_unsupported_isin_dtype(dtype):
     )
 
 
+@register_meta(aten.embedding_dense_backward)
+def meta_embedding_dense_backward(
+    grad_output,
+    indices,
+    num_weights,
+    padding_idx,
+    scale_grad_by_freq,
+):
+    grad_weight = grad_output.new_empty((num_weights, grad_output.size(-1)))
+    return grad_weight
+
+
 @register_meta(aten._embedding_bag_backward)
 def meta_embedding_bag_backward(
     grad,
diff --git a/torch/_prims_common/wrappers.py b/torch/_prims_common/wrappers.py
@@ -307,6 +307,8 @@ def maybe_check_copy_devices(out):
                 result = fn(*args, is_out=(out is not None), **kwargs)  # type: ignore[arg-type]
             else:
                 result = fn(*args, **kwargs)
+            if result is NotImplemented:
+                return NotImplemented
             assert (
                 (isinstance(result, TensorLike) and is_tensor)
                 or (
