Upload folder using huggingface_hub

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00001000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe743aafc4e44b8ed34b1baefaf66ed668a2984d78833af0dbfb4d9bc4273d85
+size 4205813723

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00002000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f32046e81803bcb0cbbd77836ee00e4dd6a81f14387f515fb33627fbbb43d80c
+size 4205813723

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00003000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e572dc1417532d473f46f78f4178e513b47f2d95b8607ba943b8f596f8b4228
+size 4205813723

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00003534.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:538b1d82b75a25ef4b59da7e5312d60ac356d48e1ef7856be7c764a096c85c13
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00004000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f335e3d50acc98aa32eefc0ff3c61392043753cd68f323982989908fd2bd2d6d
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00005000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb6fdf3ae40dd0a87ac20b4682873c2215882b180024f67aafe81bb6e16e1d94
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00006000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:572d4e27d8f19b3dbd6d6681b1b564137c9b573e82ca26510421214f6c15424d
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00006139.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4cfab10e4b7ca011913a2ccc018893caff63004ce45667485a6862c53e2cc42
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00007000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9470bcf094d6ddfd8c613d08584a753e366a11953eb029f993ab3977157b70c0
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00008000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ae4a5c19cec77723c5ccc6b960aad8e6ed33f167975dce5efbf355366dce2b6
+size 4205813787

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00008653.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b354792c613c9a78a430e6f07e0bbc2ca4acab34d8e31fcceb5be2e21d19581
+size 4205813851

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00009000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a29a60defa308d66d57455e7b84dfae1fd0fafd77e8d93f53cf4ec896915657a
+size 4205813851

390M_dynamic_proper_sqrt_lr/blumodel_muon_dynamic_sqrt_step00009042.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76616c00871132c9b2b74c23f58d777b4fff3099f2a3eda592d53d38d34eda40
+size 4205813851

check_target_loss.py

@@ -6,13 +6,11 @@ mid-run, e.g. train_390M_dynamicb.py). Two cases, both reported:
   - target not yet crossed -> fit L(T) = L0 + A*T^-alpha (Eq.14) to the
     post-warmup smoothed trajectory and solve for T algebraically, flagging
     how far past the observed range the extrapolation reaches.
-
 Everything runs on a TOKEN axis, not steps: once the batch size changes, a step
 no longer stands for a fixed amount of data, so L(S) kinks at every switch and
 a predicted step can't be compared across runs. For a fixed batch T = B*S, so
 L0 + A*S^-alpha and L0 + A'*T^-alpha are the same curve (same L0 and alpha,
 A' = A*B^alpha) -- fixed-batch results are unchanged.
-
 Batch size per step: the CSV's `global_batch` column if present (the dynamic
 trainer logs it), else --batch-size, else parsed from the filename (_1Mbs).
 Tokens consumed before each step: the CSV's `tokens` column if present, else
@@ -361,3 +359,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+

train_390M.py

@@ -1,4 +1,3 @@
-
 #!/usr/bin/env python3
 # =============================================================================
 # train_muon.py — the MUON arm of train.py. Same model, same data order, same
@@ -460,11 +459,6 @@ class ModelConfig:
     # Run length in TOKENS. The batch changes during the run, so the step count
     # (max_steps, and with it the WSD phases) is derived from this in main().
     total_tokens: int = 14_000_000_000    # what the fixed-batch arms ran: 17,802 x 786,432
-    # The train shards are cut into windows of data_grain * T tokens, the way
-    # the fixed-batch arms' loader cut them with micro-batch B = data_grain
-    # (6 in the 786K arm), so this run trains on the same tokens in the same
-    # order whatever its batch, micro-batch or GPU count (see DataLoaderLite).
-    data_grain: int = 6
 
     # Batch-size schedule by tokens consumed: ((start_tokens, tokens/step), ...).
     # From tripartite_fit.py (BATCH_SIZE_STUDY.md §4.3). For each target loss
@@ -479,9 +473,10 @@ class ModelConfig:
     #            6.554B -> 8.324B   B_opt(2.4501) = 800K
     #            8.324B -> end      B_opt(2.4376) = 807K  (the last target fitted)
     # Zhou et al. (arXiv:2601.05034, §4.3.2) also switch on tokens consumed and
-    # keep the LR fixed (LR ~ sqrt(B) gave them nothing); train_390M_dynamicb.py
-    # is that arm. THIS arm tests the alternative: the LR follows sqrt(batch)
-    # at every switch (see lr_ref_batch below). The fitted values are
+    # keep the LR fixed (LR ~ sqrt(B) gave them nothing); train_390M_dynamicb_skip.py
+    # is that arm. THIS arm is the same run -- same loader, same batch schedule,
+    # same switch points -- except that the LR follows sqrt(batch) at every
+    # switch (see lr_ref_batch below). The fitted values are
     # not whole numbers of sequences; resolve_batch_schedule() rounds them for
     # the GPU count and main() prints what runs. Empty = global_batch throughout.
     batch_schedule: tuple = ((0,             578_286),
@@ -539,8 +534,10 @@ class ModelConfig:
                                            # target-loss checkpoints are never pruned
     # Separate from the AdamW arm's dir: the two checkpoint formats differ (Muon
     # momentum vs Adam m/v) and auto-resume picks the newest file it finds.
-    ckpt_dir: str = "ckpt_blu_muon_dynamic"
-    ckpt_prefix: str = "blumodel_muon_dynamic"
+    ckpt_dir: str = "ckpt_blu_muon_dynamic_sqrt"
+    ckpt_prefix: str = "blumodel_muon_dynamic_sqrt"
+
+
 
     # ---- runtime knobs (not in the C++ struct) ----
     fused_rms_rope: bool = True
@@ -594,6 +591,7 @@ class BatchStage:
 
 def resolve_batch_schedule(cfg, world_size):
     """cfg.batch_schedule -> [BatchStage], split over `world_size` GPUs.
+
     A step is micro_B * grad_accum sequences on every rank, so a batch must be
     a whole number of sequences per rank (a multiple of T * world_size) that
     factors as micro_B * grad_accum with micro_B <= cfg.B. A target that does
@@ -638,6 +636,7 @@ def keep_batch_schedule(saved, stages, cfg, world_size):
     """On resume: the batches a checkpoint's run was running, split over this
     many GPUs, or None if the schedule config changed since (its starts or
     targets differ from `stages`) or a batch no longer splits evenly.
+
     resolve_batch_schedule() rounds for the GPU count, so a run restarted on a
     different count would otherwise switch to that count's rounding mid-run.
     Going down by a factor (8 -> 4, 2, 1) always keeps the batches.
@@ -707,29 +706,31 @@ def print_batch_schedule(stages, runs, steps_label, ref_batch):
 # 3. Data — port of the C++ DataLoaderLite (Tensor-Implementations/include/
 #    DataLoader.h), including its pinned-host / side-stream prefetch ring.
 #
+# THIS FILE'S LOADER is the one the 390M dynamic-batch run started with, plus a
+# skip for shards shorter than one step (flux_20b's shard_350001), so a resume
+# of that run continues on exactly the data it would have read. It drops each
+# shard's tail per step, which leaves it a little ahead of the fixed-batch
+# arms' data (~+24M tokens at 7.2B). train_390M_dynamicb.py has the loader that
+# reads the fixed-batch arms' token stream instead (restart from scratch, or
+# --realign-data).
+#
 # Shards are HEADERLESS flat uint16 token files. A shard is "val" if "val"
 # appears anywhere in its basename, "train" otherwise — this covers both an
 # explicit flux_train_*/flux_val_* naming and an unmarked-train naming like
 # flux_20b_shard_*.bin alongside a separately named *val* shard. For a batch
 # at position p:
 #     x = tokens[p : p+B*T]      y = tokens[p+1 : p+B*T+1]
-#
-# The token stream. Every shard is cut into windows of grain*T tokens from its
-# start; each window that fits (+1 token for the shifted targets) is used, the
-# shard's tail after the last one is not. That is exactly the set of tokens the
-# fixed-batch arms' loader (train_390M_bs.py: rank r reads windows r, r+W, ...
-# of B*T tokens) read with micro-batch B = grain. The stream is those windows'
-# sequences, shard after shard, and each micro-step takes the next
-# world_size*B sequences of it, rank r the r-th run of B, crossing into the
-# next shard where one runs out. So which tokens a step reads depends only on
-# how many tokens came before it — not on B, grad_accum or world_size: the
-# batch schedule, and a resume on another GPU count, change how steps are cut
-# from the stream, never the stream itself. After E tokens this run is E
-# tokens into the stream, which is where the fixed-batch arms' ranks are
-# spread around (theirs drift apart by a few M tokens across shards).
-#
-# The position (shard, offset) is the same on every rank, and a checkpoint
-# stores it to resume from.
+# Each micro-step reads one contiguous block of world_size*B*T tokens starting
+# at `base`, rank r taking p = base + r*B*T, and base strides by the block. An
+# optimizer step is grad_accum such blocks, so it reads one contiguous range of
+# its batch's tokens; a shard tail that cannot fit a whole STEP is dropped and
+# the loader wraps to the next shard. All ranks move through the blocks
+# together, so the position (shard, base) is one value on every rank, and
+# which tokens a step reads depends only on the batch sizes, not on how they
+# are split into world_size, B and grad_accum: a run resumed on another GPU
+# count reads the same tokens per step. The batch schedule changes the split
+# between steps without skipping or repeating a token, and a checkpoint stores
+# the position to resume from directly.
 #
 # The prefetch mirrors the C++ ring: NBUF pinned staging buffers and NBUF device
 # tensors, H2D issued on a side stream, CUDA events ordering producer against
@@ -740,18 +741,20 @@ class DataLoaderLite:
     NBUF = 4
 
     def __init__(self, B, T, split, data_root, device, master=True,
-                 max_tokens_per_shard=100_000_000, rank=0, world_size=1, grain=None):
+                 max_tokens_per_shard=100_000_000, rank=0, world_size=1):
         assert split in ("train", "val")
-        # B is also the largest micro-batch set_micro_batch() can switch to:
-        # the buffers below are sized for it.
+        # B is also the largest micro-batch set_step() can switch to: the
+        # buffers below are sized for it. A step is grad_accum micro-steps.
         self.B = self.B_max = B
+        self.grad_accum = 1
         self.T = T
         self.split = split
+        self.master = master
         self.max_tokens = max_tokens_per_shard
         self.device = device
-        # DDP: every rank reads the SAME stream but a DISJOINT run of each
-        # micro-step's sequences (offset by rank), so a training step's
-        # world_size*grad_accum micro-batches never repeat a window.
+        # DDP: every rank reads the SAME shard list but a DISJOINT slice of each
+        # block (offset by rank), so a training step's world_size*grad_accum
+        # micro-batches never repeat a window.
         self.rank = rank
         self.world_size = world_size
 
@@ -769,14 +772,6 @@ class DataLoaderLite:
         if master:
             print(f"found {len(self.shards)} shards for split {split}")
 
-        # Tokens of each shard the stream uses: its whole windows (see above).
-        G = (grain or B) * T
-        self.usable = [max(0, (min(os.path.getsize(p) // 2, max_tokens_per_shard) - 1) // G) * G
-                       for p in self.shards]
-        if not any(self.usable):
-            raise RuntimeError(f"no {split} shard holds one window of {G} tokens")
-        self.tok, self.tok_shard = None, None
-
         BT = B * T
         self.pin_x = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
         self.pin_y = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
@@ -786,63 +781,73 @@ class DataLoaderLite:
         self.h2d_done = [torch.cuda.Event() for _ in range(self.NBUF)]
         self.h2d_recorded = [False] * self.NBUF
         self.where = [(0, 0)] * self.NBUF        # (shard_index, token_offset)
-        self.block = [(0, 0)] * self.NBUF        # stream position each slot's micro-step started at
+        self.block = [(0, 0)] * self.NBUF        # (shard_index, base) of each slot's block
         self.reset()
 
-    # -- the stream ----------------------------------------------------------
-    def _tokens(self, s):
-        if s != self.tok_shard:
-            self.tok, self.tok_shard = np.memmap(self.shards[s], dtype=np.uint16, mode="r"), s
-        return self.tok
-
-    def _advance(self, n):
-        """Move n sequences along the stream. -> [(shard, offset, n_seq), ...],
-        the pieces of shards they lie in (two when they cross a shard end)."""
-        pieces = []
-        while n > 0:
-            if self.off >= self.usable[self.cur]:
-                self.cur, self.off = (self.cur + 1) % len(self.shards), 0
-                continue
-            k = min(n, (self.usable[self.cur] - self.off) // self.T)
-            pieces.append((self.cur, self.off, k))
-            self.off += k * self.T
-            n -= k
-        return pieces
-
-    def stream_offset(self, shard, off, tokens):
-        """How far (shard, off) is past `tokens` tokens into the stream (signed,
-        modulo one pass over the shards)."""
-        d = (sum(self.usable[:shard]) + min(off, self.usable[shard]) - tokens) % sum(self.usable)
-        return d - sum(self.usable) if d > sum(self.usable) // 2 else d
+    # -- shard plumbing ------------------------------------------------------
+    def _load_shard(self, i):
+        self.cur = i
+        arr = np.memmap(self.shards[i], dtype=np.uint16, mode="r")
+        self.tok = arr
+        self.ntok = min(len(arr), self.max_tokens)
+
+    def _advance_shard(self):
+        self._load_shard((self.cur + 1) % len(self.shards))
+        self.base = 0
+
+    def _step_tokens(self):
+        return self.B * self.T * self.world_size * self.grad_accum
+
+    def _next_block(self):
+        """-> base of the next block, moving past it. A block that starts a
+        step first checks the WHOLE step fits (+1 for the last shifted
+        target), and leaves the shard's tail for the next shard if not — or
+        the whole shard, if it is smaller than a step (flux_20b's
+        shard_350001 holds 574,132 tokens)."""
+        if self.micro == 0:
+            for _ in range(len(self.shards) + 1):
+                if self.base + self._step_tokens() + 1 <= self.ntok:
+                    break
+                self._advance_shard()
+                if self.master and self._step_tokens() + 1 > self.ntok:
+                    print(f"[data] skipping {os.path.basename(self.shards[self.cur])}: "
+                          f"{self.ntok:,} tokens, less than one step "
+                          f"({self._step_tokens():,})", flush=True)
+            else:
+                raise RuntimeError(f"no {self.split} shard holds one step of "
+                                   f"{self._step_tokens():,} tokens")
+        base = self.base
+        self.base += self.B * self.T * self.world_size
+        self.micro = (self.micro + 1) % self.grad_accum
+        return base
 
     def reset(self):
         self.seek(0, 0)
 
     # -- position (the same on every rank) -----------------------------------
+    # Meant for between optimizer steps, where the next block starts a step.
     def state(self):
-        """-> (shard_index, offset) of the next micro-step next_batch() will
-        return. Prefetched ones it has not returned yet do not count as consumed."""
-        return self.block[self.buf] if self.primed else (self.cur, self.off)
+        """-> (shard_index, base) of the next block next_batch() will return.
+        Prefetched blocks it has not returned yet do not count as consumed."""
+        return self.block[self.buf] if self.primed else (self.cur, self.base)
 
-    def seek(self, shard, off):
+    def seek(self, shard, base):
         """Continue from a state(). Drops anything already prefetched."""
         self.copy_stream.synchronize()
-        self.cur, self.off = shard, -(-off // self.T) * self.T     # a sequence boundary
+        self._load_shard(shard)
+        self.base = base
+        self.micro = 0          # index in its step of the next block to stage
         self.primed = False
         self.buf = 0
 
-    def set_micro_batch(self, B):
-        """Serve B sequences per rank per micro-step from the next one on."""
+    def set_step(self, B, grad_accum=1):
+        """From the next step on, serve grad_accum micro-steps of B sequences
+        per rank per step."""
         if not 1 <= B <= self.B_max:
             raise ValueError(f"micro-batch {B} outside 1..{self.B_max}")
-        if B != self.B:
+        if (B, grad_accum) != (self.B, self.grad_accum):
             self.seek(*self.state())
-            self.B = B
-
-    def skip_tokens(self, n):
-        """Fast-forward n tokens of the stream without touching the GPU."""
-        self.seek(*self.state())
-        self._advance(n // self.T)
+            self.B, self.grad_accum = B, grad_accum
 
     # -- prefetch ------------------------------------------------------------
     def _stage(self, idx):
@@ -857,17 +862,15 @@ class DataLoaderLite:
         if self.h2d_recorded[idx]:
             self.h2d_done[idx].synchronize()
 
-        self.block[idx] = (self.cur, self.off)
-        self._advance(self.rank * B)                               # lower ranks' runs
-        pieces = self._advance(B)                                  # this rank's
-        self._advance((self.world_size - 1 - self.rank) * B)       # higher ranks'
-        self.where[idx] = pieces[0][:2]
-        at = 0
-        for s, off, k in pieces:
-            buf = np.asarray(self._tokens(s)[off: off + k * self.T + 1]).astype(np.int64)
-            self.pin_x[idx][at: at + k * self.T].copy_(torch.from_numpy(buf[:-1]))
-            self.pin_y[idx][at: at + k * self.T].copy_(torch.from_numpy(buf[1:]))
-            at += k * self.T
+        # Recorded BEFORE _next_block()'s does-the-step-fit check, so a seek()
+        # to it after set_step() decides that again for the new step size.
+        self.block[idx] = (self.cur, self.base)
+        base = self._next_block()
+        pos = base + BT * self.rank
+        self.where[idx] = (self.cur, pos)
+        buf = np.asarray(self.tok[pos: pos + BT + 1]).astype(np.int64)
+        self.pin_x[idx][:BT].copy_(torch.from_numpy(buf[:-1]))
+        self.pin_y[idx][:BT].copy_(torch.from_numpy(buf[1:]))
 
         with torch.cuda.stream(self.copy_stream):
             self.dev_x[idx][:B].view(-1).copy_(self.pin_x[idx][:BT], non_blocking=True)
@@ -898,6 +901,13 @@ class DataLoaderLite:
         return (self.dev_x[i][:self.B], self.dev_y[i][:self.B],
                 self.where[i][0], self.where[i][1])
 
+    def skip_batches(self, n):
+        """Fast-forward `n` micro-steps without touching the GPU (resume from a
+        checkpoint that has no stored position)."""
+        self.seek(*self.state())
+        for _ in range(n):
+            self._next_block()
+
 
 # =============================================================================
 # 4. Model.
@@ -1469,15 +1479,6 @@ def build_argparser():
     ap.add_argument("--total-tokens", type=float, default=None,
                     help="run length in tokens (default 14e9); the step count is "
                          "derived from it and the batch schedule")
-    ap.add_argument("--data-grain", type=int, default=None,
-                    help="cut the train shards into windows of this many sequences, "
-                         "as the fixed-batch arms' loader did with that micro-batch "
-                         "(default 6, the 786K arm's)")
-    ap.add_argument("--realign-data", action="store_true",
-                    help="on resume from a checkpoint an earlier loader wrote, move "
-                         "the data position back to tokens-seen into the stream, "
-                         "where the fixed-batch arms are; the tokens in between are "
-                         "seen twice. Default: continue from the stored position.")
     ap.add_argument("--max-steps", type=int, default=None,
                     help="caps the run for benchmarking; warmup/stable shrink "
                          "proportionally so the WSD decay tail stays non-empty "
@@ -1518,7 +1519,7 @@ def build_argparser():
                     help="per-micro-batch loss + depth probes (env: BLU_DIAG=1)")
     ap.add_argument("--warmup-measure", type=int, default=3,
                     help="steps excluded from the steady-state throughput summary")
-    ap.add_argument("--csv", default="blu_muon_log.csv",
+    ap.add_argument("--csv", default="blu_muon_log_dy_sq.csv",
                     help="CSV log path; appended (not truncated) on resume. The "
                          "first 8 columns match the AdamW arm's CSV exactly, so "
                          "the two files plot against each other directly; "
@@ -1552,7 +1553,7 @@ def build_argparser():
 def config_from_args(args):
     cfg = ModelConfig()
     for f in ("B", "T", "d_model", "n_layers", "q_heads", "kv_heads", "ffn_hidden",
-              "vocab_size", "global_batch", "data_grain", "val_freq", "val_loss_steps",
+              "vocab_size", "global_batch", "val_freq", "val_loss_steps",
               "ckpt_dir", "ckpt_prefix", "ckpt_freq", "ckpt_keep",
               "muon_lr", "muon_min_lr", "muon_momentum", "muon_ns_steps",
               "muon_adjust_lr", "target_losses", "target_loss_ema_beta",
@@ -1834,7 +1835,7 @@ def main():
     # the same one.
     train_loader = DataLoaderLite(cfg.B, cfg.T, "train", args.data_root, device,
                                   master=master_process, rank=ddp_rank,
-                                  world_size=ddp_world_size, grain=cfg.data_grain)
+                                  world_size=ddp_world_size)
     # rank=0, world_size=1 UNCONDITIONALLY, on every rank: val_loss_steps
     # windows at fixed token offsets 0, BT, 2BT, ... Sharding this by
     # ddp_rank/ddp_world_size like the train loader would make the reported
@@ -1882,21 +1883,6 @@ def main():
                 # are independent of B and world_size, so this is exact.
                 tokens_seen = progress["tokens_seen"]
                 train_loader.seek(*progress["data_pos"])
-                # 0 when this loader wrote the checkpoint: E tokens in, the run
-                # is E tokens into the stream. An earlier loader dropped more
-                # of each shard's tail, which left the run ahead of that.
-                ahead = train_loader.stream_offset(*progress["data_pos"], tokens_seen)
-                if args.realign_data and ahead:
-                    train_loader.reset()
-                    train_loader.skip_tokens(tokens_seen)
-                    if master_process:
-                        print(f"[resume] --realign-data: moved the data position back "
-                              f"{ahead:,} tokens to {tokens_seen:,} tokens into the "
-                              f"stream; those tokens are seen a second time")
-                elif master_process and ahead:
-                    print(f"[resume] the data position is {ahead:+,} tokens from "
-                          f"{tokens_seen:,} tokens into the stream (left by an "
-                          f"earlier loader); it stays at that offset from here")
                 # Keep running the batches the run started with when this GPU
                 # count can (always true going 8 -> 4, 2, 1), rather than this
                 # count's own rounding: changing GPUs then changes nothing but
@@ -1916,10 +1902,11 @@ def main():
                           f"follows the schedule above")
             else:
                 # A checkpoint from a trainer without a batch schedule: assume
-                # steps [0, start_step) ran this schedule, and go that many
-                # tokens into the stream.
-                _, tokens_seen = walk_schedule(stages, float("inf"), start_step)
-                train_loader.skip_tokens(tokens_seen)
+                # steps [0, start_step) ran this schedule and replay it.
+                done, tokens_seen = walk_schedule(stages, float("inf"), start_step)
+                for s, n in done:
+                    train_loader.set_step(s.micro_B, s.grad_accum)
+                    train_loader.skip_batches(n * s.grad_accum)
                 if master_process:
                     print(f"[resume] checkpoint has no data position; replayed this "
                           f"schedule to step {start_step} ({tokens_seen:,} tokens)")
@@ -2030,7 +2017,7 @@ def main():
         # carry straight on.
         new_stage = stage_at(stages, tokens_seen)
         if stage is None or (new_stage.batch, new_stage.micro_B) != (stage.batch, stage.micro_B):
-            train_loader.set_micro_batch(new_stage.micro_B)
+            train_loader.set_step(new_stage.micro_B, new_stage.grad_accum)
             grad_accum = new_stage.grad_accum
             tokens_per_step = new_stage.batch
             lr_scale = lr_batch_scale(new_stage.batch, ref_batch)
@@ -2330,8 +2317,6 @@ if __name__ == "__main__":
 
 
 
-
-