From 472d009cc1da84d712df6d62569fd83d9f25cdfe Mon Sep 17 00:00:00 2001
From: Daniel King <43149077+dakinggg@users.noreply.github.com>
Date: Thu, 27 Jun 2024 13:56:58 -0700
Subject: [PATCH] Fix 4 gpu tests (#1311)

---
 tests/a_scripts/inference/test_convert_composer_to_hf.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
index 1b2f791995..2ea9470324 100644
--- a/tests/a_scripts/inference/test_convert_composer_to_hf.py
+++ b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -1179,7 +1179,7 @@ def test_convert_and_generate_meta(
 @pytest.mark.world_size(4)
 @pytest.mark.gpu
 @pytest.mark.parametrize('num_experts', [2, 4, 8])
-@pytest.mark.parametrize('sharding_strategy', ['FULL_SHARD', 'HYBRID_SHARD'])
+@pytest.mark.parametrize('sharding_strategy', ['FULL_SHARD'])
 def test_mptmoe_huggingface_conversion_callback(
     tmp_path: pathlib.Path,
     num_experts: int,
@@ -1281,7 +1281,6 @@ def test_mptmoe_huggingface_conversion_callback(
         make_tiny_ft_dataset(path=tiny_dataset_path, size=dataset_size)
 
     dataloader_cfg = {
-        'name': 'finetuning',
         'dataset': {
             'hf_name': tiny_dataset_folder_path,
             'split': 'train',
@@ -1350,7 +1349,6 @@ def test_mptmoe_huggingface_conversion_callback(
         save_weights_only=True,
     )
     trainer.fit()
-    #self.state.outputs = self.state.model(self.state.batch)
     batch = trainer.state.batch
     model_output_logits = trainer.state.model(batch).logits
 
@@ -1428,7 +1426,6 @@ def test_mptmoe_huggingface_conversion_callback(
             loaded_model_logits = loaded_model(
                 input_ids=batch.get('input_ids', None),
                 attention_mask=batch.get('attention_mask', None),
-                prefix_mask=batch.get('bidirectional_mask', None),
                 sequence_id=batch.get('sequence_id', None),
                 inputs_embeds=batch.get('inputs_embeds', None),
             ).logits