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

use accelerator prepare for dataloader

Browse Source
This commit is contained in:
Wing Lian
2023-07-19 22:58:16 -04:00
parent 4ab9ab79fd
commit 2e295c9f94
1 changed files with 8 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/axolotl/utils/trainer.py
View File

@@ -101,12 +101,14 @@ class AxolotlTrainer(Trainer):
if self.args.sample_packing: if self.args.sample_packing:
train_sampler = self._get_train_sampler() train_sampler = self._get_train_sampler()
return MultipackDistributedDataloader( return self.accelerator.prepare(
self.train_dataset, MultipackDistributedDataloader(
batch_size=self._train_batch_size, self.train_dataset,
seq_max_length=self.args.max_seq_length, batch_size=self._train_batch_size,
collate_fn=self.data_collator, seq_max_length=self.args.max_seq_length,
sampler=train_sampler, collate_fn=self.data_collator,
sampler=train_sampler,
)
) )
return super().get_train_dataloader() return super().get_train_dataloader()