| |
| |
| |
| |
| @@ -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()), |
| |
| |
| |
| |
| @@ -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) |
| |
| |
| |
| |
| @@ -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) |
|
|
