tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

test v2batch w/ flex attn

Browse Source
This commit is contained in:
bursteratom
2025-02-13 00:11:45 -05:00
parent 0ef1f011fe
commit 82d04ea060
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/axolotl/core/trainer_builder.py
View File

@@ -830,7 +830,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
kwargs.pop("max_length") kwargs.pop("max_length")
elif use_batch_sampler_collator: elif use_batch_sampler_collator:
if self.cfg.flex_attention is True: if self.cfg.flex_attention is True:
collator = FlexBatchSamplerDataCollatorForSeq2Seq collator = V2BatchSamplerDataCollatorForSeq2Seq
elif self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES: elif self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES:
collator = V2BatchSamplerDataCollatorForSeq2Seq collator = V2BatchSamplerDataCollatorForSeq2Seq
elif ( elif (