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)