torch dmoe tests gpu oom (#1216)

mosaicml · May 16, 2024 · 3a15082 · 3a15082
1 parent 38ae65b
commit 3a15082
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/models/layers/test_dmoe.py b/tests/models/layers/test_dmoe.py
@@ -63,10 +63,10 @@ def _get_torch_dtype(fp16: bool, bf16: bool) -> Optional[torch.dtype]:
 )
 @pytest.mark.gpu
 @pytest.mark.world_size(2)
-@pytest.mark.parametrize('moe_num_experts', [1, 2, 8])
+@pytest.mark.parametrize('moe_num_experts', [1, 8])
 @pytest.mark.parametrize('mlp_type', ['glu', 'mlp'])
 @pytest.mark.parametrize('moe_world_size', [1, 2])
-@pytest.mark.parametrize('moe_normalize_expert_weights', [1, 2.0])
+@pytest.mark.parametrize('moe_normalize_expert_weights', [1.0])
 @pytest.mark.parametrize('two_d_input', [True, False])
 def test_dmoe(
     moe_num_experts: int,