Updated pytorch and disabled sparse tests

.github/workflows/pr-gpu.yaml

@@ -21,14 +21,14 @@ jobs:
       fail-fast: false
       matrix:
         include:
-        - name: "python3.11-pytorch2.5.1-gpus1"
+        - name: "python3.11-pytorch2.6.0-gpus1"
           gpu_num: 1
           python_version: 3.11
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
-        - name: "python3.11-pytorch2.5.1-gpus2"
+          container: mosaicml/pytorch:2.6.0_cu124-python3.11-ubuntu22.04
+        - name: "python3.11-pytorch2.6.0-gpus2"
           gpu_num: 2
           python_version: 3.11
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.6.0_cu124-python3.11-ubuntu22.04
     steps:
     - name: Run PR GPU tests
       uses: mosaicml/ci-testing/.github/actions/pytest-gpu@v0.1.2

.pre-commit-config.yaml

@@ -32,7 +32,7 @@ repos:
     additional_dependencies:
     - toml
 - repo: https://github.com/hadialqattan/pycln
-  rev: v2.1.2
+  rev: v2.5.0
   hooks:
   - id: pycln
     args: [. --all]

megablocks/layers/arguments.py

@@ -73,6 +73,18 @@ class Arguments:
     moe_zloss_in_fp32: bool = False
 
     def __post_init__(self):
+        # Sparse MLP is not supported with triton >=3.2.0
+        # TODO: Remove this once sparse is supported with triton >=3.2.0
+        if self.__getattribute__('mlp_impl') == 'sparse':
+            try:
+                import triton
+                if triton.__version__ >= '3.2.0':
+                    raise ValueError(
+                        'Sparse MLP is not supported with triton >=3.2.0. Please use mlp_impl="grouped" instead.',
+                    )
+            except ImportError:
+                raise ImportError('Triton is required for sparse MLP implementation')
+
         if self.__getattribute__('mlp_impl') == 'grouped':
             grouped_gemm.assert_grouped_gemm_is_available()
 

pyproject.toml

@@ -3,7 +3,7 @@
 
 # build requirements
 [build-system]
-requires = ["setuptools < 70.0.0", "torch >= 2.5.1, < 2.5.2"]
+requires = ["setuptools < 70.0.0", "torch >= 2.6.0, < 2.6.1"]
 build-backend = "setuptools.build_meta"
 
 # Pytest

setup.py

@@ -62,8 +62,8 @@
 install_requires = [
     'numpy>=1.21.5,<2.1.0',
     'packaging>=21.3.0,<24.2',
-    'torch>=2.5.1,<2.5.2',
-    'triton>=2.1.0',
+    'torch>=2.6.0,<2.6.1',
+    'triton>=3.2.0,<3.3.0',
     'stanford-stk==0.7.1',
 ]
 

tests/layers/dmoe_test.py

@@ -53,6 +53,16 @@ def construct_moes(
     mlp_impl: str = 'sparse',
     moe_zloss_weight: float = 0,
 ):
+    # All tests are skipped if triton >=3.2.0 is installed since sparse is not supported
+    # TODO: Remove this once sparse is supported with triton >=3.2.0
+    if mlp_impl == 'sparse':
+        try:
+            import triton
+            if triton.__version__ >= '3.2.0':
+                pytest.skip('Sparse MLP is not supported with triton >=3.2.0')
+        except ImportError:
+            pass
+
     init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
     args = Arguments(
         hidden_size=hidden_size,

tests/layers/glu_test.py

@@ -23,6 +23,16 @@ def construct_dmoe_glu(
     mlp_impl: str = 'sparse',
     memory_optimized_mlp: bool = False,
 ):
+    # All tests are skipped if triton >=3.2.0 is installed since sparse is not supported
+    # TODO: Remove this once sparse is supported with triton >=3.2.0
+    if mlp_impl == 'sparse':
+        try:
+            import triton
+            if triton.__version__ >= '3.2.0':
+                pytest.skip('Sparse MLP is not supported with triton >=3.2.0')
+        except ImportError:
+            pass
+
     init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
     args = Arguments(
         hidden_size=hidden_size,

tests/layers/moe_test.py

@@ -41,6 +41,15 @@ def construct_moe(
     moe_top_k: int = 1,
     moe_zloss_weight: float = 0,
 ):
+    # All tests are skipped if triton >=3.2.0 is installed since sparse is not supported
+    # TODO: Remove this once sparse is supported with triton >=3.2.0
+    try:
+        import triton
+        if triton.__version__ >= '3.2.0':
+            pytest.skip('Sparse MLP is not supported with triton >=3.2.0')
+    except ImportError:
+        pass
+
     init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
     args = Arguments(
         hidden_size=hidden_size,