added test

Author: jcaip

test/sparsity/test_fast_sparse_training.py

@@ -68,6 +68,48 @@ def test_runtime_weight_sparsification(self):
         for name, mod in model_c.named_modules():
             assert not isinstance(mod, SemiSparseLinear)
 
+    @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "pytorch 2.4+ feature")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    def test_runtime_weight_sparsification_compile(self):
+        # need this import inside to not break 2.2 tests
+        from torch.sparse import SparseSemiStructuredTensorCUSPARSELT
+        input = torch.rand((128, 128)).half().cuda()
+        grad = torch.rand((128, 128)).half().cuda()
+        model = TestModel().half().cuda()
+        model_c = copy.deepcopy(model)
+
+        for name, mod in model.named_modules():
+            if isinstance(mod, torch.nn.Linear):
+                sparse = SparseSemiStructuredTensorCUSPARSELT.prune_dense_static_sort(mod.weight.detach()).to_dense()
+                mod.weight = nn.Parameter(sparse)
+
+        model = torch.compile(model, fullgraph=True)
+        dense_result = model(input)
+
+        # map from fqn to replacement linear module
+        sparse_config = {
+            "linear1": SemiSparseLinear,
+            "linear2": SemiSparseLinear,
+        }
+
+        swap_linear_with_semi_sparse_linear(model_c, sparse_config)
+        model_c = torch.compile(model_c, fullgraph=True)
+        sparse_result = model_c(input)
+
+        assert torch.allclose(dense_result, sparse_result, rtol=1e-1, atol=1e-1)
+
+        dense_result.backward(grad)
+        sparse_result.backward(grad)
+
+        # check grad
+        assert torch.allclose(model.linear1.weight.grad, model_c.linear1.weight.grad, rtol=1e-1, atol=1e-1)
+        assert torch.allclose(model.linear2.weight.grad, model_c.linear2.weight.grad, rtol=1e-1, atol=1e-1)
+
+        # check that swap back works
+        swap_semi_sparse_linear_with_linear(model_c)
+        for name, mod in model_c.named_modules():
+            assert not isinstance(mod, SemiSparseLinear)
+
 
 if __name__ == "__main__":
     unittest.main()

torchao/sparsity/training/README.md

@@ -1,6 +1,9 @@
 # Accelerated Sparse Training
 
-This folder contains an implementation of accelerated sparse training, utilizing the runtime semi-structured (2:4) sparsification [kernels]() present in core.
+This folder contains an implementation of accelerated sparse training.
+<!--For more information about our API and how it works, please see our blog post. (Will add link when its public)-->
+
+Special thanks to @danthe3rd for writing the runtime semi-structured (2:4) sparsification [kernels](https://github.com/pytorch/pytorch/pull/122350) in core.
 
 ### Quickstart
 
@@ -55,5 +58,3 @@ For VIT-L MLP shapes on a NVIDIA A100 we see the following results:
 
 Times are in microseconds (us).
 ```
-
-For more information about our API and how it works, please see our blog post.