make test smaller

mosaicml · Mar 1, 2023 · 2e2d298 · 2e2d298
1 parent 8d2280a
commit 2e2d298
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/examples/llm/tests/test_flash_triton_torch.py b/examples/llm/tests/test_flash_triton_torch.py
@@ -6,16 +6,16 @@
 from composer.utils import reproducibility
 from omegaconf import OmegaConf as om
 
-from examples.llm.src.models.layers.attention import (  # type: ignore
-    FlashCausalAttention, TorchCausalAttention, TritonFlashCausalAttention)
-
 
 def allclose_helper(t0, t1, rtol=1e-2, atol=1e-2):
     return torch.allclose(t0, t1, rtol=rtol, atol=atol)
 
 
-# @pytest.mark.gpu
+@pytest.mark.gpu
 def test_flash_torch(device='cuda'):
+    from examples.llm.src.models.layers.attention import (  # type: ignore
+        FlashCausalAttention, TorchCausalAttention)
+
     reproducibility.seed_all(7)
 
     cfg = om.create({
@@ -73,14 +73,17 @@ def gen_tca_mask():
     assert allclose_helper(x0.grad, x1.grad)
 
 
-# @pytest.mark.gpu
+@pytest.mark.gpu
 @pytest.mark.parametrize('attn_clip_qkv,attn_qk_ln', [
     (False, False),
     (False, True),
     (True, False),
     (True, True),
 ])
 def test_flash_triton(attn_clip_qkv, attn_qk_ln, device='cuda'):
+    from examples.llm.src.models.layers.attention import (  # type: ignore
+        FlashCausalAttention, TritonFlashCausalAttention)
+
     reproducibility.seed_all(7)
 
     cfg = om.create({

diff --git a/examples/llm/tests/test_model.py b/examples/llm/tests/test_model.py
@@ -186,7 +186,7 @@ def test_full_forward_and_backward_gpt2_small(prefixlm, batch_size=2):
 
     device = 'cpu'
     neo_cfg.device = device
-    neo_cfg.max_seq_len = 1024
+    neo_cfg.max_seq_len = 256
 
     if prefixlm:
         neo_cfg.model.name = 'hf_prefix_lm'