add validation for batch flattening

src/axolotl/core/trainer_builder.py

@@ -2006,11 +2006,7 @@ def build_collator(
                 collator = MultiModalChatDataCollator
                 kwargs["processor"] = self.processor
                 kwargs["chat_template"] = training_args.chat_template
-            elif (
-                self.cfg.flash_attention
-                and self.cfg.micro_batch_size > 1
-                and not self.cfg.sample_packing
-            ):
+            elif self.cfg.batch_flattening:
                 collator = DataCollatorWithFlattening
                 collator_args.pop(0)
                 kwargs.pop("pad_to_multiple_of", None)

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -696,6 +696,8 @@ class Config:
     curriculum_sampling: Optional[bool] = None
     multipack_real_batches: Optional[bool] = None
 
+    batch_flattening: Optional[bool] = None
+
     # for PoSE context length extension
     use_pose: Optional[bool] = None
     pose_split_on_token_ids: Optional[List[int]] = None
@@ -923,6 +925,19 @@ def check_sample_packing_wo_flash(cls, data):
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_batch_flattening_fa(cls, data):
+        if data.get("batch_flattening"):
+            if not data.get("flash_attention"):
+                raise ValueError("batch_flattening requires flash attention")
+            if data.get("sample_packing"):
+                raise ValueError("batch_flattening not compatible with sample_packing")
+            if data.get("micro_batch_size") == 1:
+                LOG.warning("batch_flattening has no effect with micro_batch_size == 1")
+
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_sample_packing_w_rl(cls, data):