Blame for iclr-debug-generator-decay27b-20260917/source.patch

diff --git a/evaluation/eval_random_gen_ppl_checkpoint.py b/evaluation/eval_random_gen_ppl_checkpoint.py
index 2d5c741..eb4eb89 100644
--- a/evaluation/eval_random_gen_ppl_checkpoint.py
+++ b/evaluation/eval_random_gen_ppl_checkpoint.py
@@ -5,6 +5,7 @@ import json
 import math
 import sys
 from collections import Counter
+from contextlib import contextmanager
 from pathlib import Path
 
 REPO_ROOT = Path(__file__).resolve().parents[1]
@@ -22,14 +23,35 @@ from utils.misc import set_manual_seed
 from utils.registry import trainers
 
 
+def seed_evaluation(seed, reference_dtype):
+    # The shared seed helper enables TF32, so precision must be set afterward.
+    set_manual_seed(seed)
+    if reference_dtype in ("float32", "float64"):
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+
+
+@contextmanager
+def reference_precision_context(reference_dtype):
+    previous_dtype = torch.get_default_dtype()
+    try:
+        # Transformers' causal-mask helper creates a tensor without a dtype
+        # before converting it. Its FP64 minimum overflows a default FP32 tensor.
+        if reference_dtype == "float64":
+            torch.set_default_dtype(torch.float64)
+        yield
+    finally:
+        torch.set_default_dtype(previous_dtype)
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--ckpt", required=True)
     parser.add_argument("--codec")
     parser.add_argument("--validation-source")
     parser.add_argument("--ref-model", default="gpt2-large")
-    parser.add_argument("--ref-dtype", choices=("bfloat16", "float32"), default="bfloat16",
-                        help="Legacy default retained; new ICLR runs explicitly use float32")
+    parser.add_argument("--ref-dtype", choices=("bfloat16", "float32", "float64"), default="bfloat16",
+                        help="Reference model and scoring dtype; float32/float64 disable TF32")
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--n-samples", type=int, default=128)
     parser.add_argument("--batch-size", type=int, default=8)
@@ -40,10 +62,6 @@ def main():
     parser.add_argument("--save-all-generations", action="store_true")
     args = parser.parse_args()
 
-    if args.ref_dtype == "float32":
-        torch.backends.cuda.matmul.allow_tf32 = False
-        torch.backends.cudnn.allow_tf32 = False
-
     device = "cuda"
     checkpoint = torch.load(args.ckpt, map_location="cpu", weights_only=False)
     cfg = checkpoint["args"]
@@ -58,7 +76,7 @@ def main():
     else:
         cfg.optimization.stochastic_encode = {"enabled": False}
 
-    set_manual_seed(args.seed)
+    seed_evaluation(args.seed, args.ref_dtype)
     Trainer = hydra.utils.get_class(trainers[cfg.task])
     trainer = Trainer(cfg, device)
     trainer.model.load_state_dict(checkpoint["model"])
@@ -129,18 +147,19 @@ def main():
     torch.cuda.empty_cache()
 
     ref_tokenizer = AutoTokenizer.from_pretrained(args.ref_model)
-    ref_model = AutoModelForCausalLM.from_pretrained(
-        args.ref_model, torch_dtype=getattr(torch, args.ref_dtype)
-    ).to(device).eval()
-    result = generation_perplexity(
-        ref_model,
-        ref_tokenizer,
-        prompts,
-        generations,
-        score_only_generated=True,
-        batch_size=args.batch_size,
-        device=device,
-    )
+    with reference_precision_context(args.ref_dtype):
+        ref_model = AutoModelForCausalLM.from_pretrained(
+            args.ref_model, torch_dtype=getattr(torch, args.ref_dtype)
+        ).to(device).eval()
+        result = generation_perplexity(
+            ref_model,
+            ref_tokenizer,
+            prompts,
+            generations,
+            score_only_generated=True,
+            batch_size=args.batch_size,
+            device=device,
+        )
 
     payload = {
         "checkpoint": str(Path(args.ckpt).resolve()),
diff --git a/scripts/run_iclr_debug.py b/scripts/run_iclr_debug.py
index da25b0c..ce5ae4e 100644
--- a/scripts/run_iclr_debug.py
+++ b/scripts/run_iclr_debug.py
@@ -23,7 +23,15 @@ from utils.iclr_training import stage_config, export_training_checkpoint
 
 def make_plan(args):
     root = args.output.resolve()
-    stages = ["tokenizer", "generator"] if args.variant == "two-stage" else ["decoder-mstok"]
+    stages = {"two-stage": ["tokenizer", "generator"], "generator": ["generator"],
+              "decoder-mstok": ["decoder-mstok"]}[args.variant]
+    external_tokenizer = getattr(args, "tokenizer_checkpoint", None)
+    external_steps = getattr(args, "tokenizer_steps", None)
+    if args.variant == "generator":
+        if external_tokenizer is None or external_steps is None or external_steps <= 0:
+            raise ValueError("Generator-only training requires --tokenizer-checkpoint and positive --tokenizer-steps")
+    elif external_tokenizer is not None or external_steps is not None:
+        raise ValueError("External tokenizer options require --variant generator")
     if args.tokenizer_only and args.variant != "two-stage":
         raise ValueError("--tokenizer-only requires --variant two-stage")
     if args.pilot_steps is not None and not 32 <= args.pilot_steps <= 500:
@@ -32,13 +40,18 @@ def make_plan(args):
     for stage in stages:
         checkpoint, tokenizer_steps = None, None
         if stage == "generator":
-            tokenizer_steps = configs["tokenizer"].iclr_debug.stop_step
-            checkpoint = root / "tokenizer" / f"milestone-iter-{tokenizer_steps}.pt"
+            if args.variant == "generator":
+                tokenizer_steps = external_steps
+                checkpoint = Path(external_tokenizer).resolve()
+            else:
+                tokenizer_steps = configs["tokenizer"].iclr_debug.stop_step
+                checkpoint = root / "tokenizer" / f"milestone-iter-{tokenizer_steps}.pt"
         budget = args.generator_positions if stage == "generator" else args.tokenizer_positions
         if stage == "decoder-mstok":
             budget = args.joint_positions
         configs[stage] = stage_config(stage, args.data_dir, root, target_positions=budget,
-            horizon_positions=args.joint_lr_horizon_positions if stage == "decoder-mstok" else None,
+            horizon_positions=(args.joint_lr_horizon_positions if stage == "decoder-mstok" else
+                               getattr(args, "generator_lr_horizon_positions", None) if stage == "generator" else None),
             pilot_steps=args.pilot_steps, wandb_enabled=not args.no_wandb,
             tokenizer_checkpoint=checkpoint, tokenizer_steps=tokenizer_steps)
         if stage == "generator":
@@ -65,7 +78,7 @@ def latest(root):
     return max(candidates, key=lambda x: x[0], default=(0, None))
 
 
-def generation_evaluation(controller, config, step, directory):
+def generation_evaluation(controller, config, step, directory, reference_dtype="float32"):
     stage_root = Path(config.experiment_dir)
     output = stage_root / "evaluation" / f"step-{step}"
     output.mkdir(parents=True, exist_ok=True)
@@ -78,19 +91,25 @@ def generation_evaluation(controller, config, step, directory):
             controller.run(f"{config.iclr_debug.stage}-eval-{step}-seed{seed}", [sys.executable,
                 "evaluation/eval_random_gen_ppl_checkpoint.py", "--ckpt", directory / "ncp.pt",
                 "--codec", directory / "vqvae.pt", "--validation-source", config.dataset.validate_source,
-                "--ref-model", "gpt2-large", "--ref-dtype", "float32", "--seed", seed,
+                "--ref-model", "gpt2-large", "--ref-dtype", reference_dtype, "--seed", seed,
                 "--n-samples", 128, "--batch-size", 4, "--temperature", 1,
                 "--top-k", 0, "--top-p", 1, "--save-all-generations", "--out", path],
                 path.with_suffix(".log"), env=dict(controller.env, CUDA_VISIBLE_DEVICES="0"), timeout=3600)
         row = json.loads(path.read_text())
-        if (row["reference_dtype"] != "float32" or row["protocol"]["n_levels"] != 16
+        if (row["reference_dtype"] != reference_dtype or row.get("tf32_allowed") is not False
+                or row["protocol"]["n_levels"] != 16
                 or row["checkpoint_step"] != step or row["seed"] != seed
                 or row["requested_samples"] != 128 or row["scored_samples"] != 128):
             raise ValueError(f"Incomplete/mismatched evaluation: {path}")
     controller.run(f"summarize-{step}", [sys.executable, "evaluation/summarize_gen_ppl.py",
         "--input-glob", output / "random-seed-[0-4].json", "--out", output / "summary.json"],
         output / "summarize.log", timeout=120)
-    write_json(output / "EVAL_DONE.json", dict(step=step, reference_dtype="float32", sampling="supplied-level0-untruncated"))
+    summary_path = output / "summary.json"
+    summary = json.loads(summary_path.read_text())
+    summary.update(reference_dtype=reference_dtype, tf32_allowed=False)
+    write_json(summary_path, summary)
+    write_json(output / "EVAL_DONE.json", dict(step=step, reference_dtype=reference_dtype,
+                                              tf32_allowed=False, sampling="supplied-level0-untruncated"))
     (output / "EVAL_FAILED.json").unlink(missing_ok=True)
 
 
@@ -130,6 +149,19 @@ def run(args):
         raise FileNotFoundError("Resume requires this launcher's manifest")
     expected = json.loads((ROOT / "config/iclr-debug/full-owt-audit.json").read_text())
     data_preflight(args.data_dir, expected)
+    tokenizer_provenance = None
+    if args.variant == "generator":
+        import torch
+        checkpoint = Path(configs["generator"].iclr_debug.tokenizer_checkpoint)
+        payload = torch.load(checkpoint, map_location="cpu", mmap=True, weights_only=False)
+        if (OmegaConf.select(payload["args"], "iclr_debug.stage") != "tokenizer"
+                or int(payload["step"]) != configs["generator"].iclr_debug.tokenizer_steps):
+            raise ValueError("External tokenizer checkpoint stage/step mismatch")
+        for key in ("codec", "dataset", "pre_tokenizer"):
+            if OmegaConf.to_container(payload["args"][key], resolve=True) != OmegaConf.to_container(configs["generator"][key], resolve=True):
+                raise ValueError(f"External tokenizer changes {key}")
+        tokenizer_provenance = dict(path=str(checkpoint), step=int(payload["step"]), sha256=sha256(checkpoint))
+        del payload
     root.mkdir(parents=True, exist_ok=True)
     lock = (root / "controller.lock").open("a")
     fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
@@ -141,6 +173,8 @@ def run(args):
                     configs={s: OmegaConf.to_container(c, resolve=True) for s, c in configs.items()},
                     sources=sources, dataset=expected, hf_repo=args.hf_repo,
                     generation_evaluation=not args.no_generation_eval,
+                    evaluation_reference_dtype=args.eval_ref_dtype,
+                    external_tokenizer=tokenizer_provenance,
                     versions={name: version(name) for name in ("torch", "transformers", "numpy", "hydra-core")})
     if args.resume:
         if json.loads((root / "manifest.json").read_text()) != manifest:
@@ -170,7 +204,7 @@ def run(args):
             if args.pilot_steps:
                 # Force a real save/reload/continuation in every stage.
                 targets = [args.pilot_steps // 2, args.pilot_steps]
-            if stage == "generator" and not (root / "tokenizer" / "STAGE_DONE.json").is_file():
+            if stage == "generator" and args.variant == "two-stage" and not (root / "tokenizer" / "STAGE_DONE.json").is_file():
                 raise RuntimeError("Tokenizer stage must finish before generator handoff")
             for target in targets:
                 done, checkpoint = latest(stage_root)
@@ -195,7 +229,7 @@ def run(args):
                     export_training_checkpoint(checkpoint, exports)
                 if not args.no_generation_eval and target in config.training.generation_eval_steps:
                     try:
-                        generation_evaluation(controller, config, target, exports)
+                        generation_evaluation(controller, config, target, exports, args.eval_ref_dtype)
                     except InterruptedError:
                         raise
                     except Exception as exc:
@@ -223,11 +257,16 @@ def run(args):
 
 def main():
     p = argparse.ArgumentParser(description=__doc__)
-    p.add_argument("--variant", choices=("two-stage", "decoder-mstok"), required=True)
+    p.add_argument("--variant", choices=("two-stage", "generator", "decoder-mstok"), required=True)
     p.add_argument("--data-dir", type=Path, required=True)
     p.add_argument("--output", type=Path, required=True)
     p.add_argument("--tokenizer-positions", type=int, default=135_000_000_000)
     p.add_argument("--generator-positions", type=int, default=135_000_000_000)
+    p.add_argument("--generator-lr-horizon-positions", type=int,
+                   help="Generator cosine-decay horizon; hold minimum LR after this many input positions")
+    p.add_argument("--tokenizer-checkpoint", type=Path, help="Frozen tokenizer training checkpoint for --variant generator")
+    p.add_argument("--tokenizer-steps", type=int, help="Expected step of the external tokenizer checkpoint")
+    p.add_argument("--eval-ref-dtype", choices=("float32", "float64"), default="float32")
     p.add_argument("--generator-microbatch", type=int, choices=(32, 64), default=32,
                    help="32x16 (measured fastest) or 64x8; both keep global batch 4096")
     p.add_argument("--joint-positions", type=int, default=135_000_000_000)
diff --git a/tests/test_iclr_training.py b/tests/test_iclr_training.py
index 2bc7aee..d19b507 100644
--- a/tests/test_iclr_training.py
+++ b/tests/test_iclr_training.py
@@ -90,6 +90,51 @@ def test_plan_handoff_and_pilot(tmp_path):
     assert pilot["generator"].training.generation_eval_steps == []
 
 
+def test_generator_only_short_decay_plan(tmp_path):
+    args = SimpleNamespace(output=tmp_path / "fresh-generator", data_dir=tmp_path,
+        variant="generator", tokenizer_only=False, pilot_steps=None, no_wandb=False,
+        tokenizer_checkpoint=tmp_path / "tokenizer-128747.pt", tokenizer_steps=128747,
+        generator_positions=135_000_000_000, generator_lr_horizon_positions=27_000_000_000,
+        joint_lr_horizon_positions=None)
+    configs = make_plan(args)
+    assert set(configs) == {"generator"}
+    cfg = configs["generator"]
+    assert cfg.iclr_debug.tokenizer_checkpoint == str(args.tokenizer_checkpoint.resolve())
+    assert cfg.iclr_debug.tokenizer_steps == 128747
+    assert cfg.training.total_iters == 128747
+    assert cfg.optimization.generator_lr_decay_iters == 25750
+    assert cfg.optimization.generator_warmup_iters == 225
+    assert cfg.optimization.generator_min_lr == 1e-5
+    assert cfg.training.resume_checkpoint is None
+    args.tokenizer_steps = None
+    with pytest.raises(ValueError, match="requires"):
+        make_plan(args)
+
+
+def test_generator_lr_stays_at_minimum_after_short_horizon_and_resume(tmp_path):
+    cfg = stage_config("generator", tmp_path, tmp_path, horizon_positions=27_000_000_000)
+    model = SimpleNamespace(ncp=torch.nn.Linear(1, 1), vqvae=torch.nn.Linear(1, 1).requires_grad_(False))
+    optimizer, scheduler = make_optimizer_scheduler(model, cfg)
+    horizon = cfg.optimization.generator_lr_decay_iters
+    # Exercise the actual scheduler on both sides of the boundary and after reload.
+    scheduler.last_epoch = horizon - 2
+    rates = []
+    for _ in range(4):
+        optimizer.step()
+        scheduler.step()
+        rates.append(optimizer.param_groups[0]["lr"])
+    assert rates[0] > cfg.optimization.generator_min_lr
+    assert rates[1:] == [cfg.optimization.generator_min_lr] * 3
+    opt2, sched2 = make_optimizer_scheduler(model, cfg)
+    opt2.load_state_dict(optimizer.state_dict())
+    sched2.load_state_dict(scheduler.state_dict())
+    for step in (horizon + 100, 51499, 128747):
+        sched2.last_epoch = step - 1
+        opt2.step()
+        sched2.step()
+        assert opt2.param_groups[0]["lr"] == cfg.optimization.generator_min_lr
+
+
 @pytest.mark.parametrize("stage", ["tokenizer", "decoder-mstok"])
 def test_strict_checkpoint_exact_continuation(stage, tmp_path):
     torch.set_num_threads(1)