more fixes for dataloader integration

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

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")

View File

@@ -91,12 +91,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()
@@ -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