Upload 2 files

code/train.py

@@ -320,12 +320,10 @@ for step in range(start_step, end_step):
     accum_loss = 0.0
     for micro_i in range(GRAD_ACCUM):
         x, y = get_batch(shard_pool, MICRO_BATCH, CONTEXT_LEN)
-        sync_ctx = model.no_sync() if micro_i < GRAD_ACCUM - 1 else contextlib.nullcontext()
-        with sync_ctx:
-            with torch.autocast(device_type="cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
-                _, loss = model(x, y)
-                loss = loss / GRAD_ACCUM
-            loss.backward()
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
+            _, loss = model(x, y)
+            loss = loss / GRAD_ACCUM
+        loss.backward()
         accum_loss += loss.item()
         del x, y, loss