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

more fixes for dataloader integration

Browse Source
This commit is contained in:
Wing Lian
2023-07-18 10:50:40 -04:00
parent 762f1b08db
commit 41d4992029
2 changed files with 14 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
scripts/finetune.py
View File

@@ -231,7 +231,7 @@ def train(
cfg.pretraining_dataset, cfg.pretraining_dataset,
tokenizer, tokenizer,
max_tokens=cfg.sequence_len, max_tokens=cfg.sequence_len,
seed=cfg.seed, seed=cfg.seed or 42,
) )
# https://discuss.huggingface.co/t/how-to-use-huggingface-trainer-streaming-datasets-without-wrapping-it-with-torchdatas-iterablewrapper/25230 # https://discuss.huggingface.co/t/how-to-use-huggingface-trainer-streaming-datasets-without-wrapping-it-with-torchdatas-iterablewrapper/25230
train_dataset = train_dataset.with_format("torch") train_dataset = train_dataset.with_format("torch")

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

@@ -91,13 +91,15 @@ 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(
MultipackDistributedDataloader(
self.train_dataset, self.train_dataset,
batch_size=self._train_batch_size, batch_size=self._train_batch_size,
seq_max_length=self.args.max_seq_length, seq_max_length=self.args.max_seq_length,
collate_fn=self.data_collator, collate_fn=self.data_collator,
sampler=train_sampler, sampler=train_sampler,
) )
)
return super().get_train_dataloader() return super().get_train_dataloader()
def get_eval_dataloader( def get_eval_dataloader(
@@ -157,7 +159,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
train_dataset, train_dataset,
num_replicas=1, num_replicas=1,
rank=0, rank=0,
seed=cfg.seed, seed=cfg.seed or 42,
) )
data_loader = MultipackDistributedDataloader( data_loader = MultipackDistributedDataloader(
train_dataset, train_dataset,
@@ -170,12 +172,11 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
), ),
sampler=sampler, sampler=sampler,
) )
data_loader_len = len(data_loader)
LOG.info(f"data_loader_len: {data_loader_len}")
total_num_steps = int( total_num_steps = int(
math.ceil( math.ceil(
len(data_loader) data_loader_len * cfg.micro_batch_size * cfg.num_epochs / cfg.batch_size
* cfg.micro_batch_size
* cfg.num_epochs
/ cfg.batch_size
) )
) )
else: else:
@@ -262,8 +263,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
training_arguments_kwargs["save_safetensors"] = cfg.save_safetensors training_arguments_kwargs["save_safetensors"] = cfg.save_safetensors
training_args = AxolotlTrainingArguments( # pylint: disable=unexpected-keyword-arg training_args = AxolotlTrainingArguments( # pylint: disable=unexpected-keyword-arg
max_steps=total_num_steps max_steps=total_num_steps, # this is helpful in case we don't actually know total # of steps
* cfg.num_epochs, # this is helpful in case we don't actually know total # of steps
per_device_train_batch_size=cfg.micro_batch_size, per_device_train_batch_size=cfg.micro_batch_size,
per_device_eval_batch_size=cfg.eval_batch_size per_device_eval_batch_size=cfg.eval_batch_size
if cfg.eval_batch_size is not None if cfg.eval_batch_size is not None