fix tests

mosaicml · Apr 12, 2024 · 0a5026e · 0a5026e
1 parent b0d2849
commit 0a5026e
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/models/layers/test_dmoe.py b/tests/models/layers/test_dmoe.py
@@ -239,6 +239,10 @@ def test_fwd_equal_dmoe(seqlen: int, precision: str, mlp_type: str):
 
     torch_dmoe_config = copy.deepcopy(mb_dmoe_config)
     torch_dmoe_config.ffn_config['ffn_type'] = 'torch_dmoe'
+    del torch_dmoe_config.ffn_config['moe_world_size']
+    del torch_dmoe_config.ffn_config['fc_type']
+    del torch_dmoe_config.ffn_config['moe_loss_weight']
+    del torch_dmoe_config.ffn_config['return_bias']
 
     mb_dmoe_model = MPTForCausalLM(mb_dmoe_config).to(device=device,
                                                       dtype=dtype)