Repository: facebookresearch/xformers
Issue category: Bug Report

Issue:
Test of SwiGLU is passing when it should fail
# 🐛 Bug

Test is passing when it should fail if you corrupt output of swiglu_packedw.

## To Reproduce

Steps to reproduce the behavior:

In the file `xformers/csrc/swiglu/swiglu_packedw.cpp`
```c++
class SwiGLUPackedWeights
    : public torch::autograd::Function<SwiGLUPackedWeights> {
 public:
  static at::Tensor forward(
      torch::autograd::AutogradContext* ctx,
      const at::Tensor& x,
      const at::Tensor& w1w2,
      const std::optional<at::Tensor>& b1b2,
      const at::Tensor w3,
      const std::optional<at::Tensor>& b3) {
    at::AutoDispatchBelowADInplaceOrView g;
    auto w1 = w1w2[0];
    auto w2 = w1w2[1];
    std::optional<at::Tensor> b1, b2;
    if (b1b2.has_value()) {
      b1 = b1b2.value()[0];
      b2 = b1b2.value()[1];
    }
    at::Tensor x1, x2, x4;
    std::tie(x1, x2, x4) = dual_gemm_silu_identity_mul(x, w1, b1, w2, b2);
    auto x5 = torch::nn::functional::linear(
        x4, w3, b3.has_value() ? b3.value() : at::Tensor());

    ctx->save_for_backward({x, w1w2, w3, x1, x2});
    ctx->saved_data["has_b1b2"] = b1b2.has_value();
    ctx->saved_data["has_b3"] = b3.has_value();

    return x5;
  }
```
 Try to replace `return x5;` to `return x5 * x5;` for example and run the tests. Tests will wrongly pass.


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Environment

```
Collecting environment information...
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.31.0
Libc version: glibc-2.31

Python version: 3.9.5 (default, Nov 23 2021, 15:27:38)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.13.0-35-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.6.77
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB

Nvidia driver version: 560.28.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          112
On-line CPU(s) list:             0-111
Thread(s) per core:              2
Core(s) per socket:              56
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           1
Model name:                      AMD EPYC Processor
Stepping:                        2
CPU MHz:                         1999.999
BogoMIPS:                        3999.99
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       1.8 MiB
L1i cache:                       3.5 MiB
L2 cache:                        28 MiB
L3 cache:                        112 MiB
NUMA node0 CPU(s):               0-111
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; LFENCE, IBPB conditional, STIBP conditional, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext ssbd ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 arat umip

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] flake8-copyright==0.2.4
[pip3] mypy==1.10.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==2.0.2
[pip3] pytorch_sphinx_theme==0.0.24
[pip3] torch==2.5.1
[pip3] torchaudio==2.5.1
[pip3] torchvision==0.20.1
[pip3] triton==3.1.0
[conda] Could not collect

```

- PyTorch Version: 2.5.1+cu124
- OS: Linux
- Python version: 3.9.5



Hints:

Patch:
diff --git a/xformers/ops/swiglu_op.py b/xformers/ops/swiglu_op.py
index 630335ac6c..bb27e43308 100644
--- a/xformers/ops/swiglu_op.py
+++ b/xformers/ops/swiglu_op.py
@@ -212,11 +212,11 @@ def info(self):
 def _eager_functional_swiglu(
     x: torch.Tensor,
     w1: torch.Tensor,
-    b1: torch.Tensor,
+    b1: Optional[torch.Tensor],
     w2: torch.Tensor,
-    b2: torch.Tensor,
+    b2: Optional[torch.Tensor],
     w3: torch.Tensor,
-    b3: torch.Tensor,
+    b3: Optional[torch.Tensor],
 ) -> torch.Tensor:
     x1 = F.linear(x, w1, b1)
     x2 = F.linear(x, w2, b2)


Test patch:
diff --git a/tests/test_swiglu.py b/tests/test_swiglu.py
index 6600488135..46b380f68a 100644
--- a/tests/test_swiglu.py
+++ b/tests/test_swiglu.py
@@ -7,7 +7,7 @@
 import functools
 import random
 from contextlib import nullcontext
-from typing import ContextManager, Optional, Sequence, cast
+from typing import ContextManager, Optional, Sequence, Union, cast
 
 import pytest
 import torch
@@ -95,8 +95,10 @@ def generate_test_shapes():
         (4728, 1536, 2736),
         # GPT-3 (small)
         (2048, 2048, 5632),
+        # TODO: Not enough memory for this shape in github CI.
+        # restore it after rearrange the code (fw, fw, bw, bw) -> (fw, bw, fw, bw)
         # Chinchilla
-        (2048, 8192, 22016),
+        # (2048, 8192, 22016),
     ]
     # Add some random shapes
     r = random.Random(0)
@@ -113,7 +115,11 @@ def generate_test_shapes():
 _dtypes = [torch.float16]
 if _is_sm80:
     _dtypes += [torch.bfloat16]
-_ops: Sequence[xsw.SwiGLUOp] = [xsw.SwiGLUFusedOp, xsw.SwiGLUPackedFusedOp]
+_ops: Sequence[Union[xsw.SwiGLUOp, None]] = [
+    xsw.SwiGLUFusedOp,
+    xsw.SwiGLUPackedFusedOp,
+    None,
+]
 
 FORWARD_ATOL = {torch.float: 2e-6, torch.half: 1e-2, torch.bfloat16: 1e-2}
 FORWARD_RTOL = {torch.float: 1e-5, torch.half: 4e-3, torch.bfloat16: 4e-3}
@@ -125,7 +131,7 @@ def generate_test_shapes():
 }
 BACKWARD_RTOL = {
     torch.float: 2e-3,
-    torch.half: 1e-2,
+    torch.half: 3e-2,
     torch.bfloat16: 4e-2,
 }
 
@@ -138,7 +144,9 @@ def create_module_cached(**kwargs) -> xsw.SwiGLU:
 @disable_tf32
 @disable_on_rocm
 @pytest.mark.parametrize("autocast", [False, True], ids=["regular", "autocast"])
-@pytest.mark.parametrize("op", _ops, ids=[x.NAME for x in _ops])
+@pytest.mark.parametrize(
+    "op", _ops, ids=[x.NAME if x is not None else "auto_selected_op" for x in _ops]
+)
 @pytest.mark.parametrize("dtype", _dtypes, ids=[str(x) for x in _dtypes])
 @pytest.mark.parametrize("device", _devices)
 @pytest.mark.parametrize("bias", [False, True], ids=["nobias", "bias"])
@@ -159,7 +167,7 @@ def test_forward_backward(
 ):
     torch.manual_seed(shape[0] * shape[1] * shape[2])
 
-    if not op.supports(
+    if op is not None and not op.supports(
         xsw.SwiGLUOpDispatch(
             device=device,
             dtype=dtype,
@@ -170,6 +178,10 @@ def test_forward_backward(
     ):
         pytest.skip("Not supported by operator")
 
+    if op is not None:
+        if pack_weights and not op.PACKED_WEIGHTS:
+            pytest.skip("Not supported combination when module.op is set manually")
+
     inp_model_dtype = torch.float if autocast else dtype
     x = torch.randn(shape[:2], device=device, dtype=inp_model_dtype)
 
@@ -200,8 +212,10 @@ def test_forward_backward(
         torch.autocast("cuda", dtype=dtype) if autocast else nullcontext(),
     )
     with cm:
-        ref = module(x)
-        out = xsw.swiglu(x, *module._ordered_params(), op=op)
+        ref = xsw._eager_functional_swiglu(x, *module._ordered_params())
+
+        module.op = op
+        out = module(x)
 
     if ref_f32 is None:
         ref_f32 = ref