feat: add support for mamba cp

Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>

Author: kmehant

tests/test_sft_trainer.py

@@ -412,6 +412,7 @@ def test_parse_arguments(job_config):
         _,
         _,
         _,
+        _,
     ) = sft_trainer.parse_arguments(parser, job_config_copy)
     assert str(model_args.torch_dtype) == "torch.bfloat16"
     assert data_args.dataset_text_field == "output"
@@ -438,6 +439,7 @@ def test_parse_arguments_defaults(job_config):
         _,
         _,
         _,
+        _,
     ) = sft_trainer.parse_arguments(parser, job_config_defaults)
     assert str(model_args.torch_dtype) == "torch.bfloat16"
     assert model_args.use_flash_attn is False
@@ -461,6 +463,7 @@ def test_parse_arguments_peft_method(job_config):
         _,
         _,
         _,
+        _,
     ) = sft_trainer.parse_arguments(parser, job_config_pt)
 
     assert isinstance(tune_config, peft_config.PromptTuningConfig)
@@ -480,6 +483,7 @@ def test_parse_arguments_peft_method(job_config):
         _,
         _,
         _,
+        _,
     ) = sft_trainer.parse_arguments(parser, job_config_lora)
     assert isinstance(tune_config, peft_config.LoraConfig)
     assert not tune_config.target_modules