Upload folder using huggingface_hub

390M_1M_8ksteps_main/blumodel_muon_step00003461.pt

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

390M_1M_8ksteps_main/blumodel_muon_step00004065.pt

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

390M_1M_8ksteps_main/blumodel_muon_step00005697.pt

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

390M_1M_8ksteps_main/blumodel_muon_step00006362.pt

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

390M_1M_8ksteps_main/blumodel_muon_step00007249.pt

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

390M_1M_8ksteps_main/blumodel_muon_step00007250.pt

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

390M_1M_8ksteps_main/blumodel_muon_step00008700.pt

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

check_target_loss.py

@@ -1,5 +1,3 @@
-
-
 #!/usr/bin/env python3
 """Find when a training run (CSV log) reaches a target loss, after smoothing.
 Two cases, both reported:
@@ -20,6 +18,11 @@ EXTRAPOLATION_WARN_FACTOR = 3.0
 CONVERGENCE_TOL = 0.05
 SWEEP_STEP = 300
 SWEEP_MIN_WINDOW = 500
+EMA_TOKEN_WINDOW = 100_000_000   # smoothing window in TOKENS, not steps -- see
+                                  # batchsize_scaling.py for why: a fixed step
+                                  # count covers a different amount of data per
+                                  # run depending on batch size, which biases
+                                  # the cross-batch-size comparison.
 
 
 def power_law(s, l0, a, alpha):
@@ -34,6 +37,19 @@ def parse_batch_size_from_filename(path):
     return value * (1_000_000 if unit == "M" else 1_000)
 
 
+def ema_span_for_token_window(batch_size, token_window):
+    """Convert a token-based smoothing window into the step-based `span`
+    pandas' .ewm() expects, so every run gets an equivalent amount of actual
+    data smoothed over regardless of its batch size."""
+    span = token_window / batch_size
+    if span < 1:
+        print(f"[WARN] EMA_TOKEN_WINDOW ({token_window:,}) < batch size "
+              f"({batch_size:,.0f}) -- window is under one step; clamping "
+              "span to 1 (effectively no smoothing).")
+        return 1.0
+    return span
+
+
 def detect_warmup_end(df):
     lr_max = df["lr"].max()
     past_warmup = df.index[df["lr"] >= 0.99 * lr_max]
@@ -122,8 +138,12 @@ def main():
     p.add_argument("csv")
     p.add_argument("--target", type=float, nargs="+", required=True)
     p.add_argument("--loss-col", default="train_loss")
-    p.add_argument("--ema-span", type=float, default=50,
-                    help="EMA span in steps for smoothing (default: 50)")
+    p.add_argument("--ema-token-window", type=float, default=EMA_TOKEN_WINDOW,
+                    help=f"EMA smoothing window in TOKENS, not steps (default: "
+                         f"{EMA_TOKEN_WINDOW:,.0f}) -- converted to a step-based "
+                         "span using --batch-size, so different batch sizes get "
+                         "an equivalent amount of data smoothed, not an "
+                         "equivalent step count")
     p.add_argument("--batch-size", type=float, default=None,
                     help="tokens per step; parsed from filename (e.g. _1Mbs) if omitted")
     p.add_argument("--fit-start-step", type=int, default=None,
@@ -155,13 +175,15 @@ def main():
     if fit_start is None:
         fit_start = detect_warmup_end(df)
 
-    smoothed_full = df[args.loss_col].ewm(span=args.ema_span, adjust=False).mean()
+    ema_span = ema_span_for_token_window(batch_size, args.ema_token_window)
+    smoothed_full = df[args.loss_col].ewm(span=ema_span, adjust=False).mean()
     post_warmup = df[df["step"] >= fit_start]
     smoothed_post = smoothed_full.loc[post_warmup.index]
     steps_post = post_warmup["step"]
     s_max_observed = steps_post.max()
 
-    print(f"loss col: {args.loss_col} | ema span: {args.ema_span} steps | "
+    print(f"loss col: {args.loss_col} | ema token window: {args.ema_token_window:,.0f} "
+          f"(~{ema_span:.1f} steps at this batch size) | "
           f"batch size: {batch_size:,.0f} tokens/step | fit start step: {fit_start}")
     print("-" * 100)
 
@@ -195,5 +217,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-

train_390M.py

@@ -464,10 +464,10 @@ class ModelConfig:
     muon_lr: float = 0.001
     muon_min_lr: float = 0.0001
     
-    max_steps: int = 26703
+    max_steps: int = 13351
 
     warmup_steps: int = 500
-    stable_steps: int = int(max_steps * 0.9) - warmup_steps           # steps held at max_lr after warmup
+    stable_steps: int = int(max_steps * 0.99) - warmup_steps           # steps held at max_lr after warmup
      
 
     # ---- Muon (hidden 2D weights); the AdamW fields above still drive the
@@ -488,7 +488,9 @@ class ModelConfig:
 
     # ---- checkpointing ----
     checkpointing: bool = True
-    ckpt_freq: int = 1000                 # save every 1800 optimizer steps
+    ckpt_freq: int = 1000                 # unused by the save trigger (milestone-based
+                                           # now, see target_losses below) — kept in case
+                                           # interval saving is wanted again later
     ckpt_keep: int = 40                    # retain the most-recent N checkpoints
     # 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.
@@ -499,6 +501,18 @@ class ModelConfig:
     fused_rms_rope: bool = True
     activation_checkpoint: bool = False
     loss_chunks: int = 1
+    # Milestones to record (step + elapsed hours) AND checkpoint on, as the
+    # EMA-smoothed train loss passes each one — the same shared target-loss
+    # set used across this project's batch-size ablation arms. Raw per-step
+    # loss is noisy, so the comparison runs against an EMA, not loss_val
+    # directly — same convention as batchsize_scaling.py's EMA_BETA=0.99.
+    # A checkpoint is saved the step each one is first crossed (replacing
+    # interval-based saving): this ablation's whole point is comparing
+    # batch sizes at matched loss levels, so that's the checkpoint that's
+    # actually useful, not whatever happens to land on a round step number.
+    target_losses: tuple = (11.035316, 2.690595, 2.554042, 2.505404,
+                            2.480678, 2.450056, 2.437589)
+    target_loss_ema_beta: float = 0.99
     # Depth probe: which layers --diag reports gradient / weight statistics for.
     probe_layers: tuple = (0, 8, 17, 26, 35)
 
@@ -1064,7 +1078,7 @@ class CheckpointManager:
                 out.append((int(m.group(1)), p))
         return sorted(out)
 
-    def save(self, step, model, opt, loss, last_val, cfg):
+    def save(self, step, model, opt, loss, last_val, cfg, milestones=None):
         tmp = self._path(step) + ".tmp"
         torch.save({"step": step, "loss": loss, "last_val": last_val,
                     # B/T/global_batch pin the data stride this checkpoint's
@@ -1075,6 +1089,12 @@ class CheckpointManager:
                     # running, so a GPU-count change alone is not a mismatch.
                     "run_cfg": {"B": cfg.B, "T": cfg.T,
                                 "global_batch": cfg.global_batch},
+                    # Target losses already crossed (loss -> (step, hours)).
+                    # Without this, a resume re-initializes the EMA from
+                    # scratch at whatever loss the model is already at, which
+                    # falsely reports every target below that value as
+                    # "reached" at the resume step, all at once.
+                    "milestones": milestones or {},
                     "model": model.state_dict(),
                     "opt": opt.state_dict()}, tmp)
         # Rename last: a crash mid-write leaves a .tmp, never a truncated
@@ -1085,18 +1105,18 @@ class CheckpointManager:
                 os.remove(stale)
 
     def load_latest(self, model, opt, cfg, master=True):
-        """-> (start_step, loss, last_val). (0, 0.0, nan) when there is
-        nothing to resume."""
+        """-> (start_step, loss, last_val, milestones). (0, 0.0, nan, {})
+        when there is nothing to resume."""
         existing = self._existing()
         if not existing:
-            return 0, 0.0, float("nan")
+            return 0, 0.0, float("nan"), {}
         _, path = existing[-1]
         return self.load(path, model, opt, cfg, master)
 
     def load(self, path, model, opt, cfg, master=True):
         """Load a specific checkpoint file (e.g. from --resume-from), rather
         than auto-discovering the latest one in self.dir.
-        -> (start_step, loss, last_val)."""
+        -> (start_step, loss, last_val, milestones)."""
         blob = torch.load(path, map_location="cpu", weights_only=True)
         model.load_state_dict(blob["model"])
         if not opt.load_state_dict(blob["opt"]):
@@ -1112,7 +1132,8 @@ class CheckpointManager:
                     print(f"         {m}")
         print(f"[resume] loaded {path}")
         return (int(blob["step"]), float(blob["loss"]),
-                float(blob.get("last_val", float("nan"))))
+                float(blob.get("last_val", float("nan"))),
+                blob.get("milestones", {}))
 
 
 # =============================================================================
@@ -1172,7 +1193,7 @@ def build_argparser():
     ap = argparse.ArgumentParser(
         description="BLU single-GPU training — PyTorch port of bluscript_single.cpp")
     ap.add_argument("--data-root",
-                    default=("/home/operator/.cache/huggingface/hub/datasets--jayasriii--flux_20b/"
+                    default=("/root/.cache/huggingface/hub/datasets--jayasriii--flux_20b/"
                              "/snapshots/ddc6cc7e43365a6b9096c30053df0fb001c94f4c"))
     ap.add_argument("--device", type=int, default=None,
                     help="CUDA device index (env: BLU_DEVICE)")
@@ -1204,6 +1225,14 @@ def build_argparser():
                     help="caps the run for benchmarking; warmup/stable shrink "
                          "proportionally so the WSD decay tail stays non-empty "
                          "(env: BLU_MAX_STEPS)")
+    ap.add_argument("--target-losses", type=float, nargs="+", default=None,
+                    help="loss milestones to record (step + elapsed hours) and "
+                         "checkpoint on, as the EMA-smoothed train loss passes "
+                         "each one; default is the shared target-loss set used "
+                         "across this project's batch-size ablation arms.")
+    ap.add_argument("--target-loss-ema-beta", type=float, default=None,
+                    help="EMA smoothing beta applied before comparing against "
+                         "--target-losses (default 0.99)")
     # ---- Muon ----
     ap.add_argument("--muon-lr", type=float, default=None,
                     help="peak LR for the hidden 2D weights. Under the default "
@@ -1263,7 +1292,7 @@ def config_from_args(args):
               "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"):
+              "muon_adjust_lr", "target_losses", "target_loss_ema_beta"):
         v = getattr(args, f, None)
         if v is not None:
             setattr(cfg, f, v)
@@ -1427,6 +1456,9 @@ def main():
         print(f"  warmup / steps  : {cfg.warmup_steps} / {cfg.max_steps}")
         print(f"  WSD stable_steps: {cfg.stable_steps}  "
               f"(decay tail={cfg.max_steps - cfg.warmup_steps - cfg.stable_steps})")
+        tl_str = ", ".join(f"{t:g}" for t in cfg.target_losses)
+        print(f"  target losses   : [{tl_str}]  (ema beta {cfg.target_loss_ema_beta:g}, "
+              f"checkpointed on crossing — ckpt_freq is unused)")
         print(f"  precision       : {args.precision}")
         print(f"  seed / determ.  : {args.seed} / {args.determinism}")
         print(f"  rms+rope        : {'fused' if cfg.fused_rms_rope else 'eager'}"
@@ -1536,15 +1568,16 @@ def main():
     # explicit broadcast. Only rank 0 ever WRITES one (below).
     start_step, ckpt = 0, None
     last_val = float("nan")
+    resumed_milestones = {}
     if cfg.checkpointing or args.resume_from:
         ckpt = CheckpointManager(cfg.ckpt_dir, cfg.ckpt_prefix, cfg.ckpt_keep)
         if args.resume_from:
             if master_process:
                 print(f"[resume] loading explicit checkpoint: {args.resume_from}")
-            start_step, resume_loss, last_val = ckpt.load(
+            start_step, resume_loss, last_val, resumed_milestones = ckpt.load(
                 args.resume_from, raw_model, opt, cfg, master_process)
         else:
-            start_step, resume_loss, last_val = ckpt.load_latest(
+            start_step, resume_loss, last_val, resumed_milestones = ckpt.load_latest(
                 raw_model, opt, cfg, master_process)
         if start_step > 0:
             if master_process:
@@ -1574,6 +1607,14 @@ def main():
     probes = [L for L in cfg.probe_layers if L < cfg.n_layers]
     steady_dts = []
     t_train_start = time.perf_counter()
+    ema_loss = None   # EMA-smoothed train loss, checked against target_losses
+    # Milestones already crossed, restored from the checkpoint on resume so a
+    # restart doesn't re-discover them all at once (see CheckpointManager.save).
+    milestones = dict(resumed_milestones)   # target loss -> (step, elapsed_hours)
+    # Remaining targets, highest loss first (the order the EMA passes them in
+    # as it decreases) — excludes anything already in `milestones`.
+    pending_targets = [t for t in sorted(cfg.target_losses, reverse=True)
+                       if t not in milestones]
 
     # ---- Ctrl-C: finish the in-flight step, save a checkpoint, exit clean ----
     # Registered here (not earlier) so a Ctrl-C during setup/data-loading still
@@ -1675,6 +1716,22 @@ def main():
                 dist.destroy_process_group()
             raise SystemExit(f"ERROR: NaN/Inf loss at step {step}")
 
+        # EMA-smoothed train loss vs. the pending milestones. loss_val is
+        # already all-reduced across ranks above, so ema_loss (and the
+        # milestones it crosses) comes out identical on every rank without an
+        # extra collective.
+        beta = cfg.target_loss_ema_beta
+        ema_loss = loss_val if ema_loss is None else beta * ema_loss + (1 - beta) * loss_val
+        elapsed_hr = (time.perf_counter() - t_train_start) / 3600.0
+        target_reached = False
+        while pending_targets and ema_loss <= pending_targets[0]:
+            t = pending_targets.pop(0)
+            milestones[t] = (step, elapsed_hr)
+            target_reached = True
+            if master_process:
+                print(f"\n[target] loss <= {t:.6f} reached at step {step} "
+                      f"({elapsed_hr:.3f} hrs)\n", flush=True)
+
         # Depth probe (--diag), read BEFORE the update, so these are the
         # pre-update gradients. Compare the profile on a spiking step against a
         # quiet one: a blow-up that begins mid-stack and grows toward the loss
@@ -1751,7 +1808,7 @@ def main():
             dist.all_reduce(flag, op=dist.ReduceOp.MAX)
             interrupted = flag.item() > 0
 
-        # ---------- periodic checkpoint (also fires once on Ctrl-C) ----------
+        # ---------- checkpoint on target-loss crossing (also end-of-run / Ctrl-C) ----------
         # Labeled step+1 (the next iteration to run) so load_latest resumes
         # exactly where this save left off — no duplicated or skipped iteration.
         # Captures weights + Adam m/v AFTER this step's update. Rank 0 only —
@@ -1760,12 +1817,19 @@ def main():
         # barrier holds every rank at this point until the save (or its
         # failure) is done, so a slow write can't let other ranks race ahead
         # into the next step's collectives before the checkpoint is durable.
-        if cfg.checkpointing and ((step + 1) % cfg.ckpt_freq == 0
+        #
+        # target_reached replaces interval-based saving: this ablation's whole
+        # point is comparing batch sizes at matched loss levels, so the
+        # checkpoint that's actually useful is the one at each milestone, not
+        # whatever happens to land on a round step number. target_reached is
+        # derived from loss_val, which is already all-reduced above, so it's
+        # identical on every rank without its own collective.
+        if cfg.checkpointing and (target_reached
                                   or step == cfg.max_steps - 1
                                   or interrupted):
             if master_process:
                 try:
-                    ckpt.save(step + 1, raw_model, opt, loss_val, last_val, cfg)
+                    ckpt.save(step + 1, raw_model, opt, loss_val, last_val, cfg, milestones)
                 except Exception as e:                   # noqa: BLE001 — advisory
                     print(f"[WARN] checkpoint save at step {step + 1} failed: {e}")
             if ddp:
@@ -1788,6 +1852,20 @@ def main():
     if log_file:
         log_file.close()
 
+    # ---- target-loss milestone summary ----
+    if master_process:
+        print("\n" + "=" * 52)
+        print("  TARGET LOSS MILESTONES")
+        print("=" * 52)
+        print(f"  {'target loss':>12}   {'step':>10}   {'hours':>10}")
+        for t in sorted(cfg.target_losses, reverse=True):
+            if t in milestones:
+                step_hit, hrs = milestones[t]
+                print(f"  {t:>12.6f}   {step_hit:>10}   {hrs:>10.3f}")
+            else:
+                print(f"  {t:>12.6f}   {'not reached':>10}   {'-':>10}")
+        print("=" * 52)
+
     # ---- steady-state summary ----
     # dt/throughput are this rank's own timing; under DDP all ranks run in
     # lockstep (grad sync + the checkpoint barrier), so rank 0's numbers are
@@ -1855,5 +1933,6 @@ if __name__ == "__main__":
 
 
 
+