Record fresh generator run: cosine LR decay over 27B then minimum through 135B

iclr-debug-generator-decay27b-20260917/README.md

@@ -0,0 +1,18 @@
+# Fresh generator: decay over 27B, then hold minimum LR
+
+Started September 17, 2026 as a controlled schedule experiment after stopping `iclr-debug-two-stage-135b-20260916`.
+
+The generator and optimizer are initialized fresh with the same seeds as the original run. The frozen tokenizer is the exact same completed tokenizer from step 128,747; its SHA-256 is pinned in manifest.json. Model architecture, dataset, corruption, objective, batch, optimizer parameters, and training precision are unchanged. The generator has not resumed from the old generator checkpoint.
+
+- Total generator budget: 135B input positions, 128,747 updates.
+- Global batch: 4,096 sequences of length 256 (32/GPU × 16 accumulation × 8 GPUs).
+- Warmup: 225 updates, to peak LR 5e-4.
+- Cosine decay: minimum LR 1e-5 reached at update 25,750 (~27B positions, first 20% of budget).
+- LR remains 1e-5 for all remaining updates. Scheduler boundary and resume tests passed.
+- Training: BF16, torch.compile loss scope, fullgraph=True, Inductor/default.
+- Generation evaluation: GPT-2 Large FP64 reference scoring, TF32 disabled; generation remains FP32. Five seeds ×128 samples, true level 0 supplied, temperature 1, no top-k/top-p truncation. First scheduled generation evaluation is at step 12,875 (~13.5B).
+- Milestone checkpoints and evaluations upload automatically to this HF repository.
+
+The source snapshot starts at commit 1b5226d4bbdd34653cb5fdd58d436355fb76ac8c with the accompanying source patch. Runtime source hashes, data checksums, versions, configuration, and tokenizer checkpoint identity are recorded in manifest.json. The isolated checkout must remain unchanged for resume. The launch script supports --resume with otherwise identical arguments.
+
+This tests whether earlier LR decay improves generation quality; no specific PPL target is guaranteed. The original run's best measured Gen PPL was 174.9645 at generator step 12,875. Its latest (45k), best validation (44k), results, and best Gen-PPL checkpoint provenance are archived in HF commit b2c358f9247554fce43ffd2ac8ff02f20170403a.

iclr-debug-generator-decay27b-20260917/generator/config.yaml

@@ -0,0 +1,223 @@
+task: mstok-next-concept
+experiment: iclr-debug-generator-decay27b-20260917-generator
+experiment_dir: /home/ubuntu/mstok-results/iclr-debug-generator-decay27b-20260917/generator
+dataset:
+  train_source: /home/ubuntu/data/full_owt/train_gpt2.bin
+  validate_source: /home/ubuntu/data/full_owt/valid_gpt2.bin
+  pad_token_id: 50257
+pre_tokenizer:
+  name: hf
+  tokenizer: hf
+  model_id: gpt2
+  special_tokens:
+    eos_token: <|endoftext|>
+    pad_token: <|pad|>
+codec:
+  n_layers: 6
+  context_length: 256
+  embed_dim: 384
+  in_vocab_size: 50304
+  vqvae_vocab_size:
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  compression_factor: 1
+  dropout: 0.1
+  pre_quant_groupnorm: 4
+  pre_quant_dropout: 0.2
+  vector_quantizer_config:
+    clss: multiscale_residual_vector_quantizer
+    decay: 0.99
+    epsilon: 1.0e-05
+    commitment_cost: 0.25
+    learned_l1_sampling: true
+    learned_all_sampling: false
+    quant_resi:
+      enabled: true
+      ratio: 0.5
+      share_mode: 0
+      learnable_ratio: false
+    levels:
+      use_manual_levels: true
+      manual_levels:
+      - 1
+      - 4
+      - 9
+      - 16
+      - 25
+      - 36
+      - 49
+      - 64
+      - 81
+      - 100
+      - 121
+      - 144
+      - 169
+      - 196
+      - 225
+      - 256
+    aux:
+      fine_drop:
+        prob: 0.5
+        min_keep: 1
+generator:
+  n_layer: 12
+  n_head: 12
+  bias: true
+  dropout: 0.1
+  n_embd: 768
+  context_length: 256
+  vocab_size:
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  - 16384
+  attn_pattern: block_diagonal
+  use_positional_encoding: true
+  use_level_encoding: true
+  prefix_len: 0
+  use_rope: true
+  rope_base: 10000
+  use_qk_norm: true
+  post_upsample_conv:
+    enabled: true
+    kernel_size: 3
+  use_relu2: true
+  use_flex_attention: true
+  shared_output_head: true
+  shared_head_per_level_bias: false
+  shared_head_per_level_scale: false
+  shared_head_adapter_rank: 0
+objectives:
+  codec_weight: 1.0
+  ncp_weight: 1.0
+  soft_assignment_temperature: 1.0
+  prediction_temperature: 1.0
+  mstok_weight: 0.25
+  residual_weight: 1.0
+  reconstruction_weight: 1.0
+  teacher_ema_decay: 0.999
+  teacher_ema_warmup_steps: 1000
+  mstok_warmup_steps: 500
+optimization:
+  codec_lr: 0.001
+  codec_min_lr: 0.0001
+  codec_warmup_iters: 0
+  codec_lr_decay_iters: 25750
+  generator_lr: 0.0005
+  generator_min_lr: 1.0e-05
+  generator_warmup_iters: 225
+  generator_lr_decay_iters: 25750
+  beta_1: 0.9
+  codec_beta_2: 0.99
+  generator_beta_2: 0.99
+  weight_decay: 0.1
+  max_grad_norm: 1.0
+  level_loss_alpha: 1.0
+  grad_accumulation_steps: 16
+  corruption:
+    mode: per_level
+    per_level_probs:
+    - 0.85
+    - 0.8321428571
+    - 0.8142857143
+    - 0.7964285714
+    - 0.7785714286
+    - 0.7607142857
+    - 0.7428571429
+    - 0.725
+    - 0.7071428571
+    - 0.6892857143
+    - 0.6714285714
+    - 0.6535714286
+    - 0.6357142857
+    - 0.6178571429
+    - 0.6
+    skip_level0: true
+training:
+  log_interval: 10
+  eval_interval: 1000
+  checkpoint_interval: 1000
+  eval_batch_size: 4
+  val_iters: 25
+  keep_last: 3
+  seed: 55
+  codec_initialization_seed: 42
+  generator_initialization_seed: 55
+  total_iters: 128747
+  batch_size: 32
+  expected_world_size: 8
+  expected_global_batch_size: 4096
+  resume_checkpoint: null
+  milestone_steps:
+  - 954
+  - 4769
+  - 12875
+  - 25750
+  - 51499
+  - 77248
+  - 102997
+  - 128747
+  generation_eval_steps:
+  - 12875
+  - 25750
+  - 51499
+  - 77248
+  - 102997
+  - 128747
+torch_compile:
+  enable: true
+  scope: loss
+  mode: default
+  dynamic: false
+  fullgraph: true
+  backend: inductor
+mixed_precision:
+  enable: true
+wandb:
+  entity: mstok
+  project: iclr-debug
+  group: full-owt-ctx256-16sq-135b-v1
+  enable: true
+  id: iclr-debug-generator-decay27b-20260917-generator
+  resume: allow
+  gradients_and_params:
+    enable: false
+    log: all
+    log_freq: 1000
+export:
+  vqvae_config_template: /home/ubuntu/mstok-runs/iclr-debug-generator-decay27b-20260917/config/repro-ctx256/vqvae.yaml
+  ncp_config_template: /home/ubuntu/mstok-runs/iclr-debug-generator-decay27b-20260917/config/repro-ctx256/ncp-sharedhead.yaml
+iclr_debug:
+  stage: generator
+  target_positions: 135000000000
+  lr_horizon_positions: 27000000000
+  stop_step: 128747
+  pilot: false
+  tokenizer_checkpoint: /home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/tokenizer/milestone-iter-128747.pt
+  tokenizer_steps: 128747
+  budget_unit: input positions including padding; log non-padding tokens separately

iclr-debug-generator-decay27b-20260917/git-provenance.json

@@ -0,0 +1,3 @@
+{
+  "commit": "1b5226d4bbdd34653cb5fdd58d436355fb76ac8c"
+}

iclr-debug-generator-decay27b-20260917/launch.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+set -euo pipefail
+cd /home/ubuntu/mstok-runs/iclr-debug-generator-decay27b-20260917
+export OMP_NUM_THREADS=1
+export HF_HUB_DISABLE_PROGRESS_BARS=1
+exec /home/ubuntu/miniconda3/envs/mstok/bin/python -u scripts/run_iclr_debug.py \
+  --variant generator \
+  --data-dir /home/ubuntu/data/full_owt \
+  --output /home/ubuntu/mstok-results/iclr-debug-generator-decay27b-20260917 \
+  --tokenizer-checkpoint /home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/tokenizer/milestone-iter-128747.pt \
+  --tokenizer-steps 128747 \
+  --generator-lr-horizon-positions 27000000000 \
+  --eval-ref-dtype float64 \
+  --hf-repo iskhare/iclr-debug \
+  "$@"

iclr-debug-generator-decay27b-20260917/manifest.json

@@ -0,0 +1,517 @@
+{
+  "format": "iclr-debug-production-v1",
+  "variant": "generator",
+  "configs": {
+    "generator": {
+      "task": "mstok-next-concept",
+      "experiment": "iclr-debug-generator-decay27b-20260917-generator",
+      "experiment_dir": "/home/ubuntu/mstok-results/iclr-debug-generator-decay27b-20260917/generator",
+      "dataset": {
+        "train_source": "/home/ubuntu/data/full_owt/train_gpt2.bin",
+        "validate_source": "/home/ubuntu/data/full_owt/valid_gpt2.bin",
+        "pad_token_id": 50257
+      },
+      "pre_tokenizer": {
+        "name": "hf",
+        "tokenizer": "hf",
+        "model_id": "gpt2",
+        "special_tokens": {
+          "eos_token": "<|endoftext|>",
+          "pad_token": "<|pad|>"
+        }
+      },
+      "codec": {
+        "n_layers": 6,
+        "context_length": 256,
+        "embed_dim": 384,
+        "in_vocab_size": 50304,
+        "vqvae_vocab_size": [
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384
+        ],
+        "compression_factor": 1,
+        "dropout": 0.1,
+        "pre_quant_groupnorm": 4,
+        "pre_quant_dropout": 0.2,
+        "vector_quantizer_config": {
+          "clss": "multiscale_residual_vector_quantizer",
+          "decay": 0.99,
+          "epsilon": 1e-05,
+          "commitment_cost": 0.25,
+          "learned_l1_sampling": true,
+          "learned_all_sampling": false,
+          "quant_resi": {
+            "enabled": true,
+            "ratio": 0.5,
+            "share_mode": 0,
+            "learnable_ratio": false
+          },
+          "levels": {
+            "use_manual_levels": true,
+            "manual_levels": [
+              1,
+              4,
+              9,
+              16,
+              25,
+              36,
+              49,
+              64,
+              81,
+              100,
+              121,
+              144,
+              169,
+              196,
+              225,
+              256
+            ]
+          },
+          "aux": {
+            "fine_drop": {
+              "prob": 0.5,
+              "min_keep": 1
+            }
+          }
+        }
+      },
+      "generator": {
+        "n_layer": 12,
+        "n_head": 12,
+        "bias": true,
+        "dropout": 0.1,
+        "n_embd": 768,
+        "context_length": 256,
+        "vocab_size": [
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384,
+          16384
+        ],
+        "attn_pattern": "block_diagonal",
+        "use_positional_encoding": true,
+        "use_level_encoding": true,
+        "prefix_len": 0,
+        "use_rope": true,
+        "rope_base": 10000,
+        "use_qk_norm": true,
+        "post_upsample_conv": {
+          "enabled": true,
+          "kernel_size": 3
+        },
+        "use_relu2": true,
+        "use_flex_attention": true,
+        "shared_output_head": true,
+        "shared_head_per_level_bias": false,
+        "shared_head_per_level_scale": false,
+        "shared_head_adapter_rank": 0
+      },
+      "objectives": {
+        "codec_weight": 1.0,
+        "ncp_weight": 1.0,
+        "soft_assignment_temperature": 1.0,
+        "prediction_temperature": 1.0,
+        "mstok_weight": 0.25,
+        "residual_weight": 1.0,
+        "reconstruction_weight": 1.0,
+        "teacher_ema_decay": 0.999,
+        "teacher_ema_warmup_steps": 1000,
+        "mstok_warmup_steps": 500
+      },
+      "optimization": {
+        "codec_lr": 0.001,
+        "codec_min_lr": 0.0001,
+        "codec_warmup_iters": 0,
+        "codec_lr_decay_iters": 25750,
+        "generator_lr": 0.0005,
+        "generator_min_lr": 1e-05,
+        "generator_warmup_iters": 225,
+        "generator_lr_decay_iters": 25750,
+        "beta_1": 0.9,
+        "codec_beta_2": 0.99,
+        "generator_beta_2": 0.99,
+        "weight_decay": 0.1,
+        "max_grad_norm": 1.0,
+        "level_loss_alpha": 1.0,
+        "grad_accumulation_steps": 16,
+        "corruption": {
+          "mode": "per_level",
+          "per_level_probs": [
+            0.85,
+            0.8321428571,
+            0.8142857143,
+            0.7964285714,
+            0.7785714286,
+            0.7607142857,
+            0.7428571429,
+            0.725,
+            0.7071428571,
+            0.6892857143,
+            0.6714285714,
+            0.6535714286,
+            0.6357142857,
+            0.6178571429,
+            0.6
+          ],
+          "skip_level0": true
+        }
+      },
+      "training": {
+        "log_interval": 10,
+        "eval_interval": 1000,
+        "checkpoint_interval": 1000,
+        "eval_batch_size": 4,
+        "val_iters": 25,
+        "keep_last": 3,
+        "seed": 55,
+        "codec_initialization_seed": 42,
+        "generator_initialization_seed": 55,
+        "total_iters": 128747,
+        "batch_size": 32,
+        "expected_world_size": 8,
+        "expected_global_batch_size": 4096,
+        "resume_checkpoint": null,
+        "milestone_steps": [
+          954,
+          4769,
+          12875,
+          25750,
+          51499,
+          77248,
+          102997,
+          128747
+        ],
+        "generation_eval_steps": [
+          12875,
+          25750,
+          51499,
+          77248,
+          102997,
+          128747
+        ]
+      },
+      "torch_compile": {
+        "enable": true,
+        "scope": "loss",
+        "mode": "default",
+        "dynamic": false,
+        "fullgraph": true,
+        "backend": "inductor"
+      },
+      "mixed_precision": {
+        "enable": true
+      },
+      "wandb": {
+        "entity": "mstok",
+        "project": "iclr-debug",
+        "group": "full-owt-ctx256-16sq-135b-v1",
+        "enable": true,
+        "id": "iclr-debug-generator-decay27b-20260917-generator",
+        "resume": "allow",
+        "gradients_and_params": {
+          "enable": false,
+          "log": "all",
+          "log_freq": 1000
+        }
+      },
+      "export": {
+        "vqvae_config_template": "/home/ubuntu/mstok-runs/iclr-debug-generator-decay27b-20260917/config/repro-ctx256/vqvae.yaml",
+        "ncp_config_template": "/home/ubuntu/mstok-runs/iclr-debug-generator-decay27b-20260917/config/repro-ctx256/ncp-sharedhead.yaml"
+      },
+      "iclr_debug": {
+        "stage": "generator",
+        "target_positions": 135000000000,
+        "lr_horizon_positions": 27000000000,
+        "stop_step": 128747,
+        "pilot": false,
+        "tokenizer_checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/tokenizer/milestone-iter-128747.pt",
+        "tokenizer_steps": 128747,
+        "budget_unit": "input positions including padding; log non-padding tokens separately"
+      }
+    }
+  },
+  "sources": {
+    "config/repro-ctx256/alignment-generator.yaml": "f62f36d078a3a373a19ea716a3ebc7eda5ae2fb7d670c5b935c734826bb76895",
+    "config/repro-ctx256/alignment-joint.yaml": "bb511950d20335e4237e406464ab26760a7fd9c9c9472c0507addd88d33d5178",
+    "config/repro-ctx256/alignment-tokenizer.yaml": "2d32955296f7d15990f284328c1cf8f52071abdba28033a65e272bb49cf5dd56",
+    "config/repro-ctx256/alignment.yaml": "3a486fa2ed28eca4fc406232669399b0625d274376ac43271f4ccf187b1e9077",
+    "config/repro-ctx256/joint.yaml": "8b4889111687a52e2d7d71995bda45a7f01f2777674c88191f338dca4c1732d2",
+    "config/repro-ctx256/mstok-semantic-eostok-matched10ep.yaml": "ae8a72c0c9f1cb5d2e301d83a1cdc08985a1f895bf627a5f75b1b3ce627f5e58",
+    "config/repro-ctx256/mstok-semantic-eostok.yaml": "2c9c48b21325c0e479990701b28116d78f611101724846747d173e52145fe15a",
+    "config/repro-ctx256/mstok-semantic-gear-matched10ep.yaml": "164cb8ac39fa407bb337f559ad8a9c62a34708fb900b8829d3ccf182961f2d0c",
+    "config/repro-ctx256/mstok-semantic-gear.yaml": "8cb8c998551daef0bf7fd9ec77b13b34c06662028d8820767e1432b5a9e8be54",
+    "config/repro-ctx256/mstok-semantic.yaml": "e0b00dbfa3eb838103f657c19bde20e36375c7fbe692623452ea093003f1ac98",
+    "config/repro-ctx256/mstok-w1-pilot.yaml": "21ee22eeccaaea2563c92ccc84f6df3871ea5da6a5ce980656088ea577dc1438",
+    "config/repro-ctx256/mstok-w1.yaml": "b6bf6079114b3c997fab89e85a53318bb2ec5dae55678513599d84152eab701b",
+    "config/repro-ctx256/mstok.yaml": "d9bb216f4eeb58857d72be399d5f8847d197ce70d0e504879f0788045483dc87",
+    "config/repro-ctx256/ncp-sharedhead.yaml": "86c9efbcb2b4f74b7253aa3bc4393de4a6a05e20ea30da573d1ca094ec600645",
+    "config/repro-ctx256/ncp.yaml": "93663654c4ce5ad0931db052edf507fdcb7c0ff9fc9722f4018820deee5c5acf",
+    "config/repro-ctx256/substitution-generator.yaml": "064f55887b9b2b8f3b0b52c094d4a193510dc1148ad7a5c7f53e317971f87c7c",
+    "config/repro-ctx256/substitution-joint.yaml": "b1576f4ecfe6fc6029a0cd868271b80a1d988b7f3667c6b1a8d24a6a4015df3b",
+    "config/repro-ctx256/substitution-tokenizer.yaml": "e29be5046f0fafbfe4c06973be4cd036a004e845348911cc0e6e05a7fca5d425",
+    "config/repro-ctx256/substitution.yaml": "ebba1c7e664de0e84913bdfd38e794837781c42c1b48cde90135ff9417856a95",
+    "config/repro-ctx256/vqvae.yaml": "05fde46e57c3e6b411e8a6afcac251f178e4b8a0856a0b99332bb284fb4710ce",
+    "data/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "data/serialized_dataset.py": "5622c2d94236cd1b9b4ca18e8d5c272a86b54d193d99048fb9b7cd2687e8dbc6",
+    "data/tinysentences/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "data/tinysentences/clean.py": "918a9a91aae484a5e06d6d212c2c0974946a4032e061c62a7b65851735bfc8dd",
+    "data/tinysentences/clean_parsed.py": "009a087f967e171bc6098484dbfd58635302900f6b07629629898f1d00fa70a3",
+    "data/tinysentences/generate_meta.py": "662714f1370e05567d412add168a7c3a05f4b1c9adb047fad374900bc1936a7e",
+    "data/tinysentences/generate_metadata_char.py": "afd3336ea03ff0653ccaf41c121d9b93debf7a380db6cd881b62b308ee6ee68c",
+    "data/tinysentences/generate_training_data_word.py": "82e2edcb7ea4db0609ccf2b047a0b207e7b8fcff4fc0b4ceee521ff278e8a8e8",
+    "data/tinysentences/parse_summaries.py": "2317159364dbbe7206f252d2939cab5d375395e0a536be0a1f60c5df88b31702",
+    "data/tinysentences/prepare.py": "35263b287f93337fda4207ed1991e9e8a820ab75139411bb7dd5f0a8a25bcbd9",
+    "data/tinysentences/sentences_only.py": "5e4d72bf612df27dad46cf606a40402606d897043ae24274dc2dfa88ec0b3e79",
+    "data/tinysentences/stats.py": "2665e9e6fac4a8f6c3e5bd8b07cbfd941ba6359a43981d89bdef58c136fdd3a9",
+    "data/tinysentences/summarize/merge_summaries_char.py": "bce94f4d9a824c2d25e45448023d59b7fbc42aba22b33efbc067cdbba93e1934",
+    "data/tinysentences/summarize/merge_summaries_word.py": "f6079110332cc38b2dd3571c30bedcca97098c39498f45eec1ef6e60cce7ac81",
+    "data/tinysentences/summarize/summarize_gemini.py": "282a438efce0dee351fd9833a59c56e4480bed4314dae40844201794ada3bd96",
+    "data/tinysentences/summarize/summarize_gpt.py": "fc08cbc37d6161e5b0057a1ede43272d96f7274524917d6c4d42c22f505c6e09",
+    "data/tinysentences/summarize/summarize_llama.py": "7f3f8230652a5b88ea2856b273d497cd5b633580b430c909e4f996cd0232832d",
+    "data/tinysentences/summarize/summarize_qwen.py": "b07e3c226ab0fdf5b21b025263adfa7c2625d640f0fac3c787bc4a2c9e32107f",
+    "data/tinysentences/summarize/summarize_together.py": "8c26ea052f2a64c2653d39f1ea5c13983dfbc9547f21cd7d7ea9d9b39033bffe",
+    "data/tinysentences/test.py": "9749244da77c078f550a78751664b45bb92067682f3aa27c2338419f1186c844",
+    "data/tinysentences/test_regex.py": "46652e189a0d5f5c0c7a2c90b857082004bf6727ef1d3ad42b631ac7bb0aa88f",
+    "data/tinysentences/tinysentences_char_dataset.py": "152ab9c30aec669a9db14df9e90412597870105eff9cfb737b876aad2323a206",
+    "data/tinysentences/tinysentences_dataset.py": "f90251d21871f6a7ad5ebd2ad0a6e59cd0b02ca245f7e46f93e6d970bf657bc7",
+    "data/tinysentences/tinysentences_vae.py": "30f7adde5413348faf5066448706747fe26a7fedd510ffea720aa0bf8283f6f8",
+    "data/tinysentences/to_csv.py": "c06c74648ed95b4494835c631c94fe1a00a9ff47a272f5e38546a7f94ff9be8b",
+    "data/tinysentences/tokenize/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "data/tinysentences/tokenize/common.py": "e877c74df2a48d15991492b126cd0c5fe26731d61f3aa36d83ebf2cb9bda34a0",
+    "data/tinysentences/tokenize/tokenizer.py": "c16abc50b92a163f6978542b36b03c9a14aa00153c73117c4d52084d5de25c2a",
+    "data/tinysentences/tokenize/train_bpe.py": "fa8e55978dd5ca90117200648b468fb22822bf6981f15f0394c815e72eae1bdb",
+    "data/tinystories/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "data/tinystories/download.py": "5afce37e912f6965b6722f0bc028076a655fea78304f1a213e4862dbe223cdcb",
+    "data/tinystories/stats.py": "216650fe19b78e8431c4512ff87a5c84af64d9e56ab8fdaf3952ca035725ea1b",
+    "data/wikitext/parse_wikitext.py": "ab583e0328f82afc3ba242410e8076be42ef1a16a51d66f3eb46a713cca591d0",
+    "evaluation/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "evaluation/eval_random_gen_ppl_checkpoint.py": "edf98c0975f54af690fff4b5448f7cbf1c5b95fae5e85421926a663a26dc4c0a",
+    "evaluation/gen_ppl.py": "e6853b14d27f93407505eb7607a1bd9fed49a7add1511ceca2b4a3076d533859",
+    "evaluation/gen_ppl_ncp.py": "96d68adee372e5186d0470659171eb03829e3d37f611e8f5917b45fb61e3a4ad",
+    "evaluation/gen_ppl_ntp.py": "b2638622bc041fa88e9d538be4229b3723ebec19c5b99fd6e5b73db4776dc952",
+    "evaluation/guess.py": "c27803225dd21b2deff1e7e494b9e3666180b93776340164ce6e953dacc50656",
+    "evaluation/ppl.py": "f3b581cb4849b0b2aa518729b8338838c7055697c3a26028dedb819c335ac838",
+    "evaluation/run_genppl_multiseed.py": "f8c9ac982f4caa4335b5f4d1a22d70f56dc11a706106a3dfb31d344eb8c93c2e",
+    "evaluation/run_genppl_release_ts.py": "d1bc5c00089118d9cf6c05c77459329ad40c2f8bc347b5aeedda8f5318de8796",
+    "evaluation/run_genppl_sweep.py": "ade4f39e373125848fdd32938e8d7bf86a0ec5108cca9bd75f0a966d092a798e",
+    "evaluation/show_8k_samples.py": "30ad825d5a24375e178750dfaaac45bc486bc3587f2e7d6d305d53bf37f8d226",
+    "evaluation/show_low_sample.py": "d696f1e54ca67e87565dce5864e0e6e2ad4194108cfb1c9c1e06eebd97008da6",
+    "evaluation/summarize_gen_ppl.py": "307f3d0d85884879295d5268aaf7f399f654b1b330aaa5728d6824ad24c6d15a",
+    "evaluation/sweep_gen_ppl_checkpoints.py": "3a9141e27d84b49c5493194dfc5a70bfa0a915cdf5ceb7a46afcf27cc6f33acb",
+    "evaluation/sweep_gen_ppl_ncp.py": "d91c2959cfed51d224ccb1c0d0f2c3b9a35c4027cc9ab8c65aa7603d587de761",
+    "evaluation/sweep_gen_ppl_ncp_scaling.py": "05249c1e00c5eb05042b6ea5b1ca008ecb8f0e4a77537ee755a27c153c25a17c",
+    "evaluation/sweep_gen_ppl_ncp_yolo.py": "4b26ef25b27af2f7fcbcb5ff3c312d47d2a4746ba759155d28802e978a681755",
+    "evaluation/sweep_gen_ppl_ntp.py": "14a04074741ee44580ec6557870d42534dd172d26c5be3e943ed5baf1b727303",
+    "evaluation/sweep_random_llama_yolo.py": "4062b54289fcd9407d15feb1ffd241b8722fd717c700350c22472fecec5ae46c",
+    "evaluation/val_oracle_ppl.py": "1a23f3c2897b6a6a3d239ed533979de35352d303bae094f4006bef4156f935c0",
+    "evaluation/val_vqvae_ppl.py": "4c293c502d3faf982d4ad66566b7a037faa7d28191d929de1621cb67f5ae8b5f",
+    "models/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "models/common.py": "99aa0da356963a282936413d738ee1da72a3e3f277fecc28eddb553051582a4b",
+    "models/masked_lm.py": "80279614d74efc7b4c25306789d74e2f71045c54e4e4f485ee365f664705d7f6",
+    "models/next_concept.py": "743667aac3bdceb7bb7476a580700185b2638ddd9575e6357db57acc7d699339",
+    "models/next_token.py": "d0d292b9f5c032a351cbca1a8f0a6d41c3062890cb3e3f5b67d4e23a7cffeb77",
+    "models/quant.py": "b7b73abd4e55e10454cb680dea82124c38c222d37236c4c76b5c7d21123ac1b1",
+    "models/semantic_input.py": "b06658d4bc641397cbb973255eadfd2d6087fde244d801d3c70d29e3ce841bf2",
+    "models/vae.py": "02e2213a8878920374baf27dd10dfc03847dce866557c97e816c427d230c9ded",
+    "models/vqvae.py": "8eb3503eb1d990121b7d6138a4f8216aecfc72246d4a92bbecf98025c8a3d504",
+    "scripts/benchmark_iclr_debug.py": "dba86e8d578eec2188c633a479ccf673d58f3e5580b9074869cae1a352b8a043",
+    "scripts/benchmark_mstok_compile.py": "aad909da6f05ba25ac3d8f6b3cf450103bed55d0b93b1490764c0ef54f35b606",
+    "scripts/eval_owt_ckpts_gpt2_large.py": "4977a2d614698f60032fb6fd7fef614eb503d5b3883f499146a4fc048214a509",
+    "scripts/evaluate_mstok_sidecar.py": "9ba33e865fa0fad34bd2ab375d7fe38f37c441e2d96348f0e389e29dd9643455",
+    "scripts/evaluations/sweep_gen_ppl_shared_vocab_ncp.py": "2311ac71f9beba282426faa21390a5898c7d183a95f529755d6e1d6e8fc24050",
+    "scripts/neurips/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "scripts/neurips/eval_gen_ppl_ncp_owt.py": "b92eff1ce198eb4ac8c88da7113b41d095c12acad92db541685175993db3535c",
+    "scripts/neurips/sweep_gen_ppl_owt_gpt2.py": "02adb1629cfa40ff51a75152e4da9c9ee00ead42b7ecda9e78cf7101e6826a67",
+    "scripts/preflight_owtsmall_mstok_ctx256.py": "a4f80921435b6cd136d089ec39bfa5442359d0dffbeefaaa1019178b8364657a",
+    "scripts/prepare_full_owt.py": "b1cb6b74397625c2b0643c4ee744ca3c89dfa6bbb7eeb7e9681e3a36f91a2c24",
+    "scripts/r1_lab/diagnose_at_40k.py": "b1582dcc29ccebc9255eaa77a7d56a7dbdef589ec5d5a9b0af150dede98874ab",
+    "scripts/r1_lab/inspect_summaries.py": "4154e8a6c1b3f68df5c623a57a1a7632b594c46c5f6d1d85cb08c476d6aa9cea",
+    "scripts/r1_lab/overnight_pipeline.py": "53f4e03ce1a5fb43fbca3ae455288f45ae8350b0a16f3d5318763011de6a8f96",
+    "scripts/r1_lab/prime_rope_qknorm_score_cache.py": "a791cdec03b415bfccdfd8d89c7e659aff8a2a5696d16ab7b4350ff608d0391b",
+    "scripts/r1_lab/sweep_gen_ppl_r1lab.py": "c7c0acf9aa6f07bb0a985cf0d5e5692c696e658ebb0470b7c92d9332565f9c06",
+    "scripts/r1_lab/sweep_gen_ppl_r1lab_v15.py": "524aa8680aa1b9fea2689509848c19e5a74a3ad72e78b0388bbca7acd43be22e",
+    "scripts/r1_lab/sweep_gen_ppl_rope_qknorm.py": "7822df057355b0a3a71ea5145742107c6e30f906859096237ff61d765fe36d50",
+    "scripts/r1_lab/train_variant.py": "9d0a2091784609ac9728226afd82f56fd4c1ac2dc08e89e1b17b60d34bb3d55e",
+    "scripts/run_alignment.py": "6312cd9d0976899763589fdcbd258c7681f6e9e96d7396cf35ff022d4695b862",
+    "scripts/run_iclr_debug.py": "440fd11c54e1f0c30c01b2555fc4109e98b296fe564ebabb82510ae3b64927a2",
+    "scripts/run_mstok_semantic.py": "8a7a38a50dd6694950bfa0e7b79718ec653c7e70be99fe13cbf2ee05c28f46a0",
+    "scripts/run_mstok_w1_pilot.py": "35224e462d0d11eb3e4326b0daa8bad7c99e1e30cdc6383310f8e2c8b5bb708f",
+    "scripts/run_substitution.py": "2cd22eaa58dab8b5231984eee340a3a2bdd0b88cc83333f7133447471e7f73f1",
+    "scripts/smoke_test_masked_l0.py": "d24f1270b3219a31212b68bcbb786f649c7c72f96e53f602d87ad6c17f3f73e5",
+    "scripts/test0_diagnostics.py": "ed389c666beef015ccac004f314a5a1428def303b46ac0e442f0236337344549",
+    "scripts/test0_r1_similarity.py": "4a486878bebd418a3b5454ddbf02585726ccbacc2700fdf591aa524766f9c9c9",
+    "scripts/test0_r1_uniqueness.py": "3ee59d4397d5b268151dbb8ca779c7f9bf8461d129d97fcd1a517c48ec2a69a1",
+    "scripts/test0_residual_norms.py": "3c4f68e5960e84796dc0bca4fd4583bf3b598552654f9b565fd5fcabd420f9e1",
+    "scripts/test_gen_equiv.py": "ae0e9c513422ee58bdc050ed6218ee54fc555dfb15ea8968c9ffc62a7a7ea99a",
+    "scripts/training_budget.py": "bb94c19c6dfbdd542946b10476b8d6958b8a189abd141cf62288fdb5fa9a2d38",
+    "scripts/vqvae_utilization.py": "5f90327f1507a5042b2f77b119fe4fc0df4e2131e5bbc21069dffd1fadc059c5",
+    "scripts/xsum/codec_ceiling.py": "ddb751e9152d5c88966e588d89cdceee4c5d91c4139aaedf369f65a45ecf2276",
+    "scripts/xsum/eval_merge.py": "6e313d6ed25227d0418b7422d58e09ceda4fc6dfae7d0fd663417b306241135c",
+    "scripts/xsum/eval_q0.py": "82fe04dccce28d7fc9c6ec8b4c9c91404a49527f9437913f44f51a6853ba9abc",
+    "scripts/xsum/eval_shard.py": "07192f6ba6219da1e8123ceeb99748c86da85322859ed5fa25cebcc3816f7ea8",
+    "scripts/xsum/eval_shard_cfg.py": "95ac4a8e2e3a73a21fa4303c62df80c086f23bcfe5ab34eb07ae90b585052ecb",
+    "scripts/xsum/eval_test.py": "7f92d578c31893e788e5fa7d9eef3098ba5602b3642873ce2ce031342c8be868",
+    "scripts/xsum/eval_uncond_rep.py": "ba4740b29bc091d866e2dfd91881e9a7d5c46a1bfa993d13eb63c0467397ec28",
+    "scripts/xsum/make_warmstart.py": "78793d6ef11fbf658c2837ef12cb843fc4817add8ba73c822b5e408901de3318",
+    "scripts/xsum/mbr_fast.py": "3db6e21d2dcce2e3bcf6f22474535d7eba0232b071b4271d2b172f787aed5c85",
+    "scripts/xsum/mbr_select.py": "107a061014c62395b7f5fd986fba9e55dd738054e2474022005cb7e9e0c9b8fb",
+    "scripts/xsum/mbr_shard_cfg.py": "3e7162375028c391276149eff288e9de40c874a3108ac53de0248198f4d08791",
+    "scripts/xsum/pool_append.py": "e50252f48c4d8d905be0de3eb56ca1d720f8f8545c6fdcb58ed27f0360c4e2ad",
+    "scripts/xsum/prep_xsum_le48.py": "3688ed1abc49d18f75d34f8382ff68a5afbd753b36cefdc0652d0bffb6a5f347",
+    "scripts/xsum/prep_xsum_mix30.py": "4bd47e52637c4b6bca5fd4956c1787b506a7f51b8a5bb0b68f19f728c4d48afb",
+    "scripts/xsum/prep_xsum_mix_split.py": "ec2fe74f08068ae8fda4ed59a86537a1d3f48276d205099bbfe16f8428ca36d5",
+    "scripts/xsum/prep_xsum_pass_uncond.py": "82ea50248d811e5182e63d0b4dfd6f995dd2f47eb963b5bb1c5d09bc533db7c1",
+    "scripts/xsum/prep_xsum_test_all.py": "7389fbdf5612ba524801c5b04109f810b139152ca50101b60acbb9b1e767b256",
+    "scripts/xsum/score_checkpoints.py": "41b412c06cd92cbd4caa6b586da7ddc7e913607a47c1a3aac0ee0a36a20a36be",
+    "tokenization/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "tokenization/bpe_tokenizer.py": "0c8af84868a6aa40f981768f8ab49d28e501d55cce444edb1ad315de53c3445a",
+    "tokenization/build_word_vocab.py": "778f9666f4e7c642dc471773ee81c92c26975db37de1ebfffb605cae211351b5",
+    "tokenization/byte_tokenizer.py": "eb567a1772ae69738884ed9b94330aa39387c49e54d76e4b9752e67674fad810",
+    "tokenization/compression_ratio.py": "ed07ab1cc38d6176362631ce655bc974dd3623b87693f51279d5047e6bc6890f",
+    "tokenization/hf_tokenizer.py": "2ca8fae097f9c640bcb52b4788c37de8e0d9bb2fa971e30a5abb39f90ff97d4c",
+    "tokenization/tokenize.py": "f8ab54338613b6bd84513fdda85be173fd32274c35d97ef39a5c17d584f8108f",
+    "tokenization/tokenize_doc_level.py": "a06c81cf088fb81d2317ed1c3a1417a0e48ae0ea04962824f13f5d3beddc8902",
+    "tokenization/tokenize_vqvae.py": "14730eec3da87f6528e5fbd1166d5da79e10ac1700f7cd0be73bc40716bdebfc",
+    "tokenization/tokenizer.py": "e527d47e4ed8634c673d48a872a07c9cc8828d30988331de2d4bd3b7c21ddcb7",
+    "tokenization/train_bpe.py": "bf4e2a29eeba874024ed95f32fe720ae2320f72cb978361692db73e71be98683",
+    "tokenization/vqvae_tokenizer.py": "d7fa9d595865aa3e0b0d185bfd04721eb297fe5f657c4fca4ce69fc96aca308f",
+    "tokenization/word_tokenizer.py": "a405621bf35ed6cbb29abadd9d1d651e0783e6801fc7645b35e799d5de381ac2",
+    "train_mstok.py": "077b7c99c12daad2bbd986918d0f487411779de05a472af71124098da399b79d",
+    "trainer/alignment_trainer.py": "54f8ee7cca4d407223f2ce5146e6fc561ad748e20f2c702beb12ae040cd290c4",
+    "trainer/corruption.py": "f0f24c94652e6e9cf8e4ef30ae1a302e3232d64e765c931696d454db96e8b57d",
+    "trainer/joint_trainer.py": "8c8ebb55d21b6e223be1f7595c51f1b9257ccac40f8cff28260c1088656e2e49",
+    "trainer/mstok_trainer.py": "90712e396472d5dac9679b62af4f0b62afa6390a839d4d9bd9dae57406fe76a4",
+    "trainer/ncp_trainer.py": "73fc23f78f403dc39800b924dfb18ad42ffeab847fed30e931183c3d3acaee68",
+    "trainer/ntp_trainer.py": "5076d7aa29fac70869e1e309a15af6c4d3a91d6de8e00ac113a586841a9866e1",
+    "trainer/onpolicy.py": "6ff1fe4103d71001afe2d73e574d20c7242381e420929db14020141464ec6821",
+    "trainer/plain_two_stage.py": "4c5973ed87100b1b478aa1df0bb16cd3d1eb5844dffe5894003dd67c463f6c9a",
+    "trainer/scheduled_sampling.py": "f53515e23c2cec1231265f21a341ff77d11fa8a6129e0e382d540cc85571f5ed",
+    "trainer/semantic_mstok_trainer.py": "ebc02e80d279a445c99d8f09505cdfaf7632f1a661d6880f09cc1bd0fa90ef30",
+    "trainer/semantic_teacher.py": "7983b5da86d25954ed8267eee03d144c421305967705e0a242f4f025794ceeae",
+    "trainer/substitution_trainer.py": "33c7833b2e8dfabbb9ee56b258708055d3194b8d7453a1b8d93f59b44ff91465",
+    "utils/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "utils/arg_utils.py": "e9eb7971f6273e9630b56fd87a03725b29183a8c80abb014c4c84f8b1bdb6445",
+    "utils/benchmark.py": "ccf436ef72a4e626c7f671d8e63a683057bf502a3852e5a7142410d1dc3e5501",
+    "utils/data.py": "78cce11e777813755cd171ccaa3247f89ab99ac4079e0157e55c72ce8a5a7fe5",
+    "utils/dist.py": "b2849d3ce03a8b1737e3dd61216e76e4a9e82d63ad1ea74cddf72e2285eaca64",
+    "utils/experiment_config.py": "616a70fefaac99ab0d856a1b6c36787a28fb9af8a4c55fc6d4765012fbaaeb38",
+    "utils/iclr_training.py": "10c6857fa2319952fe0df18067730b42e3ff755244e2b059f6c95882f8d27c8e",
+    "utils/logging.py": "901e67515cc65a35f575f05d16a3b9ef4a5b2cb0ea252cd7abbc80ecce9e88e8",
+    "utils/lr.py": "4d1ea2c25340afc4bdde8b582dab36f174e0b2965fa7bafa2a7539019dbe2b7e",
+    "utils/misc.py": "be4a323116a0956a77960d1baf613f268057d6a36023663d2af9f8e4728cbe77",
+    "utils/plot_levels.py": "916b39885ced191960c26c334a3b16f9c99751f6c29c46bd69d14fee414965f3",
+    "utils/registry.py": "f7a44c6d18433100873d5df462b9f08a8bdeb11cf3d92226334c7e5ece318318",
+    "utils/train.py": "43eef988000f3a5e20856653d725337a0fe0dbe276f0cf80e1b200292abe812d",
+    "train_iclr_debug.py": "e0a0adf4a30ef55925f9017de24d434299749332e58a507a6424265f7a5f796e",
+    "config/iclr-debug/README.md": "1c757d2ff37026d08c22324ec44cfa760b0599af62c8ad9b728ada471a8f2905",
+    "config/iclr-debug/benchmark-16sq.yaml": "18c9cef7032f9d3806d8b0c8b20fe0fc7d1ea98ad1d2440e24a67ee873c97f9b",
+    "config/iclr-debug/full-owt-audit.json": "fefa394834f6e9f65944d3b514633c29061071ebe6ca60fce21a6edef9748c1e",
+    "config/iclr-debug/training-defaults.yaml": "3b144b8d53f96fe035e36edca96aac456e6cf21202373a467a3cd875fe86b474",
+    "config/iclr-debug/NCM-BUDGET-AUDIT.md": "e734caeab7edb612de7ff508a8c0ec52445530ab7a35895b5b793d6a7b752e4b",
+    "config/iclr-debug/TRAINING.md": "91555457c5cf6af078075673ce96cb440bc8e22aeae5211c252d94f92ea87ac7",
+    "config/iclr-debug/BENCHMARK.md": "2781ea506212e12a0640a4b0d2a2c881355ec85756ca4014bccb1b1aedccdf59",
+    "config/iclr-debug/HF-MODEL-CARD.md": "de1fe176ebde262eacdf32af7d3f8c6497a764dd98e5a9094dad26de4aade38e"
+  },
+  "dataset": {
+    "dataset": "hazyresearch/ncm-tokenized-datasets",
+    "revision": "951e04517fd53d57521d8e09abd6511a662787d0",
+    "dtype": "little-endian uint16",
+    "eos_id": 50256,
+    "offset_units": "token positions; [0, each EOS position + 1] with final sentinel",
+    "sampling": "length-weighted within-segment windows; pad short segments with 50257",
+    "caveat": "EOS segments are not independently verified source-document identities",
+    "splits": {
+      "train": {
+        "bytes": 18070847356,
+        "tokens": 9035423678,
+        "sha256": "b24b78a07f56974fd0c9855dc16f25b836d2cb172be23eb5d3c650cd6d73f56a",
+        "min_token_id": 0,
+        "max_token_id": 50256,
+        "segments": 8009770,
+        "trailing_tokens": 0,
+        "eos_only_segments": 0,
+        "length_percentiles": {
+          "min": 132.0,
+          "p25": 416.0,
+          "p50": 718.0,
+          "p75": 1249.0,
+          "p99": 7610.0,
+          "max": 131288.0
+        },
+        "segments_shorter_than_256": 673011,
+        "expected_padding_fraction_at_256": 1.5955357567300495e-05,
+        "offsets_status": "created",
+        "offsets_sha256": "7e84564946b8cf8190dbdc1e86bdf2fe8cfd9235a517328c80b4df038ea0028b",
+        "offset_entries": 8009771
+      },
+      "valid": {
+        "bytes": 9186834,
+        "tokens": 4593417,
+        "sha256": "1f2ee17e08327a39d4e2417948d02c7b998d626d2dcb2661197a6e208f58266f",
+        "min_token_id": 0,
+        "max_token_id": 50256,
+        "segments": 3999,
+        "trailing_tokens": 0,
+        "eos_only_segments": 0,
+        "length_percentiles": {
+          "min": 143.0,
+          "p25": 433.0,
+          "p50": 734.0,
+          "p75": 1292.5,
+          "p99": 7122.719999999997,
+          "max": 24868.0
+        },
+        "segments_shorter_than_256": 363,
+        "expected_padding_fraction_at_256": 1.6064002171981108e-05,
+        "offsets_status": "created",
+        "offsets_sha256": "09a748ce1dcec7de36c77c48e39ccbf5c0293cc073689fa7c5f48bdfb9bb9445",
+        "offset_entries": 4000
+      }
+    },
+    "overlap": {
+      "exact_train_segment_matches": 0,
+      "unique_validation_segments_in_train": 0,
+      "unique_validation_segments": 3999,
+      "near_duplicates_checked": false
+    }
+  },
+  "hf_repo": "iskhare/iclr-debug",
+  "generation_evaluation": true,
+  "evaluation_reference_dtype": "float64",
+  "external_tokenizer": {
+    "path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/tokenizer/milestone-iter-128747.pt",
+    "step": 128747,
+    "sha256": "0078f567fdf3b261aba666f8d591bedc32403b1c626eed9ef05c4f5c23a2cedd"
+  },
+  "versions": {
+    "torch": "2.13.0+cu130",
+    "transformers": "4.44.2",
+    "numpy": "1.26.4",
+    "hydra-core": "1.3.6"
+  }
+}

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

@@ -0,0 +1,304 @@
+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)

iclr-debug-generator-decay27b-20260917/tests/test_evaluation_precision.py

@@ -0,0 +1,50 @@
+import pytest
+import torch
+
+from evaluation import eval_random_gen_ppl_checkpoint as evaluator
+
+
+@pytest.mark.parametrize("dtype, expected_tf32", [("float32", False), ("float64", False), ("bfloat16", True)])
+def test_seed_helper_cannot_override_requested_evaluation_precision(monkeypatch, dtype, expected_tf32):
+    previous = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32)
+    seeds = []
+
+    def seed_with_training_defaults(seed):
+        seeds.append(seed)
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+    monkeypatch.setattr(evaluator, "set_manual_seed", seed_with_training_defaults)
+    try:
+        evaluator.seed_evaluation(17, dtype)
+        assert seeds == [17]
+        assert torch.backends.cuda.matmul.allow_tf32 is expected_tf32
+        assert torch.backends.cudnn.allow_tf32 is expected_tf32
+    finally:
+        torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = previous
+
+
+def test_fp64_reference_causal_mask_and_logits():
+    from transformers import GPT2Config, GPT2LMHeadModel
+
+    previous_dtype = torch.get_default_dtype()
+    with evaluator.reference_precision_context("float64"):
+        config = GPT2Config(vocab_size=8, n_positions=8, n_embd=8, n_layer=1, n_head=1)
+        config._attn_implementation = "sdpa"
+        model = GPT2LMHeadModel(config).double().eval()
+        with torch.no_grad():
+            logits = model(
+                input_ids=torch.tensor([[1, 2, 3], [1, 2, 0]]),
+                attention_mask=torch.tensor([[1, 1, 1], [1, 1, 0]]),
+            ).logits
+        assert logits.dtype == torch.float64
+        assert torch.isfinite(logits).all()
+    assert torch.get_default_dtype() == previous_dtype
+
+
+def test_reference_precision_context_restores_dtype_on_failure():
+    previous_dtype = torch.get_default_dtype()
+    with pytest.raises(RuntimeError):
+        with evaluator.reference_precision_context("float64"):
+            raise RuntimeError("scoring failed")
+    assert torch.get_default_dtype() == previous_dtype