diff --git a/scripts/mid_train.py b/scripts/mid_train.py index 202682da..811df9e8 100644 --- a/scripts/mid_train.py +++ b/scripts/mid_train.py @@ -15,6 +15,7 @@ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" import time import wandb import torch +from contextlib import nullcontext from nanochat.common import compute_init, compute_cleanup, print0, DummyWandb, get_base_dir from nanochat.tokenizer import get_token_bytes @@ -217,13 +218,18 @@ while True: torch.cuda.synchronize() t0 = time.time() for micro_step in range(grad_accum_steps): - with autocast_ctx: - loss = model(x, y) - train_loss = loss.detach() # for logging - loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here - loss.backward() - x, y = next(train_loader) # prefetch the next batch while the GPU is busy with forward/backward - progress = max(progress, approx_progress) # only increase progress monotonically + # Avoid redundant all-reduce during gradient accumulation. + sync_needed = (not ddp) or (micro_step == grad_accum_steps - 1) + sync_context = nullcontext() if sync_needed else model.no_sync() + with sync_context: + with autocast_ctx: + loss = model(x, y) + if micro_step == 0: + train_loss = loss.detach() # log once per iteration + loss = loss / grad_accum_steps # normalize loss + loss.backward() + x, y = next(train_loader) # prefetch next batch while backward runs + progress = max(progress, approx_progress) # step the optimizers lrm = get_lr_multiplier(progress) for opt in optimizers: