Updated ci/cd

Signed-off-by: romit <[email protected]>

Author: romitjain

tests/artifacts/language_models/__init__.py

@@ -20,4 +20,4 @@
 ### Constants used for model path
 PREDEFINED_MODEL_PATH = os.path.join(os.path.dirname(__file__))
 MAYKEYE_TINY_LLAMA_CACHED = os.path.join(PREDEFINED_MODEL_PATH, "maykeye-tinyllama-v0")
-TINYMIXTRAL_MOE = "Isotonic/TinyMixtral-4x248M-MoE"
+TINYMIXTRAL_MOE = "Isotonic/TinyMixtral-4x248M-MoE"

tests/artifacts/predefined_data_configs/tokenize_and_apply_input_masking_streaming.yaml

@@ -8,7 +8,7 @@ datasets:
     data_handlers:
       - name: tokenize_and_apply_input_masking
         arguments:
-          remove_columns: all
+          # remove_columns: all
           batched: false
           fn_kwargs:
             input_column_name: input

tests/test_sft_trainer.py

@@ -1766,9 +1766,7 @@ def test_run_moe_ft_and_inference_ep1_kernels(dataset_path, ep_degree):
         sft_trainer.train(
             model_args, data_args, train_args, fast_moe_config=fast_moe_config
         )
-        _test_run_inference(
-            checkpoint_path=_get_hf_converted_path(tempdir)
-        )
+        _test_run_inference(checkpoint_path=_get_hf_converted_path(tempdir))
 
 
 @pytest.mark.skipif(
@@ -1801,27 +1799,17 @@ def test_run_moe_lora_and_inference(dataset_path, target_modules, ep_degree):
         lora_args.target_modules = target_modules
         fast_moe_config = FastMoeConfig(fast_moe=FastMoe(ep_degree=ep_degree))
 
-        if target_modules == "all-linear":
-            with pytest.raises(ValueError):
-                sft_trainer.train(
-                    model_args,
-                    data_args,
-                    train_args,
-                    lora_args,
-                    fast_moe_config=fast_moe_config,
-                )
-        else:
-            sft_trainer.train(
-                model_args,
-                data_args,
-                train_args,
-                lora_args,
-                fast_moe_config=fast_moe_config,
-            )
-            _test_run_inference(
-                checkpoint_path=_get_checkpoint_path(tempdir),
-                base_model_name_or_path=TINYMIXTRAL_MOE,
-            )
+        sft_trainer.train(
+            model_args,
+            data_args,
+            train_args,
+            lora_args,
+            fast_moe_config=fast_moe_config,
+        )
+        _test_run_inference(
+            checkpoint_path=_get_checkpoint_path(tempdir),
+            base_model_name_or_path=TINYMIXTRAL_MOE,
+        )
 
 
 @pytest.mark.skipif(
@@ -1936,7 +1924,8 @@ def _get_hf_converted_path(dir_path):
     hf_converted_dir = [
         d
         for d in os.listdir(final_checkpoint_path)
-        if os.path.isdir(os.path.join(final_checkpoint_path, d)) and re.match(r"^safetensors-\d+$", d)
+        if os.path.isdir(os.path.join(final_checkpoint_path, d))
+        and re.match(r"^safetensors-\d+$", d)
     ]
     hf_converted_dir.sort(key=lambda name: int(name.split("-")[-1]))
 

tuning/config/acceleration_configs/acceleration_framework_config.py

@@ -240,6 +240,9 @@ def get_framework(self):
             try:
                 with NamedTemporaryFile("w") as f:
                     self.to_yaml(f.name)
+
+                    AccelerationFramework.active_plugins = []
+                    AccelerationFramework.plugins_require_custom_loading = []
                     return AccelerationFramework(f.name)
             except ValueError as e:
                 (msg,) = e.args