restrict use cases for flattening

axolotl-ai-cloud · Dec 12, 2024 · 7bd5478 · 7bd5478
1 parent c13a148
commit 7bd5478
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
@@ -2006,7 +2006,11 @@ def build_collator(
                 collator = MultiModalChatDataCollator
                 kwargs["processor"] = self.processor
                 kwargs["chat_template"] = training_args.chat_template
-            elif self.cfg.flash_attention:
+            elif (
+                self.cfg.flash_attention
+                and self.cfg.micro_batch_size > 1
+                and not self.cfg.sample_packing
+            ):
                 collator = DataCollatorWithFlattening
                 collator_args.pop(0)
                 kwargs.pop("pad_to_multiple_of", None)