tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-09 22:02:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix gradient accumulation for variable length sequences

Browse Source
This commit is contained in:
kibitzing 2025-10-15 08:56:58 +00:00
parent 67aaca98f5
commit b48d210795
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
scripts/chat_sft.py
View File

@ -34,7 +34,7 @@ from tasks.smoltalk import SmolTalk
# SFT Hyperparameters # SFT Hyperparameters
run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb) run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb)
# input model options # input model options
source = "mid" # base|mid , which checkpoint to load the model from (base model or midtrained model) source = "base" # base|mid , which checkpoint to load the model from (base model or midtrained model)
model_tag = None # model tag to load the model from (base model or midtrained model) model_tag = None # model tag to load the model from (base model or midtrained model)
step = None # step to load the model from (base model or midtrained model) step = None # step to load the model from (base model or midtrained model)
# compute/precision # compute/precision
@ -208,18 +208,24 @@ for step in range(num_iterations):
break break
# evaluate the gradient # evaluate the gradient
total_loss_sum = torch.tensor(0.0, device=device) # sum of losses
num_tokens = torch.tensor(0, device=device) # the number of "active" tokens of supervision seen num_tokens = torch.tensor(0, device=device) # the number of "active" tokens of supervision seen
for micro_step in range(grad_accum_steps): for micro_step in range(grad_accum_steps):
train_inputs, train_targets = next(train_iter) train_inputs, train_targets = next(train_iter)
with autocast_ctx: with autocast_ctx:
loss = model(train_inputs, train_targets) loss = model(train_inputs, train_targets, loss_reduction='sum')
train_loss = loss.detach() # for logging total_loss_sum += loss.detach() # for logging
loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here
loss.backward() # accumulate the gradient loss.backward() # accumulate the gradient
num_tokens += (train_targets >= 0).sum() num_tokens += (train_targets >= 0).sum()
if ddp: if ddp:
dist.all_reduce(total_loss_sum, op=dist.ReduceOp.SUM)
dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM) # sum over ranks dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM) # sum over ranks
# Scale gradients by total number of tokens
for param in model.parameters():
if param.grad is not None:
param.grad.div_(num_tokens.item())
# learning rate scheduler # learning rate scheduler
lrm = get_lr_multiplier(step) lrm = get_lr_multiplier(step)
for opt in optimizers: for opt in optimizers:
@ -232,8 +238,8 @@ for step in range(num_iterations):
model.zero_grad(set_to_none=True) model.zero_grad(set_to_none=True)
# logging # logging
train_loss_item = train_loss.item()
num_tokens_item = num_tokens.item() num_tokens_item = num_tokens.item()
train_loss_item = total_loss_sum.item() / num_tokens_item
print0(f"Step {step:05d}/{num_iterations:05d} | Training loss: {train_loss_item:.6f}| lrm: {lrm:.6f}| num_tokens: {num_tokens_item:,}") print0(f"Step {step:05d}/{num_iterations:05d} | Training loss: {train_loss_item:.6f}| lrm: {lrm:.6f}| num_tokens: {num_tokens_item:,}")
wandb_run.log({ wandb_run.log({
"step": step, "step": step,