plugin rename

Signed-off-by: 1000850000 user <aaron.chew1@ibm.com>

tests/acceleration/test_acceleration_dataclasses.py

@@ -27,8 +27,8 @@
     FastKernelsConfig,
     FusedLoraConfig,
 )
-from tuning.config.acceleration_configs.instruct_lab_config import (
-    InstructLabConfig,
+from tuning.config.acceleration_configs.attention_and_distributed_packing import (
+    AttentionAndDistributedPackingConfig,
     PaddingFree,
 )
 from tuning.config.acceleration_configs.quantized_lora_config import (
@@ -70,7 +70,7 @@ def test_dataclass_parse_successfully():
     assert isinstance(cfg.bnb_qlora, BNBQLoraConfig)
 
     # 3. Specifing "--padding_free" will parse a PaddingFree class
-    parser = transformers.HfArgumentParser(dataclass_types=InstructLabConfig)
+    parser = transformers.HfArgumentParser(dataclass_types=AttentionAndDistributedPackingConfig)
     (cfg,) = parser.parse_args_into_dataclasses(
         ["--padding_free", "huggingface"],
     )

tests/acceleration/test_acceleration_framework.py

@@ -42,8 +42,8 @@
     FastKernelsConfig,
     FusedLoraConfig,
 )
-from tuning.config.acceleration_configs.instruct_lab_config import (
-    InstructLabConfig,
+from tuning.config.acceleration_configs.attention_and_distributed_packing import (
+    AttentionAndDistributedPackingConfig,
     PaddingFree,
 )
 from tuning.config.acceleration_configs.quantized_lora_config import (
@@ -467,8 +467,9 @@ def test_framework_intialized_properly_foak():
 
 
 @pytest.mark.skipif(
-    not is_fms_accelerate_available(plugins="ilab"),
-    reason="Only runs if fms-accelerate is installed along with instruct-lab plugin",
+    not is_fms_accelerate_available(plugins="aadp"),
+    reason="Only runs if fms-accelerate is installed along with \
+        attention_and_distributed_packing plugin",
 )
 def test_framework_initialize_and_trains_with_ilab():
     """
@@ -489,7 +490,7 @@ def test_framework_initialize_and_trains_with_ilab():
         data_args.dataset_text_field = None
 
         # initialize a config
-        instruct_lab_config = InstructLabConfig(
+        attention_and_distributed_packing_config = AttentionAndDistributedPackingConfig(
             padding_free=PaddingFree(method="huggingface")
         )
 
@@ -512,7 +513,8 @@ def test_framework_initialize_and_trains_with_ilab():
                     model_args,
                     data_args,
                     train_args,
-                    instruct_lab_config=instruct_lab_config,
+                    attention_and_distributed_packing_config=\
+                        attention_and_distributed_packing_config,
                 )
 
         # spy inside the train to ensure that the ilab plugin is called
@@ -522,8 +524,9 @@ def test_framework_initialize_and_trains_with_ilab():
 
 
 @pytest.mark.skipif(
-    not is_fms_accelerate_available(plugins="ilab"),
-    reason="Only runs if fms-accelerate is installed along with instruct-lab plugin",
+    not is_fms_accelerate_available(plugins="aadp"),
+    reason="Only runs if fms-accelerate is installed along with \
+        attention_and_distributed_packing plugin",
 )
 def test_padding_free_plugin_raises_error_with_untokenized_dataset():
     """
@@ -547,7 +550,7 @@ def test_padding_free_plugin_raises_error_with_untokenized_dataset():
         data_args.dataset_text_field = "output"
 
         # initialize a config
-        instruct_lab_config = InstructLabConfig(
+        attention_and_distributed_packing_config = AttentionAndDistributedPackingConfig(
             padding_free=PaddingFree(method="huggingface")
         )
 
@@ -570,5 +573,27 @@ def test_padding_free_plugin_raises_error_with_untokenized_dataset():
                         model_args,
                         data_args,
                         train_args,
-                        instruct_lab_config=instruct_lab_config,
+                        attention_and_distributed_packing_config=\
+                        attention_and_distributed_packing_config,
                     )
+
+def test_error_raised_with_paddingfree_and_flash_attn_disabled():
+    """Ensure error raised when padding-free is not used with flash attention"""
+    with pytest.raises(
+        ValueError,
+        match="`--padding_free` argument was called without enabling flash attention, \
+            ensure `use_flash_attn = True` to use padding-free flash attention",
+    ):
+        attention_and_distributed_packing_config = \
+            AttentionAndDistributedPackingConfig(
+                padding_free=PaddingFree(method="huggingface")
+            )
+        model_args = copy.deepcopy(MODEL_ARGS)
+        model_args.use_flash_attn = False
+        sft_trainer.train(
+            model_args,
+            DATA_ARGS,
+            TRAIN_ARGS,
+            attention_and_distributed_packing_config=\
+                attention_and_distributed_packing_config
+        )

tuning/config/acceleration_configs/__init__.py

@@ -15,5 +15,5 @@
 # Local
 from .acceleration_framework_config import AccelerationFrameworkConfig
 from .fused_ops_and_kernels import FusedOpsAndKernelsConfig
-from .instruct_lab_config import InstructLabConfig
+from .attention_and_distributed_packing import AttentionAndDistributedPackingConfig
 from .quantized_lora_config import QuantizedLoraConfig

tuning/config/acceleration_configs/acceleration_framework_config.py

@@ -22,7 +22,7 @@
 
 # Local
 from .fused_ops_and_kernels import FastKernelsConfig, FusedLoraConfig
-from .instruct_lab_config import PaddingFree
+from .attention_and_distributed_packing import PaddingFree
 from .quantized_lora_config import AutoGPTQLoraConfig, BNBQLoraConfig
 from tuning.utils.import_utils import is_fms_accelerate_available
 
@@ -104,7 +104,7 @@ class AccelerationFrameworkConfig:
         ConfigAnnotation(
             path="training.attention",
             experimental=True,
-            required_packages=["ilab"],
+            required_packages=["aadp"],
         ),
     ] = None
 

tuning/config/acceleration_configs/instruct_lab_config.py → tuning/config/acceleration_configs/attention_and_distributed_packing.py

@@ -17,7 +17,7 @@ def __post_init__(self):
 
 
 @dataclass
-class InstructLabConfig:
+class AttentionAndDistributedPackingConfig:
 
     padding_free: PaddingFree = None
 

tuning/sft_trainer.py

@@ -45,7 +45,7 @@
 from tuning.config.acceleration_configs import (
     AccelerationFrameworkConfig,
     FusedOpsAndKernelsConfig,
-    InstructLabConfig,
+    AttentionAndDistributedPackingConfig,
     QuantizedLoraConfig,
 )
 from tuning.config.tracker_configs import (
@@ -87,7 +87,7 @@ def train(
     exp_metadata: Optional[Dict] = None,
     quantized_lora_config: Optional[QuantizedLoraConfig] = None,
     fusedops_kernels_config: Optional[FusedOpsAndKernelsConfig] = None,
-    instruct_lab_config: Optional[InstructLabConfig] = None,
+    attention_and_distributed_packing_config: Optional[AttentionAndDistributedPackingConfig] = None,
 ):
     """Call the SFTTrainer
 
@@ -115,7 +115,7 @@ def train(
         fusedops_kernels_config: tuning.config.acceleration_configs.FusedOpsAndKernelsConfig \
             Should be used in combination with quantized_lora_config. Also currently 
             fused_lora and fast_kernels must used together (may change in future). \
-        instruct_lab_config: Used for padding free and multipack.
+        attention_and_distributed_packing_config: Used for padding-free attention and multipack.
     """
 
     train_args, logger = set_log_level(train_args, "sft_trainer_train")
@@ -130,6 +130,12 @@ def train(
     ):
         raise ValueError("gradient_accumulation_steps has to be an integer >= 1")
 
+    if (attention_and_distributed_packing_config.padding_free is not None and
+        model_args.use_flash_attn is False
+    ):
+        raise ValueError("`--padding_free` argument was called without enabling \
+            flash attention, ensure `use_flash_attn = True` to use padding-free flash attention")
+
     task_type = "CAUSAL_LM"
     additional_metrics = {}
 
@@ -181,7 +187,7 @@ def train(
             trainer_callbacks.append(cb)
 
     framework = AccelerationFrameworkConfig.from_dataclasses(
-        quantized_lora_config, fusedops_kernels_config, instruct_lab_config
+        quantized_lora_config, fusedops_kernels_config, attention_and_distributed_packing_config
     ).get_framework()
 
     model_loader = AutoModelForCausalLM.from_pretrained
@@ -439,7 +445,7 @@ def get_parser():
             AimConfig,
             QuantizedLoraConfig,
             FusedOpsAndKernelsConfig,
-            InstructLabConfig,
+            AttentionAndDistributedPackingConfig,
         )
     )
     parser.add_argument(
@@ -503,7 +509,7 @@ def parse_arguments(parser, json_config=None):
             aim_config,
             quantized_lora_config,
             fusedops_kernels_config,
-            instruct_lab_config,
+            attention_and_distributed_packing_config,
         ) = parser.parse_dict(json_config, allow_extra_keys=True)
         peft_method = json_config.get("peft_method")
         exp_metadata = json_config.get("exp_metadata")
@@ -519,7 +525,7 @@ def parse_arguments(parser, json_config=None):
             aim_config,
             quantized_lora_config,
             fusedops_kernels_config,
-            instruct_lab_config,
+            attention_and_distributed_packing_config,
             additional,
             _,
         ) = parser.parse_args_into_dataclasses(return_remaining_strings=True)
@@ -544,7 +550,7 @@ def parse_arguments(parser, json_config=None):
         aim_config,
         quantized_lora_config,
         fusedops_kernels_config,
-        instruct_lab_config,
+        attention_and_distributed_packing_config,
         exp_metadata,
     )
 
@@ -565,7 +571,7 @@ def main():
             aim_config,
             quantized_lora_config,
             fusedops_kernels_config,
-            instruct_lab_config,
+            attention_and_distributed_packing_config,
             exp_metadata,
         ) = parse_arguments(parser, job_config)
 
@@ -577,7 +583,7 @@ def main():
             model_args %s, data_args %s, training_args %s, trainer_controller_args %s, \
             tune_config %s, file_logger_config, %s aim_config %s, \
             quantized_lora_config %s, fusedops_kernels_config %s, \
-            instruct_lab_config %s exp_metadata %s",
+            attention_and_distributed_packing_config %s exp_metadata %s",
             model_args,
             data_args,
             training_args,
@@ -587,7 +593,7 @@ def main():
             aim_config,
             quantized_lora_config,
             fusedops_kernels_config,
-            instruct_lab_config,
+            attention_and_distributed_packing_config,
             exp_metadata,
         )
     except Exception as e:  # pylint: disable=broad-except
@@ -629,7 +635,7 @@ def main():
             exp_metadata=metadata,
             quantized_lora_config=quantized_lora_config,
             fusedops_kernels_config=fusedops_kernels_config,
-            instruct_lab_config=instruct_lab_config,
+            attention_and_distributed_packing_config=attention_and_distributed_packing_config,
         )
     except (MemoryError, OutOfMemoryError) as e:
         logger.error(traceback.format_exc())