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

est total tokens, fix field loop

Browse Source
This commit is contained in:
Wing Lian
2023-07-18 11:30:07 -04:00
parent 41d4992029
commit 66774011c4
2 changed files with 37 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -155,6 +155,16 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
if cfg.sample_packing: if cfg.sample_packing:
train_dataset = train_dataset.map(add_position_ids) train_dataset = train_dataset.map(add_position_ids)
eval_dataset = eval_dataset.map(add_position_ids) eval_dataset = eval_dataset.map(add_position_ids)
if cfg.sample_packing_eff_est:
total_num_tokens = sum(len(s["input_ids"]) for s in train_dataset)
total_num_steps = math.ceil(
total_num_tokens
/ cfg.sample_packing_eff_est
/ 2048
* cfg.num_epochs
/ cfg.batch_size
)
else:
sampler = DistributedSampler( sampler = DistributedSampler(
train_dataset, train_dataset,
num_replicas=1, num_replicas=1,
@@ -176,7 +186,10 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
LOG.info(f"data_loader_len: {data_loader_len}") LOG.info(f"data_loader_len: {data_loader_len}")
total_num_steps = int( total_num_steps = int(
math.ceil( math.ceil(
data_loader_len * cfg.micro_batch_size * cfg.num_epochs / cfg.batch_size data_loader_len
* cfg.micro_batch_size
* cfg.num_epochs
/ cfg.batch_size
) )
) )
else: else: