diff --git a/tests/model/test_configs/test_sparse.yml b/tests/model/test_configs/test_sparse.yml new file mode 100644 index 000000000..4b350d177 --- /dev/null +++ b/tests/model/test_configs/test_sparse.yml @@ -0,0 +1,15 @@ +# Add this to your config for sparse attention every other layer +{ + "attention_config": [[["local"], "all"]], + + # sparsity config: + # (these are the defaults for local sliding window sparsity, training will work without this here, but it's left in for + # illustrative purposes) + # see https://www.deepspeed.ai/tutorials/sparse-attention/#how-to-config-sparsity-structures for + # more detailed config instructions and available parameters + + "sparsity_config": { + "block": 16, # block size + "num_local_blocks": 32, + } +} \ No newline at end of file diff --git a/tests/model/test_model_instantiation.py b/tests/model/test_model_instantiation.py index 35ddcb7cd..25119859d 100644 --- a/tests/model/test_model_instantiation.py +++ b/tests/model/test_model_instantiation.py @@ -60,8 +60,10 @@ def test_model_instantiation_small_test(self): def test_model_instantiation_medium_test(self): self.run_instantiation_test(get_test_configs_with_path(["test_local_setup.yml", "test_medium.yml"]), DeepSpeedEngine) + def test_model_instantiation_small_sparse_test(self): + self.run_instantiation_test(get_test_configs_with_path(["test_local_setup.yml", "test_small.yml", "test_sparse.yml"]), DeepSpeedEngine) + if __name__ == "__main__": suite = unittest.TestSuite() - suite.addTest(TestModelInstantiation("test_model_instantiation_small_test")) - suite.addTest(TestModelInstantiation("test_model_instantiation_medium_test")) + suite.addTest(TestModelInstantiation("test_model_instantiation_small_sparse_test")) unittest.TextTestRunner(failfast=True).run(suite) \ No newline at end of file