fix: code refactor

Signed-off-by: Mehant Kammakomati <[email protected]>

Author: kmehant

tests/acceleration/test_acceleration_dataclasses.py

@@ -33,7 +33,6 @@
     FastKernelsConfig,
     FusedLoraConfig,
 )
-from tuning.config.acceleration_configs.odm import ODM, ODMConfig
 from tuning.config.acceleration_configs.quantized_lora_config import (
     AutoGPTQLoraConfig,
     BNBQLoraConfig,
@@ -97,13 +96,6 @@ def test_dataclass_parse_successfully():
     )
     assert isinstance(cfg.fast_moe, FastMoe)
 
-    # 5. Specifing "--odm" will parse an ODM class
-    parser = transformers.HfArgumentParser(dataclass_types=ODMConfig)
-    (cfg,) = parser.parse_args_into_dataclasses(
-        ["--odm", "2", "1", "train_loss", "0.1", "0.2"],
-    )
-    assert isinstance(cfg.odm, ODM)
-
 
 def test_two_dataclasses_parse_successfully_together():
     """Ensure that the two dataclasses can parse arguments successfully

tests/test_sft_trainer.py

@@ -2151,6 +2151,10 @@ def test_handler(element, **kwargs):
         _validate_training(tempdir)
 
 
+@pytest.mark.skipif(
+    not is_fms_accelerate_available(plugins="odm"),
+    reason="Only runs if fms-accelerate is installed along with online-data-mixing plugin",
+)
 @pytest.mark.parametrize(
     "datafiles, datasetconfigname, reward_type",
     [
@@ -2258,12 +2262,16 @@ def test_online_data_mixing_plugin_sample_training(
         )
         assert len(output_inference) > 0
         assert (
-            "It takes 10 days for digging a trench of 100 m long, 50 m broad and 10 m deep."
+            "It takes 10 days for digging a trench of 100 m long, 50 m broad and 10 m deep. "
             "What length of trench,\n25 m broad and 15 m deep can be dug in 30 days ?"
             in output_inference
         ), f"{output_inference} does not include the prompt"
 
 
+@pytest.mark.skipif(
+    not is_fms_accelerate_available(plugins="odm"),
+    reason="Only runs if fms-accelerate is installed along with online-data-mixing plugin",
+)
 @pytest.mark.parametrize(
     "datafiles, datasetconfigname, reward_type",
     [
@@ -2333,7 +2341,7 @@ def test_online_data_mixing_plugin_sample_training_no_validation_split(
         )
         assert len(output_inference) > 0
         assert (
-            "It takes 10 days for digging a trench of 100 m long, 50 m broad and 10 m deep."
+            "It takes 10 days for digging a trench of 100 m long, 50 m broad and 10 m deep. "
             "What length of trench,\n25 m broad and 15 m deep can be dug in 30 days ?"
             in output_inference
         ), f"{output_inference} does not include the prompt"

tuning/sft_trainer.py

@@ -632,8 +632,6 @@ def parse_arguments(parser, json_config=None):
             Configuration for padding free and packing.
         FastMoeConfig
             Configuration for accelerated MoE.
-        ODMConfig
-            Configuration for online data mixing feature.
         TrackerConfigs
             Configuration for all trackers.
         dict[str, str]