testing revisions, circleci fixes, minor changes

.circleci/config.yml

@@ -312,7 +312,7 @@ jobs:
 
 
   gpu_tests:
-    <<: *gpu_cu102
+    <<: *gpu_cu114
 
     working_directory: ~/xformers
 
@@ -324,7 +324,7 @@ jobs:
       # Cache the venv directory that contains dependencies
       - restore_cache:
           keys:
-            - cache-key-gpu-102-{{ checksum "requirements-test.txt" }}-{{ checksum ".circleci/config.yml" }}
+            - cache-key-gpu-114-{{ checksum "requirements-test.txt" }}-{{ checksum ".circleci/config.yml" }}
 
       - <<: *setup_conda
       - <<: *install_dep
@@ -338,7 +338,7 @@ jobs:
             - ~/miniconda
             - ~/venv
 
-          key: cache-key-gpu-102-{{ checksum "requirements-test.txt"}}-{{ checksum ".circleci/config.yml"}}
+          key: cache-key-gpu-114-{{ checksum "requirements-test.txt"}}-{{ checksum ".circleci/config.yml"}}
 
       - <<: *install_repo
 
@@ -354,7 +354,7 @@ jobs:
           path: test-results
 
   gpu_experimental_tests:
-    <<: *gpu_cu102
+    <<: *gpu_cu114
 
     working_directory: ~/xformers
 
@@ -366,7 +366,7 @@ jobs:
       # Cache the venv directory that contains dependencies
       - restore_cache:
           keys:
-            - cache-key-gpu-exp-102-{{ checksum "experimental/requirements.txt" }}-{{ checksum ".circleci/config.yml" }}
+            - cache-key-gpu-exp-114-{{ checksum "experimental/requirements.txt" }}-{{ checksum ".circleci/config.yml" }}
 
       - <<: *setup_conda
       - <<: *install_dep_exp
@@ -381,7 +381,7 @@ jobs:
             - ~/miniconda
             - ~/venv
 
-          key: cache-key-gpu-exp-102-{{ checksum "experimental/requirements.txt"}}-{{ checksum ".circleci/config.yml"}}
+          key: cache-key-gpu-exp-114-{{ checksum "experimental/requirements.txt"}}-{{ checksum ".circleci/config.yml"}}
 
       - <<: *install_experimental_repo
       - <<: *run_experimental_unittests

requirements-test.txt

@@ -30,4 +30,5 @@ fairscale >= 0.4.5
 triton == 2.0.0.dev20220701
 
 # Functorch: Dependency for fused layers, optional
-git+https://github.com/pytorch/functorch@v0.2.0
+# git+https://github.com/pytorch/functorch@v0.2.0
+git+https://github.com/pytorch/functorch.git

tests/test_nvfuser.py

@@ -5,6 +5,7 @@
 
 
 import logging
+from collections import OrderedDict
 from contextlib import nullcontext
 
 import pytest
@@ -13,13 +14,14 @@
 from torch.cuda.amp.autocast_mode import autocast
 
 import xformers
-from xformers.components import Activation, ResidualNormStyle
-from xformers.components.feedforward import build_feedforward
-
-_gpu_available = torch.cuda.is_available()
 
 xformers._is_functorch_available = True
 
+from xformers.components import Activation, ResidualNormStyle  # noqa : E402
+from xformers.components.feedforward import build_feedforward  # noqa : E402
+
+_gpu_available = torch.cuda.is_available()
+
 try:
     from xformers.components.nvfuser import (
         NVFusedBiasActivationDropout,
@@ -172,5 +174,17 @@ def test_nvfused_mlp(activation: Activation, device: torch.device, p: float):
     # Check fused pattern w/ unfused default (switch happens within MLP)
     mlp.cuda()
     mlp_default.cuda()
+
+    # Load same weight parameters into both models
+    default_param_dict = OrderedDict(
+        [
+            ("mlp.2.weight", v) if k == "mlp.3.weight" else (k, v)
+            for k, v in mlp_default.state_dict().items()
+        ]
+    )
+    mlp.load_state_dict(default_param_dict)
     fused_res = mlp(inputs)
     unfused_res = mlp_default(inputs)
+
+    if p == 0.0:
+        assert torch.allclose(unfused_res, fused_res, atol=1e-6, rtol=1e-2)

xformers/components/nvfuser/bias_dropout_res.py

@@ -57,7 +57,6 @@ def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
 
         # Catch a non-cuda setup, fallback to pytorch
         if not x.is_cuda:
-            print("defaulting")
             return _fn(x, self.bias, p, residual)
 
         # AOTAutograd, NVFuser backed path