Activation checkpointing is not working on SwiGLU

[warpuv]

🐛 Bug

When using torch.utils.checkpoint with SwiGLU module an error occurs.

To Reproduce

Sample code:

import torch
from torch.utils.checkpoint import checkpoint
from xformers.ops import SwiGLU


def main():
    device = 'cuda'
    dtype = torch.bfloat16

    shape = (64, 128, 256, 128)
    in_features, hidden_features, out_features = shape[1:]

    x = torch.randn(shape[:2], dtype=dtype, device=device)
    model = SwiGLU(in_features, hidden_features, out_features, bias=True, _pack_weights=True).to(device=device, dtype=dtype)
    
    y = checkpoint(model, x, use_reentrant=False) 
    #y = model(x) # backward works fine when using this variant
    y.backward(y.detach())

if __name__ == '__main__':
    main()

Error:

Traceback (most recent call last):
  File "/workspace/src/my/fused_swiglu/proof.py", line 20, in <module>
    main()
  File "/workspace/src/my/fused_swiglu/proof.py", line 17, in main
    y.backward(y.detach())
  File "/workspace/envs_python/torch241_cu124_py310/lib/python3.10/site-packages/torch/_tensor.py", line 521, in backward
    torch.autograd.backward(
  File "/workspace/envs_python/torch241_cu124_py310/lib/python3.10/site-packages/torch/autograd/__init__.py", line 289, in backward
    _engine_run_backward(
  File "/workspace/envs_python/torch241_cu124_py310/lib/python3.10/site-packages/torch/autograd/graph.py", line 769, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/workspace/envs_python/torch241_cu124_py310/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 1120, in unpack_hook
    frame.check_recomputed_tensors_match(gid)
  File "/workspace/envs_python/torch241_cu124_py310/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 857, in check_recomputed_tensors_match
    raise CheckpointError(
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: A different number of tensors was saved during the original forward and recomputation.
Number of tensors saved during forward: 5
Number of tensors saved during recomputation: 0

Expected behavior

Backward pass with activation checkpointing should work as well as without it.

Environment

/usr/lib/python3.10/runpy.py:126: RuntimeWarning: 'torch.utils.collect_env' found in sys.modules after import of package 'torch.utils', but prior to execution of 'torch.utils.collect_env'; this may result in unpredictable behaviour
  warn(RuntimeWarning(msg))
Collecting environment information...
PyTorch version: 2.4.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-35-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.6.68
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: Could not collect
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
Address sizes:                   43 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          112
On-line CPU(s) list:             0-111
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC Processor
CPU family:                      23
Model:                           1
Thread(s) per core:              2
Core(s) per socket:              56
Socket(s):                       1
Stepping:                        2
BogoMIPS:                        3999.99
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
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       1.8 MiB (56 instances)
L1i cache:                       3.5 MiB (56 instances)
L2 cache:                        28 MiB (56 instances)
L3 cache:                        8 MiB (1 instance)
NUMA node(s):                    1
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

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] torch==2.4.1+cu124
[pip3] torchaudio==2.4.1+cu124
[pip3] torchvision==0.19.1+cu124
[pip3] triton==3.0.0
[conda] Could not collect

• PyTorch Version: 2.4.1+cu124
• OS: Ubuntu 22.04.4 LTS (x86_64)
• How you installed PyTorch: pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu124
• Python version: 3.10.12
• GPU models and configuration: NVIDIA A100-SXM4