Save step 22000 generator checkpoint and matched FP32/FP64 evaluation

Generator step 22000 (23.069B input positions), matching final tokenizer, resumable state, standalone exports, configs, SHA-256 manifest, 640 saved samples, matched FP32 and FP64 GPT-2 Large scores, and runtime double-precision audit. TF32 disabled.

iclr-debug-two-stage-135b-20260916/generator/backup-manifest-step-22000.json

@@ -0,0 +1,388 @@
+{
+  "repo": "iskhare/iclr-debug",
+  "run": "iclr-debug-two-stage-135b-20260916",
+  "stage": "generator",
+  "checkpoint_step": 22000,
+  "input_positions": 23068672000,
+  "frozen_tokenizer_step": 128747,
+  "git_provenance": {
+    "commit": "1b5226d4bbdd34653cb5fdd58d436355fb76ac8c"
+  },
+  "evaluation": {
+    "checkpoint_step": 22000,
+    "training_input_positions": 23068672000,
+    "fp32": {
+      "protocol": {
+        "sampling": "random",
+        "temperature": 1.0,
+        "top_k": 0,
+        "top_p": 1.0,
+        "n_provided_levels": 1,
+        "n_levels": 16,
+        "document_aware_validation_sampling": true
+      },
+      "reference_model": "gpt2-large",
+      "checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+      "checkpoint_step": 22000,
+      "seeds": [
+        0,
+        1,
+        2,
+        3,
+        4
+      ],
+      "samples_per_seed": 128,
+      "total_scored_samples": 640,
+      "total_scored_tokens": 160526,
+      "per_seed_mean_ppl": {
+        "0": 190.84469984159884,
+        "1": 190.74649366523607,
+        "2": 201.5548702336531,
+        "3": 190.96056600638894,
+        "4": 172.39751630645821
+      },
+      "mean_gen_ppl": 189.30082921066705,
+      "gen_ppl_se": 4.7069720468787,
+      "mean_of_seed_median_ppl": 179.90716247558595,
+      "mean_entropy_nats": 4.472143790612192,
+      "reference_dtype": "float32",
+      "tf32_allowed": false,
+      "training_continued_during_evaluation": true
+    },
+    "fp64": {
+      "protocol": {
+        "sampling": "random",
+        "temperature": 1.0,
+        "top_k": 0,
+        "top_p": 1.0,
+        "n_provided_levels": 1,
+        "n_levels": 16,
+        "document_aware_validation_sampling": true
+      },
+      "reference_model": "gpt2-large",
+      "checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+      "checkpoint_step": 22000,
+      "seeds": [
+        0,
+        1,
+        2,
+        3,
+        4
+      ],
+      "samples_per_seed": 128,
+      "total_scored_samples": 640,
+      "total_scored_tokens": 160526,
+      "per_seed_mean_ppl": {
+        "0": 190.84481863176575,
+        "1": 190.74650059895347,
+        "2": 201.55483121276038,
+        "3": 190.96062394498597,
+        "4": 172.3975097731799
+      },
+      "mean_gen_ppl": 189.3008568323291,
+      "gen_ppl_se": 4.706971216847786,
+      "mean_of_seed_median_ppl": 179.9071451840254,
+      "mean_entropy_nats": 4.472143790612192,
+      "reference_dtype": "float64",
+      "tf32_allowed": false,
+      "training_continued_during_evaluation": true
+    },
+    "matched_samples": true,
+    "total_samples": 640,
+    "absolute_ppl_difference": 2.762166204206551e-05,
+    "relative_ppl_difference_percent": 1.4591411012787603e-05,
+    "previous_step": 12875,
+    "previous_fp32_gen_ppl": 174.96447636439655,
+    "improvement_vs_previous_fp32_percent": -8.193864917134563
+  },
+  "files": [
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/manual-checkpoints/checkpoint-iter-22000.pt",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/checkpoint-iter-22000.pt",
+      "size_bytes": 2565487115,
+      "sha256": "82f0fecddd66930f054cdaa2c2b7f7b21a401cd5358b204d70b7ace5cf5cac83"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp-config.yaml",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp-config.yaml",
+      "size_bytes": 3528,
+      "sha256": "56a31d9d5b68a96f9f3b394e1f93c9624c1d7a68a1911ab469c2340613038c13"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+      "size_bytes": 427469662,
+      "sha256": "14f9b9044729add9e3b0edfa24f38cf35eaa2d5cd362bc01ac5cfe9de6ff7fa9"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae-config.yaml",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae-config.yaml",
+      "size_bytes": 2204,
+      "sha256": "22bbfe0f1dd75ca79938e48f103d258bfd8194aaaf077741b67bb05e01384f10"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+      "size_bytes": 1300864085,
+      "sha256": "038ab43f06069bfd3fde68b551dd4a4e6ead45023d3c93b843aca6db628eb778"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/EVAL_DONE.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/EVAL_DONE.json",
+      "size_bytes": 129,
+      "sha256": "addff2687e757ddda333b8e4a949e7afa13fa03a20a72debf01c53316ec426e4"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/FP32_DONE.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/FP32_DONE.json",
+      "size_bytes": 105,
+      "sha256": "6a556bd691a5ad24d9c5723340e958a1875a3a317cafe28329b1f022c928e159"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-0.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-0.json",
+      "size_bytes": 850,
+      "sha256": "dbf11bb73a79f2972f15feec9ca7e93dd3a59dc5c03f59dc909f2c13ae518560"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-1.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-1.json",
+      "size_bytes": 850,
+      "sha256": "ba8df2ee32012011a57778d2d36698830cec6f7256af6b0b7fc12f5dea03d62a"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-2.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-2.json",
+      "size_bytes": 850,
+      "sha256": "9dd27c02385a30bc17123402837a9fc041888cd5077821172656f5db11a5e391"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-3.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-3.json",
+      "size_bytes": 850,
+      "sha256": "9c65960095df47e4117d40dca2ee50b0438cf14964f08c24dcb6e545fc3e1b03"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-4.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-4.json",
+      "size_bytes": 850,
+      "sha256": "3076b931bdc15ac4485d46cea2288b829228501318313d6103fb426ce2d45b9d"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/comparison.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/comparison.json",
+      "size_bytes": 2475,
+      "sha256": "84381543ba0ce4468ef888334d228487752fc6a0b36da631a3bc00d69c561012"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/eval_random_gen_ppl_checkpoint_fixed.py",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/eval_random_gen_ppl_checkpoint_fixed.py",
+      "size_bytes": 7171,
+      "sha256": "2da5a5a9fa2d367e647396880b4ea2b75fed0168c016e1ad3cff1e36e4a9876d"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/evaluation-precision.patch",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/evaluation-precision.patch",
+      "size_bytes": 3935,
+      "sha256": "98d505975ba90eeebf91f5862ed610fbfc513bd1295a524106d5f8f55529aa4f"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/evaluator-provenance.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/evaluator-provenance.json",
+      "size_bytes": 524,
+      "sha256": "07c26aa62226776391ecf781fdfecee90d747bbefca657054f3076522d1dba30"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-precision-audit.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-precision-audit.json",
+      "size_bytes": 536,
+      "sha256": "ee35feeaed0744dbf5064fb7c0e78c4897287daaa21d4666a8e3e71ed84cd40c"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-0.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-0.json",
+      "size_bytes": 137012,
+      "sha256": "027ae2f4fcd210188b4bac7bf1d4ae75d351c4089eec530adbe0380ca5a54d34"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-0.per-example.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-0.per-example.json",
+      "size_bytes": 7255,
+      "sha256": "d608686013779ae6de3d0a2345772862d8ff182532715635467fa200508f5fde"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-1.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-1.json",
+      "size_bytes": 140043,
+      "sha256": "92aa09be262cc1d8afb05e07f97895ac7ad5f8e4577e284cfe7e7ef0d50c9238"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-1.per-example.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-1.per-example.json",
+      "size_bytes": 7244,
+      "sha256": "6990fb5e425973a0cf7f2eb66965d92d1e68923ce3a608039df3b6b6be93b5f9"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-2.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-2.json",
+      "size_bytes": 141196,
+      "sha256": "1154d58d9552b29809d1cfe025fd3ca18ae4e3cc57294e0c786ebab610fcecfb"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-2.per-example.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-2.per-example.json",
+      "size_bytes": 7251,
+      "sha256": "de30bf6b2a23bc485ac3f1a6119f9e23954ba532bf9d138ac58862e6427d3aef"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-3.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-3.json",
+      "size_bytes": 136263,
+      "sha256": "3ade0d0687269503cff6725558d1f3f710841973621389ccc28a50cc27637ea9"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-3.per-example.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-3.per-example.json",
+      "size_bytes": 7247,
+      "sha256": "a9b4dcd5d83dc97762ae18d4361ffa44f7cb8769eaf93d5174e5179b1d0e1b3e"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-4.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-4.json",
+      "size_bytes": 139383,
+      "sha256": "490e7c961e43b6e3bdcf08e9486283b3de9165f00d12e6ec352a668bcb327d5c"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-4.per-example.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-4.per-example.json",
+      "size_bytes": 7255,
+      "sha256": "758bad6af54fcdf15b6d5c1cf668a6db89ce96a4dffdbd1d81627ed5f5ef5414"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+      "size_bytes": 722,
+      "sha256": "629ebde69c8bc5fb5b371d29438276e04f1cbaab1e1d1a3615a00df03806c039"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.json",
+      "size_bytes": 136525,
+      "sha256": "6b87bd3021c609ab2b49684d1f673459d5ab58a17f3e5447419fc7a8d9e22911"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.memory.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.memory.json",
+      "size_bytes": 129,
+      "sha256": "2a041635115ba0e300b79e9ad414309ca0b95951ecc602992188a05c021ca786"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.samples.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.samples.json",
+      "size_bytes": 136906,
+      "sha256": "3d46bf39ed7789b32458286ecfe257a360455c7d90d42b6627a061f56e4e9bae"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.json",
+      "size_bytes": 139558,
+      "sha256": "fe7575e3c007e2a96ccdc02d8827510c78cc6ea0dde779c8c7a35a1e7c85f2ae"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.memory.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.memory.json",
+      "size_bytes": 129,
+      "sha256": "9491308eec7b90cb306d3a9ee6e56dd4a925850a6d6d8074c0485f4bbddd6005"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.samples.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.samples.json",
+      "size_bytes": 139939,
+      "sha256": "e01ea2e9b5fea94eb80116c4458bc0d92f82a042459cfaf56b1276d49244a3bb"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.json",
+      "size_bytes": 140708,
+      "sha256": "1decd6032120c40609ce85d0bfa45b921fca0828e912d5bd1f3a637e55144e44"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.memory.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.memory.json",
+      "size_bytes": 129,
+      "sha256": "843b377468ca3f1a4b8712d237b74a573ec6701b1a7ac011fdbb19697c26b5e8"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.samples.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.samples.json",
+      "size_bytes": 141089,
+      "sha256": "94fe2d07dfa64af765b388e5719a9d8ad9fcea7dd598f931144fe38da2c69ca7"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.json",
+      "size_bytes": 135776,
+      "sha256": "7d5fbb842f5c894b475530c0cc2ce9e89a9c6c0bd9d1211396acbfa4cb3a4839"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.memory.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.memory.json",
+      "size_bytes": 129,
+      "sha256": "843b377468ca3f1a4b8712d237b74a573ec6701b1a7ac011fdbb19697c26b5e8"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.samples.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.samples.json",
+      "size_bytes": 136157,
+      "sha256": "7a196ec56f69998cbe6d496485f768c3ce42b3febda75ccb03ed7526e2a4813a"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.json",
+      "size_bytes": 138897,
+      "sha256": "0b94a48e1119c17a7c02c535de4e9a3463704c9a938d91e0f063ecf86ed65960"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.memory.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.memory.json",
+      "size_bytes": 129,
+      "sha256": "843b377468ca3f1a4b8712d237b74a573ec6701b1a7ac011fdbb19697c26b5e8"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.samples.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.samples.json",
+      "size_bytes": 139278,
+      "sha256": "ff8ddd0ee68caac99cf3c97da9fcb612638217c792a31e55e122e9613e4b7c24"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/rescore_fp64.py",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/rescore_fp64.py",
+      "size_bytes": 7246,
+      "sha256": "a0a8f7f2e4188c384ed8cbf0d8769f0c018b913027b8dc21da719556bb7ea624"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/run_evaluation.py",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/run_evaluation.py",
+      "size_bytes": 2695,
+      "sha256": "31cb1de774da3fcc09e4abfba080ce437dd5cbee95358ab02248f7c5872e708b"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/summary-fp64.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/summary-fp64.json",
+      "size_bytes": 966,
+      "sha256": "2266e4aec6d259d80c5dad3fffdc82a0e45c13377a80e26dd9e9a79ca68321da"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/summary.json",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/summary.json",
+      "size_bytes": 966,
+      "sha256": "5a813918ceffe8d1834d8b8311568cd4164d297bc8841ad29269d9796776e661"
+    },
+    {
+      "local_path": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/config.yaml",
+      "repo_path": "iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/training-config.yaml",
+      "size_bytes": 4572,
+      "sha256": "b72db37de5c634702e488e7f582115543c1ed551a01e1a4cef078a6d9dff2cc8"
+    }
+  ]
+}

iclr-debug-two-stage-135b-20260916/generator/checkpoint-iter-22000.pt

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

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/EVAL_DONE.json

@@ -0,0 +1,9 @@
+{
+  "step": 22000,
+  "reference_dtypes": [
+    "float32",
+    "float64"
+  ],
+  "matched_samples": true,
+  "total_samples": 640
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/FP32_DONE.json

@@ -0,0 +1 @@
+{"step": 22000, "reference_dtype": "float32", "sampling": "supplied-level0-untruncated", "manual": true}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/README.md

@@ -0,0 +1,16 @@
+# Generator step 22,000: matched FP32 and FP64 evaluation
+
+This checkpoint has processed 23,068,672,000 input positions (23.069B). It uses the frozen final tokenizer from step 128,747 and the original two-stage training run at commit `1b5226d4bbdd34653cb5fdd58d436355fb76ac8c`.
+
+| GPT-2 Large reference precision | Generation perplexity (mean ± seed SE) |
+| --- | --- |
+| FP32, TF32 disabled | 189.30082921 ± 4.70697205 |
+| FP64, TF32 disabled | 189.30085683 ± 4.70697122 |
+
+The two precisions score the **same 640 saved samples**: 128 samples for each seed 0–4, batch size 4, temperature 1, top-k 0, top-p 1. Generation uses FP32. The first latent level is supplied from document-aware OpenWebText validation samples; the other 15 levels are sampled. This is level-0-conditioned generation, not strictly unconditional. Each seed's PPL is exp(total reference NLL / total scored tokens); the summary reports the mean and standard error across five seeds. 160526 tokens were scored per precision. Mean within-sample token entropy: 4.47214379 nats.
+
+FP64 runtime checks verified floating weights, buffers, leaf-module outputs, final logits, and NLLs use torch.float64; autocast and TF32 were disabled. The isolated FP64 process also used FP64 as its default tensor dtype to avoid an overflow in the installed Transformers causal-mask helper. See `fp64-precision-audit.json` and per-seed sample SHA-256 hashes. FP64 − FP32 mean PPL: 0.00002762 (0.00001459%). FP32 PPL at step 12,875 was 174.96447636; lower PPL is better.
+
+Training continued during evaluation. FP64 peak allocated/reserved GPU memory: 8.055/8.537 GiB. Evaluation used GPU 7 with a 40% allocator cap. The pinned training checkout was unchanged.
+
+The full resumable checkpoint is `../../checkpoint-iter-22000.pt`. Standalone generator and matching tokenizer exports are under `../../exports/step-22000/`. The generator backup manifest includes file sizes and SHA-256 digests. Evaluation scripts and complete saved samples accompany the results.

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-0.json

@@ -0,0 +1,29 @@
+[
+  "/home/ubuntu/miniconda3/envs/mstok/bin/python",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+  "--ckpt",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "--codec",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+  "--validation-source",
+  "/home/ubuntu/data/full_owt/valid_gpt2.bin",
+  "--ref-model",
+  "gpt2-large",
+  "--ref-dtype",
+  "float32",
+  "--seed",
+  "0",
+  "--n-samples",
+  "128",
+  "--batch-size",
+  "4",
+  "--temperature",
+  "1",
+  "--top-k",
+  "0",
+  "--top-p",
+  "1",
+  "--save-all-generations",
+  "--out",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.json"
+]

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-1.json

@@ -0,0 +1,29 @@
+[
+  "/home/ubuntu/miniconda3/envs/mstok/bin/python",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+  "--ckpt",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "--codec",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+  "--validation-source",
+  "/home/ubuntu/data/full_owt/valid_gpt2.bin",
+  "--ref-model",
+  "gpt2-large",
+  "--ref-dtype",
+  "float32",
+  "--seed",
+  "1",
+  "--n-samples",
+  "128",
+  "--batch-size",
+  "4",
+  "--temperature",
+  "1",
+  "--top-k",
+  "0",
+  "--top-p",
+  "1",
+  "--save-all-generations",
+  "--out",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.json"
+]

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-2.json

@@ -0,0 +1,29 @@
+[
+  "/home/ubuntu/miniconda3/envs/mstok/bin/python",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+  "--ckpt",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "--codec",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+  "--validation-source",
+  "/home/ubuntu/data/full_owt/valid_gpt2.bin",
+  "--ref-model",
+  "gpt2-large",
+  "--ref-dtype",
+  "float32",
+  "--seed",
+  "2",
+  "--n-samples",
+  "128",
+  "--batch-size",
+  "4",
+  "--temperature",
+  "1",
+  "--top-k",
+  "0",
+  "--top-p",
+  "1",
+  "--save-all-generations",
+  "--out",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.json"
+]

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-3.json

@@ -0,0 +1,29 @@
+[
+  "/home/ubuntu/miniconda3/envs/mstok/bin/python",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+  "--ckpt",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "--codec",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+  "--validation-source",
+  "/home/ubuntu/data/full_owt/valid_gpt2.bin",
+  "--ref-model",
+  "gpt2-large",
+  "--ref-dtype",
+  "float32",
+  "--seed",
+  "3",
+  "--n-samples",
+  "128",
+  "--batch-size",
+  "4",
+  "--temperature",
+  "1",
+  "--top-k",
+  "0",
+  "--top-p",
+  "1",
+  "--save-all-generations",
+  "--out",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.json"
+]

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/command-seed-4.json

@@ -0,0 +1,29 @@
+[
+  "/home/ubuntu/miniconda3/envs/mstok/bin/python",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py",
+  "--ckpt",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "--codec",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt",
+  "--validation-source",
+  "/home/ubuntu/data/full_owt/valid_gpt2.bin",
+  "--ref-model",
+  "gpt2-large",
+  "--ref-dtype",
+  "float32",
+  "--seed",
+  "4",
+  "--n-samples",
+  "128",
+  "--batch-size",
+  "4",
+  "--temperature",
+  "1",
+  "--top-k",
+  "0",
+  "--top-p",
+  "1",
+  "--save-all-generations",
+  "--out",
+  "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.json"
+]

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/comparison.json

@@ -0,0 +1,87 @@
+{
+  "checkpoint_step": 22000,
+  "training_input_positions": 23068672000,
+  "fp32": {
+    "protocol": {
+      "sampling": "random",
+      "temperature": 1.0,
+      "top_k": 0,
+      "top_p": 1.0,
+      "n_provided_levels": 1,
+      "n_levels": 16,
+      "document_aware_validation_sampling": true
+    },
+    "reference_model": "gpt2-large",
+    "checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+    "checkpoint_step": 22000,
+    "seeds": [
+      0,
+      1,
+      2,
+      3,
+      4
+    ],
+    "samples_per_seed": 128,
+    "total_scored_samples": 640,
+    "total_scored_tokens": 160526,
+    "per_seed_mean_ppl": {
+      "0": 190.84469984159884,
+      "1": 190.74649366523607,
+      "2": 201.5548702336531,
+      "3": 190.96056600638894,
+      "4": 172.39751630645821
+    },
+    "mean_gen_ppl": 189.30082921066705,
+    "gen_ppl_se": 4.7069720468787,
+    "mean_of_seed_median_ppl": 179.90716247558595,
+    "mean_entropy_nats": 4.472143790612192,
+    "reference_dtype": "float32",
+    "tf32_allowed": false,
+    "training_continued_during_evaluation": true
+  },
+  "fp64": {
+    "protocol": {
+      "sampling": "random",
+      "temperature": 1.0,
+      "top_k": 0,
+      "top_p": 1.0,
+      "n_provided_levels": 1,
+      "n_levels": 16,
+      "document_aware_validation_sampling": true
+    },
+    "reference_model": "gpt2-large",
+    "checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+    "checkpoint_step": 22000,
+    "seeds": [
+      0,
+      1,
+      2,
+      3,
+      4
+    ],
+    "samples_per_seed": 128,
+    "total_scored_samples": 640,
+    "total_scored_tokens": 160526,
+    "per_seed_mean_ppl": {
+      "0": 190.84481863176575,
+      "1": 190.74650059895347,
+      "2": 201.55483121276038,
+      "3": 190.96062394498597,
+      "4": 172.3975097731799
+    },
+    "mean_gen_ppl": 189.3008568323291,
+    "gen_ppl_se": 4.706971216847786,
+    "mean_of_seed_median_ppl": 179.9071451840254,
+    "mean_entropy_nats": 4.472143790612192,
+    "reference_dtype": "float64",
+    "tf32_allowed": false,
+    "training_continued_during_evaluation": true
+  },
+  "matched_samples": true,
+  "total_samples": 640,
+  "absolute_ppl_difference": 2.762166204206551e-05,
+  "relative_ppl_difference_percent": 1.4591411012787603e-05,
+  "previous_step": 12875,
+  "previous_fp32_gen_ppl": 174.96447636439655,
+  "improvement_vs_previous_fp32_percent": -8.193864917134563
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/eval_random_gen_ppl_checkpoint_fixed.py

@@ -0,0 +1,182 @@
+"""Evaluate random-sampling generation PPL from a completed NCP checkpoint."""
+
+import argparse
+import json
+import math
+import sys
+from collections import Counter
+from pathlib import Path
+
+REPO_ROOT = Path('/home/ubuntu/mstok-runs/iclr-debug-two-stage-20260916')
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+import hydra
+import torch
+from omegaconf import OmegaConf
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from data.serialized_dataset import SerializedDataset
+from evaluation.gen_ppl import generation_perplexity
+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
+
+
+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", "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)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--top-k", type=int, default=0)
+    parser.add_argument("--top-p", type=float, default=1.0)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--save-all-generations", action="store_true")
+    args = parser.parse_args()
+
+    device = "cuda"
+    checkpoint = torch.load(args.ckpt, map_location="cpu", weights_only=False)
+    cfg = checkpoint["args"]
+    OmegaConf.set_struct(cfg, False)
+    if args.codec:
+        cfg.tokenizer.vqvae.checkpoint_path = args.codec
+    if args.validation_source:
+        cfg.dataset.validate_source = args.validation_source
+    cfg.optimization.token_corruption_prob = 0.0
+    if getattr(cfg.optimization, "stochastic_encode", None) is not None:
+        cfg.optimization.stochastic_encode.enabled = False
+    else:
+        cfg.optimization.stochastic_encode = {"enabled": False}
+
+    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"])
+    model = trainer.model.to(device).eval()
+    tokenizer = trainer.tokenizer
+    levels = list(model.levels)
+    prefix_len = int(cfg.model.prefix_len)
+
+    dataset = SerializedDataset(
+        cfg.dataset.validate_source,
+        int(cfg.model.context_length),
+        cfg.task,
+        tokenizer.pad_token_id,
+    )
+    full_sequences, _ = dataset.get_rand_batch(args.n_samples)
+    prefix_sequences = full_sequences[:, :prefix_len]
+    target_sequences = full_sequences[:, prefix_len:]
+
+    prompts = []
+    generations = []
+    with torch.no_grad():
+        for start in range(0, args.n_samples, args.batch_size):
+            end = min(start + args.batch_size, args.n_samples)
+            prefix = prefix_sequences[start:end].to(device)
+            target_idx = tokenizer.encode_pre_tokenized_idx(
+                target_sequences[start:end].to(device)
+            )
+            level0 = target_idx[:, : levels[0]]
+            generated_idx = model.generate(
+                level0,
+                prefix,
+                temperature=args.temperature,
+                top_k=args.top_k,
+                top_p=args.top_p,
+            )
+            reconstructions = tokenizer.decode_multiscales(
+                generated_idx, skip_special_tokens=True
+            )
+            generations.extend(
+                reconstructions[row][-1] for row in range(len(reconstructions))
+            )
+            prompts.extend(
+                tokenizer.pre_tokenizer.decode_batch(
+                    prefix.cpu().tolist(), skip_special_tokens=True
+                )
+            )
+
+    if args.save_all_generations:
+        # Preserve every generated sample even if scoring fails or some are empty.
+        samples_path = Path(args.out).with_suffix(".samples.json")
+        samples_path.parent.mkdir(parents=True, exist_ok=True)
+        samples_path.write_text(json.dumps(dict(seed=args.seed, checkpoint=args.ckpt,
+            prompts=prompts, generations=generations), indent=2) + "\n")
+    kept = [(prompt, text) for prompt, text in zip(prompts, generations) if text.strip()]
+    prompts = [prompt for prompt, _ in kept]
+    generations = [text for _, text in kept]
+
+    entropies = []
+    for text in generations:
+        ids = tokenizer.pre_tokenizer.encode(text)
+        if ids:
+            counts = Counter(ids)
+            entropies.append(
+                -sum((count / len(ids)) * math.log(count / len(ids)) for count in counts.values())
+            )
+
+    del model, trainer, tokenizer
+    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,
+    )
+
+    payload = {
+        "checkpoint": str(Path(args.ckpt).resolve()),
+        "checkpoint_step": int(checkpoint["step"]),
+        "codec": str(Path(cfg.tokenizer.vqvae.checkpoint_path).resolve()),
+        "reference_model": args.ref_model,
+        "reference_dtype": args.ref_dtype,
+        "tf32_allowed": torch.backends.cuda.matmul.allow_tf32,
+        "protocol": {
+            "sampling": "random" if args.top_k == 0 and args.top_p == 1.0 else "truncated",
+            "temperature": args.temperature,
+            "top_k": args.top_k,
+            "top_p": args.top_p,
+            "n_provided_levels": 1,
+            "n_levels": len(levels),
+            "document_aware_validation_sampling": dataset.doc_offsets is not None,
+        },
+        "seed": args.seed,
+        "requested_samples": args.n_samples,
+        "scored_samples": len(generations),
+        "skipped_empty_samples": args.n_samples - len(generations),
+        "scored_tokens": int(result.token_count.sum()),
+        "mean_ppl": float(result.mean_ppl),
+        "median_ppl": float(result.ppl.median()),
+        "mean_entropy_nats": sum(entropies) / len(entropies),
+        "sample_generations": generations if args.save_all_generations else generations[:8],
+    }
+    output_path = Path(args.out)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(json.dumps(payload, indent=2) + "\n")
+    print(json.dumps({key: value for key, value in payload.items() if key != "sample_generations"}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/evaluation-precision.patch

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

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/evaluator-provenance.json

@@ -0,0 +1,8 @@
+{
+  "base_commit": "1b5226d4bbdd34653cb5fdd58d436355fb76ac8c",
+  "evaluator_sha256": "2da5a5a9fa2d367e647396880b4ea2b75fed0168c016e1ad3cff1e36e4a9876d",
+  "change": "Disable TF32 after seeding; add float64 reference scoring support",
+  "training_checkout_unchanged": true,
+  "fp64_rescorer_sha256": "a0a8f7f2e4188c384ed8cbf0d8769f0c018b913027b8dc21da719556bb7ea624",
+  "fp64_mask_fix": "Use torch.float64 default tensor dtype in isolated FP64 scorer to avoid missing-dtype torch.full overflow in Transformers causal mask"
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-precision-audit.json

@@ -0,0 +1,24 @@
+{
+  "parameter_dtypes": [
+    "torch.float64"
+  ],
+  "floating_buffer_dtypes": [
+    "torch.float64"
+  ],
+  "observed_leaf_module_output_dtypes": [
+    "torch.float64"
+  ],
+  "audited_module_count": 329,
+  "logits_dtypes": [
+    "torch.float64"
+  ],
+  "nll_dtype": "torch.float64",
+  "autocast_enabled": false,
+  "tf32_allowed": false,
+  "attention_implementation": "sdpa",
+  "reorder_and_upcast_attn": false,
+  "gpu": 7,
+  "peak_allocated_gib": 8.05501937866211,
+  "peak_reserved_gib": 8.537109375,
+  "allocator_fraction_limit": 0.4
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-0.per-example.json

@@ -0,0 +1,392 @@
+{
+  "nll_sum": [
+    1464.3900905656703,
+    1268.4523252489776,
+    1315.5621233205584,
+    1332.39794221183,
+    1206.9263778548375,
+    1384.681874260977,
+    1380.8698006710854,
+    1316.956610847524,
+    1100.288075412007,
+    1110.8369062624454,
+    1237.5394152770118,
+    1005.9572380639675,
+    1380.7136955248511,
+    1194.0719528322488,
+    1001.943173075374,
+    1189.4977288016153,
+    1240.9615620800232,
+    1277.6675585269684,
+    1425.2421466131027,
+    2150.0273629647486,
+    1452.8711834236906,
+    1336.928605040855,
+    1277.7601285090877,
+    1502.3614338534449,
+    1484.847570118482,
+    1042.8909446528012,
+    1257.2196676908,
+    1252.2573844665371,
+    1484.2941220210346,
+    1424.503485860002,
+    1361.4034365753018,
+    1439.155116850991,
+    1390.1145243572937,
+    1849.133897385666,
+    1072.8175256658005,
+    1361.0964217124779,
+    1074.8992209383619,
+    1504.8571831712452,
+    1362.7802909888017,
+    1405.9302854708947,
+    1283.4135303366795,
+    1355.4178301314632,
+    1385.4612109983932,
+    1425.1160141816763,
+    1032.1694589824037,
+    1338.9217874249293,
+    1273.5968996470124,
+    1251.2367450526808,
+    1464.9315578723456,
+    1079.5256854544123,
+    1318.9993888238805,
+    1328.6820245042304,
+    1244.4014377769954,
+    1220.4893029090845,
+    1412.1571589824746,
+    1203.6207079256974,
+    1270.790285857638,
+    1227.7617541410127,
+    1494.1468788718887,
+    1104.6240509436839,
+    1494.4536166328376,
+    1063.7994548912604,
+    1593.7734160104192,
+    1215.850573161972,
+    1409.807325892265,
+    1345.951696347589,
+    1499.172271361063,
+    1722.5908810125577,
+    1299.6269975938797,
+    1243.37216005177,
+    1076.9682614749977,
+    1289.0213954349047,
+    1260.3022628998206,
+    1341.7135801512204,
+    1290.585757575325,
+    1396.8936411431364,
+    1397.904671665892,
+    1450.5104614625247,
+    1250.5964482392799,
+    1551.4143616793608,
+    1175.8251833799284,
+    1221.4895963107942,
+    1222.061122857294,
+    1369.791503134197,
+    1782.7862749869914,
+    1373.732283821038,
+    1723.5840535244934,
+    1560.373838963179,
+    1321.6395625565372,
+    1413.4883759896152,
+    1357.7725335387183,
+    1109.0136312236466,
+    1354.2432716745602,
+    1489.755485594454,
+    1078.9513183770716,
+    1171.8210260241663,
+    1231.9002803293176,
+    958.1657422090016,
+    1338.7019924046162,
+    1379.4428737786689,
+    1167.042298508388,
+    1394.1194020680382,
+    1158.9880694626477,
+    1274.1145147862842,
+    1019.837762462782,
+    1093.2727426614138,
+    1385.2493873878122,
+    1514.8663593017068,
+    1454.2054412411385,
+    1434.9305133602645,
+    948.0564432971278,
+    1360.7787537806726,
+    1337.0567241206968,
+    1257.2942510908415,
+    1180.0792636584677,
+    1051.7770401355947,
+    1107.6169464529864,
+    1284.782427208798,
+    1373.2282237041188,
+    1459.6261996954645,
+    1253.9382650952264,
+    1413.7236527636396,
+    1570.7224832148781,
+    1088.385649765306,
+    1322.9436179734544,
+    1248.7657739781425,
+    1312.9369946427196,
+    1267.6364970186994
+  ],
+  "token_count": [
+    244,
+    251,
+    252,
+    254,
+    253,
+    253,
+    248,
+    245,
+    254,
+    253,
+    253,
+    251,
+    242,
+    251,
+    253,
+    252,
+    252,
+    253,
+    253,
+    245,
+    253,
+    250,
+    249,
+    252,
+    252,
+    250,
+    249,
+    252,
+    254,
+    249,
+    253,
+    253,
+    248,
+    241,
+    249,
+    254,
+    254,
+    249,
+    250,
+    243,
+    251,
+    255,
+    239,
+    254,
+    242,
+    250,
+    255,
+    250,
+    251,
+    251,
+    253,
+    248,
+    254,
+    254,
+    248,
+    254,
+    252,
+    252,
+    245,
+    255,
+    253,
+    253,
+    241,
+    253,
+    251,
+    257,
+    250,
+    251,
+    253,
+    250,
+    249,
+    252,
+    252,
+    250,
+    253,
+    246,
+    254,
+    250,
+    255,
+    248,
+    254,
+    249,
+    255,
+    256,
+    238,
+    251,
+    239,
+    247,
+    255,
+    253,
+    254,
+    246,
+    248,
+    245,
+    254,
+    252,
+    252,
+    254,
+    252,
+    254,
+    254,
+    252,
+    249,
+    255,
+    254,
+    253,
+    252,
+    254,
+    250,
+    257,
+    247,
+    250,
+    254,
+    255,
+    250,
+    254,
+    249,
+    253,
+    254,
+    254,
+    253,
+    251,
+    238,
+    253,
+    251,
+    251,
+    251,
+    250
+  ],
+  "ppl": [
+    404.07428379008036,
+    156.58436235308912,
+    185.02382791990416,
+    189.74122873092713,
+    117.97349628945572,
+    238.1857630890013,
+    261.9158816494278,
+    216.0118150513673,
+    76.08436699145878,
+    80.69363570596369,
+    133.14784660574963,
+    55.02555761508786,
+    300.4942124044239,
+    116.42633742306415,
+    52.47042596638392,
+    112.19395140017421,
+    137.6137218362856,
+    156.03329302282123,
+    279.6022810061183,
+    6474.4690917375365,
+    311.86607458292127,
+    210.1274854645284,
+    169.2821384415936,
+    388.28970432434363,
+    362.220181232671,
+    64.81673191618323,
+    155.878205858847,
+    143.92255396112188,
+    345.0459669706494,
+    305.1787082095823,
+    217.24836783735614,
+    295.4087958310224,
+    271.86360856679323,
+    2148.9932613780757,
+    74.32921808243496,
+    212.43739423171814,
+    68.84700312339152,
+    421.4086994275506,
+    233.01927231775431,
+    325.6168356672514,
+    166.20156788186043,
+    203.4385593686175,
+    329.2801251210154,
+    273.33358415204407,
+    71.17652461914923,
+    211.8094711733221,
+    147.59877987085136,
+    149.14917509091362,
+    342.5373519878235,
+    73.76609031872484,
+    183.72431012256328,
+    212.2126437417713,
+    134.18484087634468,
+    122.12877193065354,
+    297.1336663685596,
+    114.28144303059102,
+    154.90601755596032,
+    130.5910197431867,
+    445.2156133806257,
+    76.08560007992898,
+    367.57643780005253,
+    67.0032367833925,
+    744.8387122282322,
+    122.20909739802065,
+    274.99756815393096,
+    188.13619342153103,
+    402.0952840143011,
+    956.147200430432,
+    170.1815150710301,
+    144.53022282876876,
+    75.5786428177496,
+    166.52813391335224,
+    148.59128095286968,
+    214.188039960954,
+    164.20727086538048,
+    292.48969507579665,
+    245.56500527905848,
+    330.97466969584553,
+    134.86844165197195,
+    520.9755289621236,
+    102.43546652607121,
+    135.0413068191696,
+    120.59002388998857,
+    210.76590426631356,
+    1791.30312076629,
+    238.18245255088718,
+    1355.1247012427423,
+    554.0766262723646,
+    178.19888273994107,
+    266.9097475324237,
+    209.67551219340928,
+    90.75698980096628,
+    235.25225254387567,
+    437.30663680401153,
+    69.95413661686723,
+    104.59371215717586,
+    132.75338726406358,
+    43.480217810821706,
+    202.81809501910703,
+    228.3495241626018,
+    98.95396267937097,
+    252.70425426108065,
+    105.06409159568571,
+    147.89868950074288,
+    55.42935355709616,
+    75.28162920665058,
+    243.96417201876255,
+    389.1795419830371,
+    335.9027773490065,
+    265.9709817109361,
+    46.44576071473974,
+    231.16113353585607,
+    193.2535101091893,
+    138.45781446802494,
+    112.20382177046713,
+    62.856508686193514,
+    85.47815353889686,
+    160.48354466984958,
+    222.83026308325788,
+    313.11165160345166,
+    142.06400956930992,
+    279.3219754906784,
+    734.8558071526454,
+    73.84140106047779,
+    194.55048747344873,
+    144.77234094434053,
+    186.94691131769494,
+    159.26125835732935
+  ]
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-1.per-example.json

@@ -0,0 +1,392 @@
+{
+  "nll_sum": [
+    1282.8641017855098,
+    1400.092761297254,
+    1169.3869786349032,
+    1109.8943496548616,
+    1197.6842517670293,
+    1516.563235951844,
+    1198.5442587586683,
+    1409.2527222203958,
+    1411.5388562372232,
+    1201.9475175038006,
+    1216.3551303670802,
+    1338.0343587862897,
+    957.3136513358281,
+    1726.106777543671,
+    1553.3948428802523,
+    1446.188028658199,
+    1408.3718497804239,
+    1474.5868460645693,
+    1341.686013673987,
+    1298.6409621644502,
+    1335.1002639022927,
+    1317.6654682506373,
+    1412.1391011604958,
+    1300.6462028717929,
+    1292.9480827405564,
+    1237.818258668689,
+    1081.281881859005,
+    1487.6426787908165,
+    1208.2342581988216,
+    1427.0822996434738,
+    1421.4609196476656,
+    1673.055810964495,
+    1385.693102921472,
+    1292.498652951742,
+    1362.2253628497378,
+    1163.3561390044583,
+    1366.2244890266861,
+    1164.5182913364174,
+    1143.624295546129,
+    1596.3803362593394,
+    1377.5231704725484,
+    1371.3988149081524,
+    1384.2233671599172,
+    1222.9846999058882,
+    1213.360848501624,
+    1204.339620936459,
+    1754.9295145729016,
+    1150.396079238962,
+    1184.1424545701188,
+    1477.4780109360058,
+    1421.5290063292227,
+    1428.4511977294937,
+    1280.3755352487854,
+    1264.008728358122,
+    1505.4416344279944,
+    1207.8243758061203,
+    1463.670713158287,
+    1108.1380084377856,
+    1313.6070812636476,
+    1325.9316990450686,
+    1412.919577752941,
+    1325.6250841042142,
+    1229.1698653844826,
+    1390.6812333132148,
+    1363.6863312626126,
+    1167.144101054263,
+    996.0012674488412,
+    1359.2510825676168,
+    1862.0211731111992,
+    1110.8300790039275,
+    1378.7680766938565,
+    1160.3917788953204,
+    1046.0304335844976,
+    1515.5216431825525,
+    1423.572816933578,
+    1231.107825416153,
+    1387.0567613101903,
+    1310.2756298474367,
+    1341.7049707018614,
+    1575.686255472655,
+    1387.115820911841,
+    1012.1517976753775,
+    1415.9072137493326,
+    1254.1160104675268,
+    1215.248257530749,
+    1004.0954044370416,
+    1347.7104063001832,
+    1203.1574587194918,
+    1075.4555928517061,
+    1322.704102969308,
+    1473.4809599560422,
+    1318.4298939065689,
+    1181.7775691075346,
+    1345.2894331624505,
+    1404.9359416008847,
+    1118.3215446615552,
+    1375.6476942117883,
+    1330.4403793519605,
+    1327.3101409890337,
+    1514.4532846640636,
+    1324.965288004103,
+    1503.098345157581,
+    1282.095140170212,
+    1252.27413755045,
+    1731.7539466058224,
+    1055.8172843954067,
+    1477.627047080227,
+    1118.3072248699593,
+    1534.2682643723435,
+    1212.6603101875326,
+    1372.1388818309006,
+    1361.7671469434376,
+    1238.381620709249,
+    1142.8747097820249,
+    1129.641135996422,
+    1357.9153028327826,
+    1353.6968180031392,
+    1332.759116404146,
+    1056.5815495524673,
+    1295.4374309064533,
+    1145.974633483295,
+    1200.6823197208182,
+    1497.7869513158494,
+    1269.982090348119,
+    1473.811351460985,
+    1249.1574710838663,
+    1069.0261099595955,
+    1311.8704411877284
+  ],
+  "token_count": [
+    253,
+    252,
+    253,
+    251,
+    254,
+    253,
+    254,
+    254,
+    256,
+    251,
+    250,
+    251,
+    186,
+    247,
+    248,
+    253,
+    253,
+    251,
+    254,
+    255,
+    252,
+    255,
+    252,
+    254,
+    252,
+    247,
+    252,
+    241,
+    240,
+    252,
+    252,
+    249,
+    250,
+    254,
+    246,
+    252,
+    255,
+    255,
+    252,
+    248,
+    248,
+    255,
+    248,
+    252,
+    252,
+    254,
+    237,
+    253,
+    248,
+    244,
+    255,
+    237,
+    249,
+    255,
+    237,
+    251,
+    255,
+    251,
+    254,
+    253,
+    254,
+    254,
+    254,
+    250,
+    251,
+    248,
+    253,
+    252,
+    245,
+    254,
+    251,
+    250,
+    252,
+    254,
+    252,
+    252,
+    250,
+    251,
+    256,
+    252,
+    253,
+    254,
+    252,
+    254,
+    253,
+    251,
+    248,
+    253,
+    254,
+    250,
+    249,
+    254,
+    251,
+    255,
+    254,
+    254,
+    253,
+    243,
+    250,
+    246,
+    252,
+    252,
+    253,
+    247,
+    248,
+    251,
+    253,
+    254,
+    253,
+    252,
+    255,
+    251,
+    254,
+    253,
+    251,
+    250,
+    252,
+    253,
+    254,
+    249,
+    254,
+    249,
+    249,
+    255,
+    254,
+    255,
+    251,
+    253
+  ],
+  "ppl": [
+    159.27130963525173,
+    258.7658648012341,
+    101.70565635036155,
+    83.25347246337188,
+    111.64144252352838,
+    401.1442509676275,
+    112.02008483362494,
+    256.78497559643205,
+    248.09795727182396,
+    120.13733273918483,
+    129.72547949231438,
+    206.606119943343,
+    171.88877107459786,
+    1083.8627382592006,
+    525.1525976676629,
+    303.7357913281958,
+    261.5661338080401,
+    355.97054629009983,
+    196.80795396520415,
+    162.83047933489433,
+    199.9399197657254,
+    175.44323824132965,
+    271.4360563938115,
+    167.44490001529283,
+    169.1433127623901,
+    150.1162440231846,
+    73.02494670741942,
+    479.52268279425164,
+    153.5934857237824,
+    288.01858191282076,
+    281.664862997611,
+    828.0716151960179,
+    255.3850513870243,
+    162.15900623500002,
+    254.0424743779495,
+    101.13867704134208,
+    212.24538765191176,
+    96.22973421593679,
+    93.5215273650981,
+    624.5413153833377,
+    258.4052045610547,
+    216.5961518536849,
+    265.4816797090937,
+    128.13877912692163,
+    123.33743855417153,
+    114.60535954735224,
+    1643.7999516056193,
+    94.35083068264116,
+    118.48281891167733,
+    426.3402556901076,
+    263.6502869187891,
+    414.5610474993274,
+    171.06958827292902,
+    142.15200942619126,
+    573.681524438938,
+    122.98339004885857,
+    311.02868717339993,
+    82.67295010107911,
+    176.21087342530006,
+    188.82803914760106,
+    260.5189238686462,
+    184.74867613848465,
+    126.37471706955924,
+    260.53180171249426,
+    228.8377579653334,
+    110.63386259015675,
+    51.252474662987105,
+    220.04971591460952,
+    1998.3685883779365,
+    79.30861554791598,
+    243.0093468846857,
+    103.7067407549328,
+    63.492032080998634,
+    390.1848660985173,
+    284.0352868444443,
+    132.33657838986053,
+    256.78188941417477,
+    184.97517610492199,
+    188.86548494595388,
+    519.4254183823319,
+    240.48824930189093,
+    53.7772006252197,
+    275.525289594176,
+    139.41632606394685,
+    121.9185009572041,
+    54.61890659026984,
+    229.13610098022906,
+    116.22908077392748,
+    68.99797368013256,
+    198.50542145640685,
+    371.51684884811954,
+    179.5886348728629,
+    110.86100037937241,
+    195.5165119441364,
+    252.45773904039297,
+    81.68257730257248,
+    229.8306342562656,
+    238.66557148075964,
+    202.19661466850604,
+    471.6863091113105,
+    192.05825969676107,
+    389.42682276050164,
+    158.7879594750588,
+    159.16410542562923,
+    1078.0173187373568,
+    67.1174026382264,
+    343.924828258336,
+    81.67797240291469,
+    430.2220368595737,
+    122.99504721853998,
+    217.22567530327794,
+    227.09470186082774,
+    131.0420405267296,
+    91.58718343072306,
+    90.06776169921584,
+    228.52856451304413,
+    215.25271788116078,
+    193.99310059406372,
+    64.0567801656744,
+    181.73688667735237,
+    91.07753306334554,
+    124.2154239617479,
+    409.61129122911797,
+    145.52122073222702,
+    331.0955139440707,
+    134.109515947679,
+    70.7440311200285,
+    178.619645719249
+  ]
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-2.per-example.json

@@ -0,0 +1,392 @@
+{
+  "nll_sum": [
+    1395.8905319607165,
+    1390.8910196293236,
+    1421.69934997863,
+    1519.6282287892745,
+    1422.604939235942,
+    1769.5007008096286,
+    978.845964172355,
+    1514.6548417781357,
+    1231.686619907808,
+    1701.5124392562404,
+    1326.3068973670752,
+    1176.928039389526,
+    1334.3269639344578,
+    1353.376264279329,
+    1284.9083244141307,
+    1816.209980128403,
+    1477.0366185893706,
+    1522.102917585078,
+    1403.3694568405544,
+    1373.824645342557,
+    1486.042091090831,
+    1467.6622403219972,
+    1379.4340369215865,
+    1408.6066541648966,
+    1232.540403188052,
+    1422.345453711965,
+    1140.6633192829,
+    1383.0520756872368,
+    1593.1117113129367,
+    1419.4001349527805,
+    1386.8793082976342,
+    1304.9622871238955,
+    1386.3508638805442,
+    1308.7039428762366,
+    1283.521113251807,
+    1499.0263135678772,
+    1365.4921650225388,
+    1513.53441184859,
+    1137.3145552259862,
+    1386.8782169208948,
+    1356.5408868828063,
+    1325.6095869699286,
+    1453.708374363745,
+    1261.7904891818948,
+    1261.1020185244697,
+    1154.5692532873945,
+    1495.6524928343458,
+    1195.6778564457175,
+    1349.8173044537077,
+    1274.307789242344,
+    1386.7147562560046,
+    1012.3928897557752,
+    1088.9361920284082,
+    1224.5600333119246,
+    1280.9218201490398,
+    1297.3798795159175,
+    1283.525916217523,
+    1661.4581114962373,
+    1173.4248411347808,
+    1454.2368008279268,
+    1185.5495774571357,
+    1357.8840828575783,
+    1149.5061327552703,
+    1316.505347066273,
+    1406.0579031496661,
+    1147.4239174636932,
+    1576.6373586555956,
+    1198.0735480512337,
+    1282.8029272182177,
+    993.1469983772247,
+    1639.8111184994743,
+    1461.1559723356504,
+    1696.2628643049816,
+    1248.971870903651,
+    1557.2016097572914,
+    1123.190341917492,
+    1575.4692219233584,
+    1489.4025396735185,
+    1464.370743455736,
+    1432.2207604035625,
+    1081.1508399143834,
+    1290.7715939535942,
+    1231.5442724216177,
+    1325.7106782946448,
+    1092.5804512518457,
+    1264.5119877524403,
+    1699.3361255588338,
+    1344.4384278357752,
+    1357.8824384664501,
+    1111.0352142074648,
+    1219.2728118047976,
+    1157.1208014780113,
+    1381.1832312127235,
+    1378.6450954196403,
+    1372.801801767232,
+    1153.0490386905005,
+    1163.7628533683146,
+    896.5742556075472,
+    1367.2145654807157,
+    1338.4425525780298,
+    1201.3827036565522,
+    992.882924055608,
+    1194.957612014015,
+    1369.8330911477558,
+    1208.30226396041,
+    1034.0477491715726,
+    1282.6986860072504,
+    1274.8136781537569,
+    1146.1770747006608,
+    1450.4928821644487,
+    1348.7634602022422,
+    1253.4292184250723,
+    1547.3923999016365,
+    1117.1517809485863,
+    1255.1028837604604,
+    1230.8062301558293,
+    1179.379203406135,
+    1252.4949457427924,
+    1319.365045039864,
+    1269.0613136957827,
+    1563.7312736985973,
+    1178.3068694321241,
+    1463.882119666736,
+    1117.0994834016751,
+    1279.9967773403318,
+    1422.5622836975022,
+    1632.6258265206163,
+    1322.536171213321
+  ],
+  "token_count": [
+    254,
+    250,
+    251,
+    253,
+    252,
+    240,
+    188,
+    253,
+    255,
+    244,
+    249,
+    253,
+    253,
+    253,
+    254,
+    241,
+    250,
+    249,
+    253,
+    253,
+    254,
+    249,
+    250,
+    252,
+    255,
+    242,
+    249,
+    255,
+    252,
+    251,
+    251,
+    253,
+    254,
+    254,
+    252,
+    254,
+    242,
+    249,
+    253,
+    253,
+    251,
+    253,
+    251,
+    253,
+    249,
+    255,
+    249,
+    252,
+    250,
+    253,
+    254,
+    249,
+    253,
+    252,
+    253,
+    253,
+    250,
+    253,
+    255,
+    251,
+    255,
+    253,
+    253,
+    255,
+    254,
+    252,
+    251,
+    253,
+    253,
+    250,
+    250,
+    249,
+    246,
+    253,
+    251,
+    253,
+    245,
+    251,
+    255,
+    239,
+    252,
+    254,
+    252,
+    255,
+    251,
+    254,
+    247,
+    249,
+    252,
+    251,
+    255,
+    253,
+    252,
+    251,
+    251,
+    255,
+    250,
+    254,
+    251,
+    255,
+    253,
+    251,
+    253,
+    253,
+    254,
+    251,
+    252,
+    253,
+    255,
+    242,
+    253,
+    252,
+    250,
+    248,
+    255,
+    255,
+    253,
+    254,
+    254,
+    249,
+    249,
+    253,
+    254,
+    251,
+    251,
+    253,
+    241,
+    237
+  ],
+  "ppl": [
+    243.6254525985763,
+    260.7505174944581,
+    288.3401433582111,
+    406.03350747145464,
+    282.9464608247346,
+    1592.275790607089,
+    182.47760789005895,
+    398.1297792872806,
+    125.22894317667507,
+    1067.8596919339548,
+    205.72364198783313,
+    104.7827840177995,
+    195.19901352324067,
+    210.46372227774273,
+    157.3848674452149,
+    1874.582030303681,
+    368.0233820601826,
+    451.6299378074352,
+    256.4451622027734,
+    228.1804898903562,
+    347.4286825059106,
+    362.9357657489316,
+    249.07053954392234,
+    267.6577052885402,
+    125.64893364042278,
+    356.90175667778186,
+    97.60973046872991,
+    226.72404794118918,
+    556.613926972841,
+    285.7109448352429,
+    250.99062075080045,
+    173.80842797763302,
+    234.64512519825067,
+    172.84197329019275,
+    162.93278454382678,
+    365.65066115304825,
+    282.1756099753833,
+    436.3529582875578,
+    89.59633790355804,
+    240.262501695781,
+    222.41508005200959,
+    188.5877818810388,
+    327.55855406571425,
+    146.54230916003408,
+    158.32766916978565,
+    92.54754957608043,
+    406.11505956968705,
+    114.97944791323927,
+    221.24467544537345,
+    153.97491211520176,
+    234.9815297845333,
+    58.31357412672631,
+    74.0022590666308,
+    128.94232502485053,
+    158.05326483924142,
+    168.67667444303913,
+    169.71213273299128,
+    711.2529326606659,
+    99.65019900194656,
+    328.2488844998735,
+    104.50281713212672,
+    214.24725864701554,
+    94.01952758951603,
+    174.6468730374083,
+    253.5753546908781,
+    94.94231662328211,
+    534.5491811791919,
+    113.91682414933534,
+    159.23280301237222,
+    53.121832052251094,
+    705.7382894916361,
+    353.5752289518374,
+    987.6985078438613,
+    139.30248378235714,
+    494.7192733368401,
+    84.73150712710225,
+    620.4758105586563,
+    377.61481194598025,
+    311.88370142992324,
+    400.4366124155894,
+    72.98698304317085,
+    161.06015503381494,
+    132.56597501143185,
+    181.06668794212104,
+    77.70426638919102,
+    145.24088684141947,
+    972.5323525505872,
+    221.262795681917,
+    218.85783742801019,
+    83.63274385925332,
+    119.27860880028436,
+    96.89229549130052,
+    240.05926686952958,
+    242.89030991707762,
+    237.30112102615212,
+    91.99745706242155,
+    105.11462416789364,
+    34.117822576912296,
+    232.07717843068914,
+    190.33664570413177,
+    115.41660514535364,
+    52.232706541557356,
+    112.52243417537444,
+    224.6087599656577,
+    116.4073325120639,
+    61.541524869320845,
+    162.40190384176225,
+    154.28310235119812,
+    89.55133691472878,
+    400.9241432057838,
+    206.66122747691438,
+    144.59337171916852,
+    487.63618223730253,
+    90.43616819342925,
+    137.27306169251568,
+    124.79733463929254,
+    105.80289469072041,
+    138.52938421952018,
+    180.251044526186,
+    163.4703751064829,
+    533.8128321337649,
+    105.35540086613753,
+    318.40222173563353,
+    85.67795443683288,
+    163.95446787106442,
+    276.6562677281303,
+    875.1375052840887,
+    265.15681283609206
+  ]
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-3.per-example.json

@@ -0,0 +1,392 @@
+{
+  "nll_sum": [
+    1326.4778324506874,
+    1292.2152341369588,
+    1331.1084774621327,
+    1206.6701966802302,
+    1508.3609232735862,
+    1202.3034113450049,
+    1180.5849638145103,
+    1363.0924762895165,
+    1104.559157582879,
+    1267.2658315258568,
+    1422.676886860936,
+    1395.2950827287805,
+    1194.9249908233428,
+    1161.3345641506319,
+    1207.0869486663687,
+    948.2832579190958,
+    1128.960106387153,
+    1572.9109783124982,
+    1274.2166431398969,
+    1257.7863388125943,
+    1184.0474165744063,
+    1194.9395427360591,
+    1232.661686470744,
+    1350.671635782147,
+    1404.9924297196264,
+    1396.9833400715927,
+    1483.5882621155174,
+    1125.0239919957262,
+    1592.9944721960528,
+    1231.6463976280356,
+    1213.2011689522046,
+    927.0018631888959,
+    1424.8595096591032,
+    1290.689824335513,
+    1575.0789672206217,
+    1289.0816943745476,
+    1412.5473659713175,
+    1344.9502729540159,
+    1115.6999776799082,
+    1046.3755166626079,
+    1264.8747754677665,
+    1296.453843084956,
+    1319.6259505247856,
+    1149.0643215003379,
+    1644.5469489893885,
+    1272.7669313374759,
+    1210.840842234679,
+    1370.4032163021252,
+    1224.9564199630643,
+    1408.592917088552,
+    1400.4205212991226,
+    1192.4888645816345,
+    1372.432409940108,
+    1734.2153919014108,
+    1530.732595279746,
+    1108.245195054079,
+    1261.980373253532,
+    1371.3415701736108,
+    1219.335653146477,
+    1408.138143509049,
+    1424.124577072114,
+    1052.2904068163612,
+    1350.6917917349926,
+    1428.5242824858947,
+    1260.994515761327,
+    1532.7450738069779,
+    1351.4084544501177,
+    1142.1703323394895,
+    1192.5474982877658,
+    1151.822131003265,
+    1372.8746360236255,
+    1284.4060460310354,
+    1367.337334357088,
+    1397.4591198746352,
+    1378.6587266852248,
+    1492.7206554972531,
+    1042.0217082837898,
+    1697.6603579796774,
+    1186.7610578474303,
+    1410.453403381914,
+    1351.1196837233276,
+    1374.73357635517,
+    1379.5061090695845,
+    1509.9910381519567,
+    1245.2493336365558,
+    1092.443299128602,
+    1345.8596227476896,
+    1601.0922510691798,
+    1292.1516172162178,
+    1557.2160293983939,
+    1251.0804184675562,
+    1532.4715593823087,
+    1272.4765073569474,
+    1273.9894134197368,
+    1405.5052756322202,
+    1227.2042357479588,
+    1427.689033168857,
+    1353.0928868502338,
+    1131.5984208188518,
+    1416.5392269623487,
+    1582.8013906477047,
+    1500.264501802741,
+    1475.3132254701907,
+    1160.2280152503727,
+    1172.3481831470108,
+    1165.7304868971219,
+    1216.6511405364645,
+    1393.875559914543,
+    1126.7280401226792,
+    1395.9589925459545,
+    1327.4653693926364,
+    1676.5023280022033,
+    1902.838621746405,
+    1219.6189252575023,
+    1315.74835692945,
+    1181.2568340659454,
+    1105.9187916341957,
+    1343.3561492250526,
+    1454.978902998197,
+    1113.1948803194996,
+    1118.7105231342407,
+    1389.5181354505996,
+    1330.4809051077493,
+    1284.1488188705398,
+    1064.8913451588658,
+    1158.0133075466647,
+    1319.711258375042,
+    1260.9373093626966
+  ],
+  "token_count": [
+    252,
+    253,
+    236,
+    251,
+    253,
+    249,
+    250,
+    252,
+    254,
+    255,
+    254,
+    240,
+    253,
+    255,
+    252,
+    252,
+    251,
+    248,
+    250,
+    253,
+    254,
+    253,
+    254,
+    254,
+    246,
+    249,
+    251,
+    254,
+    242,
+    255,
+    251,
+    249,
+    253,
+    254,
+    236,
+    253,
+    251,
+    254,
+    254,
+    254,
+    254,
+    252,
+    247,
+    248,
+    246,
+    237,
+    253,
+    254,
+    251,
+    236,
+    251,
+    253,
+    254,
+    249,
+    254,
+    250,
+    253,
+    255,
+    254,
+    250,
+    253,
+    253,
+    252,
+    253,
+    253,
+    254,
+    253,
+    251,
+    254,
+    251,
+    252,
+    253,
+    252,
+    252,
+    253,
+    237,
+    250,
+    246,
+    248,
+    246,
+    255,
+    247,
+    253,
+    230,
+    254,
+    252,
+    253,
+    252,
+    253,
+    238,
+    253,
+    242,
+    250,
+    253,
+    250,
+    251,
+    255,
+    254,
+    251,
+    252,
+    242,
+    253,
+    235,
+    252,
+    250,
+    252,
+    251,
+    253,
+    252,
+    252,
+    253,
+    241,
+    240,
+    254,
+    255,
+    253,
+    253,
+    251,
+    227,
+    253,
+    252,
+    250,
+    253,
+    249,
+    253,
+    254,
+    252,
+    253
+  ],
+  "ppl": [
+    193.21449066800125,
+    165.26828041595712,
+    281.54439919203867,
+    122.41917089705592,
+    388.3476262717193,
+    125.02675525088293,
+    112.43101743948085,
+    223.42976813376256,
+    77.37456435794272,
+    143.9793559085265,
+    270.721361544686,
+    334.8656850808404,
+    112.50792674790307,
+    95.0357553359334,
+    120.30468590765872,
+    43.078705421426854,
+    89.82371519113076,
+    568.1485845710926,
+    163.5087609280556,
+    144.24128753425708,
+    105.80565863549893,
+    112.51439810211168,
+    128.12402827039514,
+    203.89494083553583,
+    302.27927458061026,
+    273.24664797226166,
+    368.96810634349157,
+    83.86667327415492,
+    722.4309077009206,
+    125.20919181841899,
+    125.64629697927144,
+    41.384195177757505,
+    279.1797304012968,
+    161.00831366481097,
+    791.6057048426559,
+    163.23396875379686,
+    278.01602304708405,
+    199.35353669019827,
+    80.8438565342412,
+    61.53392497163,
+    145.4484823417882,
+    171.5128356902003,
+    209.058723108551,
+    102.8553751759249,
+    800.4308931821748,
+    214.93262454351137,
+    119.81299865346237,
+    220.36565903942326,
+    131.67074652308673,
+    390.96344143979076,
+    264.90323985587355,
+    111.42979169379704,
+    222.1332005651885,
+    1058.6189385015648,
+    414.2651719033203,
+    84.18197125552442,
+    146.652334832307,
+    216.5475338238791,
+    121.57532961652784,
+    279.37433192252263,
+    278.3699258102404,
+    64.02352665712934,
+    212.7011358708591,
+    283.25315473867204,
+    146.0819910314966,
+    417.560491842887,
+    208.8331051653857,
+    94.67778707640682,
+    109.40634548344664,
+    98.38937227026233,
+    232.2734223444546,
+    160.24497519201577,
+    227.22524699359178,
+    256.075596038974,
+    232.5822623464534,
+    543.7009602467585,
+    64.5917590092849,
+    993.3254611379683,
+    119.74049340905897,
+    309.06463205481356,
+    200.03824307043632,
+    261.31406058335415,
+    233.36256485544408,
+    709.9385569244976,
+    134.63352151568859,
+    76.33201598200375,
+    204.3028091736459,
+    574.5233145776454,
+    165.22672888492707,
+    694.3141288094763,
+    140.4683071125788,
+    562.5764887789388,
+    162.37460290219772,
+    153.78127157130174,
+    276.447547828864,
+    132.85521196350166,
+    270.09682943145737,
+    205.84785949805084,
+    90.77285092783964,
+    276.2171712263963,
+    692.633925956893,
+    376.11660560857456,
+    532.6841284927551,
+    99.89098342715097,
+    108.78709688566431,
+    102.09611061271865,
+    127.38521687896636,
+    247.00030337834121,
+    87.45663026392084,
+    254.55573541770664,
+    189.9761795006231,
+    1049.890545177565,
+    2775.2448520934468,
+    121.71099147206984,
+    174.12918706962205,
+    106.59102811167517,
+    79.1401649545218,
+    211.03342447762216,
+    607.6497771779484,
+    81.44922045071961,
+    84.71794738310743,
+    259.32252099049384,
+    192.2540754287219,
+    173.68167542970906,
+    67.29303235977635,
+    95.49821058929201,
+    188.09542954326852,
+    146.04896383824487
+  ]
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/fp64-seed-4.per-example.json

@@ -0,0 +1,392 @@
+{
+  "nll_sum": [
+    1555.6135330991824,
+    1323.8989381595102,
+    1357.4001829763474,
+    977.6157690283565,
+    1290.2417714124654,
+    1560.2890129227703,
+    1180.1057249302798,
+    1309.3710573889732,
+    1197.2059510581225,
+    1490.3452447047907,
+    1264.1286599788327,
+    1095.0863969061782,
+    1218.1928258551363,
+    1131.1936169742494,
+    1436.851058212294,
+    1406.5997162702147,
+    1095.7943618962122,
+    1308.4009260610935,
+    1280.404771980779,
+    1370.465177666214,
+    1341.408047455575,
+    1045.044439561996,
+    1157.259022731473,
+    1508.3450312691796,
+    1563.130751371985,
+    1119.2500353289781,
+    1103.9126556135197,
+    1530.0974210990883,
+    1155.7094293730897,
+    1336.9305268336238,
+    1646.7251024820935,
+    1069.509197484343,
+    1150.2452154306413,
+    1194.4903578182257,
+    1070.6813548193827,
+    1627.825549760531,
+    1423.7339303533913,
+    1174.0704192458402,
+    1283.481029323761,
+    1371.72956944818,
+    1360.4818047993099,
+    1301.3017308580193,
+    1051.1651336932182,
+    1313.915661447692,
+    1230.2640794890906,
+    1064.2042303670517,
+    1546.0631703737145,
+    1443.721202381376,
+    1399.3736315827502,
+    1243.3180808250395,
+    1004.0256498570718,
+    1099.8969590249646,
+    1419.1473556354708,
+    1103.0320449774922,
+    1207.845469845116,
+    1383.7337347090133,
+    1080.3834220373478,
+    1476.8324953818133,
+    1407.5764629147682,
+    1248.9482693109537,
+    1291.707392207926,
+    1203.5523268587795,
+    1580.4322292303223,
+    1106.479207243926,
+    1393.6497233631883,
+    1573.4974166500356,
+    1331.3402921873464,
+    1230.2434623689805,
+    1592.8741192708542,
+    1213.0297246241057,
+    1402.074099526489,
+    1095.7897532387067,
+    1366.1124440618719,
+    1063.9099401611022,
+    1523.4546868117418,
+    1148.197867137052,
+    1346.2473672984252,
+    1247.1411022208283,
+    1257.648609130259,
+    1274.3826106121141,
+    1059.3708549796686,
+    1364.071190802188,
+    1271.7578589406683,
+    1389.6789187746526,
+    1107.0246550042912,
+    1236.501647628898,
+    1002.2127485013282,
+    1277.8308822573151,
+    1383.0273800360035,
+    1195.3799618296719,
+    1442.8489735949888,
+    1183.6185576470284,
+    1276.8594138829578,
+    1167.7708419603287,
+    1386.5929836460086,
+    1284.891827137601,
+    1108.4830721425317,
+    1453.5743036461613,
+    1334.7508624000286,
+    1242.9676100125357,
+    1594.7571453813266,
+    1520.3701440596192,
+    1552.238644530701,
+    1301.2957664617438,
+    1204.5786466706331,
+    1077.661776967967,
+    1283.0547049530956,
+    1410.95942493366,
+    1236.1816774382764,
+    1056.3858198288453,
+    1366.097091479799,
+    1331.7089817814503,
+    1199.1317819748774,
+    1265.2959061301128,
+    1457.3195047057748,
+    1523.0119778750702,
+    1468.0463053522053,
+    1377.6446016689913,
+    1220.368870442225,
+    1407.602292429834,
+    1222.7080367121034,
+    1136.2526046519997,
+    1179.850598863316,
+    1605.9849086515592,
+    1446.1627682713229,
+    1334.7144257270136,
+    1081.001918537354,
+    1330.978536255479
+  ],
+  "token_count": [
+    250,
+    251,
+    255,
+    254,
+    252,
+    253,
+    251,
+    251,
+    255,
+    252,
+    252,
+    255,
+    251,
+    250,
+    254,
+    250,
+    246,
+    253,
+    254,
+    250,
+    251,
+    250,
+    250,
+    250,
+    247,
+    255,
+    253,
+    252,
+    253,
+    248,
+    238,
+    254,
+    252,
+    248,
+    254,
+    251,
+    254,
+    249,
+    255,
+    253,
+    252,
+    248,
+    249,
+    250,
+    251,
+    249,
+    256,
+    250,
+    251,
+    250,
+    247,
+    253,
+    252,
+    251,
+    253,
+    254,
+    250,
+    254,
+    253,
+    252,
+    253,
+    250,
+    237,
+    253,
+    253,
+    249,
+    253,
+    254,
+    245,
+    252,
+    253,
+    238,
+    252,
+    254,
+    252,
+    247,
+    254,
+    254,
+    250,
+    253,
+    251,
+    253,
+    254,
+    252,
+    252,
+    255,
+    253,
+    252,
+    236,
+    253,
+    252,
+    252,
+    254,
+    251,
+    254,
+    251,
+    254,
+    253,
+    246,
+    255,
+    250,
+    252,
+    250,
+    251,
+    255,
+    251,
+    254,
+    254,
+    255,
+    252,
+    254,
+    250,
+    253,
+    252,
+    250,
+    248,
+    251,
+    255,
+    254,
+    251,
+    254,
+    254,
+    253,
+    249,
+    254,
+    254,
+    252,
+    246
+  ],
+  "ppl": [
+    503.9384473789441,
+    195.2923685494811,
+    205.02624030941675,
+    46.94050648550723,
+    167.3365458949194,
+    476.825334466206,
+    110.12503851389472,
+    184.309748593641,
+    109.39063666646939,
+    370.2092666768868,
+    150.86472454439257,
+    73.29236617050778,
+    128.17004667837938,
+    92.27511248153435,
+    286.2581224879596,
+    277.66042245284837,
+    86.00871476483641,
+    176.18686456726928,
+    154.61894477209515,
+    240.29340816217493,
+    209.40185753476206,
+    65.37747356633136,
+    102.41529748254528,
+    417.1225756787058,
+    560.2956441786955,
+    80.57720755302466,
+    78.51511302794154,
+    433.4667808114502,
+    96.35328222947561,
+    219.38954598859664,
+    1011.3213636660325,
+    67.40142309254811,
+    96.01122777300127,
+    123.53115337418288,
+    67.71318549108146,
+    655.4753882818251,
+    271.8503400064766,
+    111.62468820619996,
+    153.43222499162505,
+    226.29874473651287,
+    221.12702731695578,
+    190.03046446957129,
+    68.13879457499704,
+    191.64843881859832,
+    134.48471048554225,
+    71.80201729967389,
+    419.60309781341067,
+    322.1073288853454,
+    263.80066241383895,
+    144.49896187926754,
+    58.25798551933231,
+    77.27873420037263,
+    279.0907856642965,
+    81.00816726244966,
+    118.40284681260466,
+    232.2398349643297,
+    75.30403267474946,
+    335.0571662940135,
+    260.7451078017342,
+    142.04500241903375,
+    164.9368733685865,
+    123.24932145051116,
+    787.2063550560632,
+    79.31566116420655,
+    246.7799207389395,
+    555.1657587304405,
+    192.90823193511008,
+    126.91000310688372,
+    666.1580917904067,
+    123.1754816052881,
+    255.13552122968525,
+    99.89889525185679,
+    226.12345862739977,
+    65.93186179274625,
+    422.1898550703103,
+    104.43599136140708,
+    200.37417335407227,
+    135.64000301048122,
+    153.02394868879304,
+    154.02045487226297,
+    68.07438578215736,
+    219.551259647209,
+    149.44384514220485,
+    248.2903640394784,
+    80.87906223932781,
+    127.6160401012432,
+    52.526363815497284,
+    159.29494268107723,
+    350.82429308480584,
+    112.71043227326562,
+    306.6144714389467,
+    109.60675855734353,
+    152.47574991747425,
+    104.84399400728219,
+    234.86890200443315,
+    167.1833216488796,
+    78.57916462894705,
+    312.7339968329345,
+    227.19677991137064,
+    130.89333885790577,
+    589.3549192381851,
+    417.05364448734775,
+    497.18121538622034,
+    178.47441983576184,
+    112.59956247951148,
+    73.22033734872778,
+    156.24049834778387,
+    258.5162012021146,
+    127.45600982334516,
+    66.15544641866553,
+    216.66123524903367,
+    205.78582163520414,
+    114.39430682065839,
+    151.56514014846093,
+    340.11303519786094,
+    464.6002368555113,
+    346.81450218140867,
+    221.96680633565134,
+    122.07087908453741,
+    272.5922968183512,
+    123.20026082187941,
+    87.6573488856898,
+    106.0002128603484,
+    632.5369190654015,
+    296.94715859725756,
+    191.47958623616944,
+    72.94386355508747,
+    223.73937157748583
+  ]
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/limited_eval.py

@@ -0,0 +1,14 @@
+import json
+from pathlib import Path
+import runpy
+import sys
+import torch
+
+torch.cuda.set_per_process_memory_fraction(0.4)
+if torch.cuda.mem_get_info()[0] < 40 * 2**30:
+    raise RuntimeError('Insufficient spare GPU memory for concurrent evaluation')
+output=Path(sys.argv[sys.argv.index('--out')+1])
+runpy.run_path('/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/eval_random_gen_ppl_checkpoint_fixed.py',run_name='__main__')
+output.with_suffix('.memory.json').write_text(json.dumps(dict(
+    gpu=7,peak_allocated_gib=torch.cuda.max_memory_allocated()/2**30,
+    peak_reserved_gib=torch.cuda.max_memory_reserved()/2**30,allocator_fraction_limit=0.4),indent=2)+'\n')

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-0.memory.json

@@ -0,0 +1,6 @@
+{
+  "gpu": 7,
+  "peak_allocated_gib": 4.060837268829346,
+  "peak_reserved_gib": 4.298828125,
+  "allocator_fraction_limit": 0.4
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-1.memory.json

@@ -0,0 +1,6 @@
+{
+  "gpu": 7,
+  "peak_allocated_gib": 4.058009147644043,
+  "peak_reserved_gib": 4.298828125,
+  "allocator_fraction_limit": 0.4
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-2.memory.json

@@ -0,0 +1,6 @@
+{
+  "gpu": 7,
+  "peak_allocated_gib": 4.056677341461182,
+  "peak_reserved_gib": 4.298828125,
+  "allocator_fraction_limit": 0.4
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-3.memory.json

@@ -0,0 +1,6 @@
+{
+  "gpu": 7,
+  "peak_allocated_gib": 4.056677341461182,
+  "peak_reserved_gib": 4.298828125,
+  "allocator_fraction_limit": 0.4
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/random-seed-4.memory.json

@@ -0,0 +1,6 @@
+{
+  "gpu": 7,
+  "peak_allocated_gib": 4.056677341461182,
+  "peak_reserved_gib": 4.298828125,
+  "allocator_fraction_limit": 0.4
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/rescore_fp64.py

@@ -0,0 +1,128 @@
+"""Score the exact saved FP32 generations with a double-precision reference."""
+import datetime
+import hashlib
+import json
+import math
+import os
+from pathlib import Path
+import statistics
+import sys
+import time
+
+os.environ.update(CUDA_VISIBLE_DEVICES='7', OMP_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false', PYTORCH_ALLOC_CONF='expandable_segments:True')
+root = Path(__file__).resolve().parent
+checkout = Path('/home/ubuntu/mstok-runs/iclr-debug-two-stage-20260916')
+sys.path.insert(0, str(checkout))
+
+def write(name, data):
+    (root/name).write_text(json.dumps(data, indent=2)+'\n')
+
+def status(stage, **kw):
+    write('fp64-status.json', dict(stage=stage, updated_utc=datetime.datetime.now(datetime.timezone.utc).isoformat(), **kw))
+
+try:
+    status('waiting_for_fp32_samples')
+    deadline=time.monotonic()+3600
+    while not (root/'FP32_DONE.json').exists():
+        current=json.loads((root/'status.json').read_text())
+        if current['stage']=='failed':
+            raise RuntimeError(f'FP32 evaluation failed: {current}')
+        if time.monotonic()>deadline:
+            raise TimeoutError('FP32 samples were not completed within one hour')
+        time.sleep(5)
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from evaluation.gen_ppl import generation_perplexity
+    torch.cuda.set_per_process_memory_fraction(0.4)
+    if torch.cuda.mem_get_info()[0] < 40*2**30:
+        raise RuntimeError('Insufficient spare GPU memory for concurrent evaluation')
+    torch.backends.cuda.matmul.allow_tf32=False
+    torch.backends.cudnn.allow_tf32=False
+    # The installed Transformers mask helper omits dtype in torch.full.
+    # Keep newly constructed reference tensors in FP64 to avoid FP32 overflow.
+    torch.set_default_dtype(torch.float64)
+    status('loading_reference')
+    model=AutoModelForCausalLM.from_pretrained('gpt2-large',torch_dtype=torch.float64).cuda().eval()
+    tokenizer=AutoTokenizer.from_pretrained('gpt2-large')
+    assert not model.config.reorder_and_upcast_attn
+    parameter_dtypes={str(p.dtype) for p in model.parameters()}
+    floating_buffer_dtypes={str(b.dtype) for b in model.buffers() if b.is_floating_point()}
+    assert parameter_dtypes=={'torch.float64'}, parameter_dtypes
+    assert floating_buffer_dtypes <= {'torch.float64'}, floating_buffer_dtypes
+    observed_dtypes=set()
+    observed_modules=set()
+    logits_dtypes=set()
+    def audit_tensor(obj):
+        if isinstance(obj,torch.Tensor) and obj.is_floating_point():
+            observed_dtypes.add(str(obj.dtype))
+            assert obj.dtype==torch.float64, obj.dtype
+        elif isinstance(obj,(tuple,list)):
+            for item in obj: audit_tensor(item)
+    handles=[]
+    for name,module in model.named_modules():
+        if not list(module.children()):
+            def hook(mod, inputs, output, name=name):
+                observed_modules.add(name)
+                audit_tensor(output)
+            handles.append(module.register_forward_hook(hook))
+    def logit_hook(mod,inputs,output):
+        logits_dtypes.add(str(output.logits.dtype))
+        assert output.logits.dtype==torch.float64
+        assert not torch.is_autocast_enabled()
+    handles.append(model.register_forward_hook(logit_hook))
+    rows=[]
+    for seed in range(5):
+        status('scoring',seed=seed)
+        source=root/f'random-seed-{seed}.samples.json'
+        sample_bytes=source.read_bytes()
+        samples=json.loads(sample_bytes)
+        previous=json.loads((root/f'random-seed-{seed}.json').read_text())
+        kept=[(p,g) for p,g in zip(samples['prompts'],samples['generations']) if g.strip()]
+        assert len(kept)==128
+        prompts,generations=map(list,zip(*kept))
+        torch.cuda.synchronize()
+        started=time.monotonic()
+        result=generation_perplexity(model,tokenizer,prompts,generations,score_only_generated=True,batch_size=4,device='cuda')
+        torch.cuda.synchronize()
+        assert result.nll_sum.dtype==torch.float64 and result.ppl.dtype==torch.float64
+        assert int(result.token_count.sum())==previous['scored_tokens']
+        assert torch.isfinite(result.nll_sum).all() and (result.token_count>0).all()
+        row=dict(previous,reference_dtype='float64',tf32_allowed=False,
+                 mean_ppl=float(result.mean_ppl),median_ppl=float(result.ppl.median()),
+                 mean_nll=float(result.mean_nll),samples_sha256=hashlib.sha256(sample_bytes).hexdigest(),
+                 samples_source=str(source),scoring_seconds=time.monotonic()-started,
+                 nll_dtype=str(result.nll_sum.dtype),logits_dtypes=sorted(logits_dtypes),
+                 fp32_mean_ppl=previous['mean_ppl'],ppl_delta_fp64_minus_fp32=result.mean_ppl-previous['mean_ppl'])
+        write(f'fp64-seed-{seed}.json',row)
+        write(f'fp64-seed-{seed}.per-example.json',dict(nll_sum=result.nll_sum.tolist(),token_count=result.token_count.tolist(),ppl=result.ppl.tolist()))
+        rows.append(row)
+        print(json.dumps({k:row[k] for k in ('seed','mean_ppl','ppl_delta_fp64_minus_fp32','scoring_seconds')}),flush=True)
+    fp32=json.loads((root/'summary.json').read_text())
+    values=[r['mean_ppl'] for r in rows]
+    fp64=dict(fp32, reference_dtype='float64',mean_gen_ppl=statistics.mean(values),
+              gen_ppl_se=statistics.stdev(values)/math.sqrt(len(values)),
+              mean_of_seed_median_ppl=statistics.mean(r['median_ppl'] for r in rows),
+              per_seed_mean_ppl={str(r['seed']):r['mean_ppl'] for r in rows})
+    write('summary-fp64.json',fp64)
+    audit=dict(parameter_dtypes=sorted(parameter_dtypes),floating_buffer_dtypes=sorted(floating_buffer_dtypes),
+               observed_leaf_module_output_dtypes=sorted(observed_dtypes),audited_module_count=len(observed_modules),
+               logits_dtypes=sorted(logits_dtypes),nll_dtype='torch.float64',autocast_enabled=False,
+               tf32_allowed=False,attention_implementation=model.config._attn_implementation,
+               reorder_and_upcast_attn=model.config.reorder_and_upcast_attn,
+               gpu=7,peak_allocated_gib=torch.cuda.max_memory_allocated()/2**30,
+               peak_reserved_gib=torch.cuda.max_memory_reserved()/2**30,allocator_fraction_limit=0.4)
+    write('fp64-precision-audit.json',audit)
+    previous=json.loads((root.parent/'step-12875-fp32/summary.json').read_text())
+    comparison=dict(checkpoint_step=22000,training_input_positions=22000*4096*256,
+                    fp32=fp32,fp64=fp64,matched_samples=True,total_samples=640,
+                    absolute_ppl_difference=fp64['mean_gen_ppl']-fp32['mean_gen_ppl'],
+                    relative_ppl_difference_percent=100*(fp64['mean_gen_ppl']/fp32['mean_gen_ppl']-1),
+                    previous_step=12875,previous_fp32_gen_ppl=previous['mean_gen_ppl'],
+                    improvement_vs_previous_fp32_percent=100*(1-fp32['mean_gen_ppl']/previous['mean_gen_ppl']))
+    write('comparison.json',comparison)
+    status('complete',seeds=list(range(5)),total_samples=640)
+    write('EVAL_DONE.json',dict(step=22000,reference_dtypes=['float32','float64'],matched_samples=True,total_samples=640))
+    print(json.dumps(comparison,indent=2),flush=True)
+except BaseException as exc:
+    status('failed',error=str(exc))
+    raise

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/run_evaluation.py

@@ -0,0 +1,44 @@
+import datetime
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+root=Path(__file__).resolve().parent
+checkout=Path('/home/ubuntu/mstok-runs/iclr-debug-two-stage-20260916')
+export=root.parent.parent/'exports/step-22000'
+env=dict(os.environ,CUDA_VISIBLE_DEVICES='7',OMP_NUM_THREADS='1',TORCHINDUCTOR_COMPILE_THREADS='2',
+         PYTORCH_ALLOC_CONF='expandable_segments:True',TOKENIZERS_PARALLELISM='false')
+
+def status(stage,**extra):
+    (root/'status.json').write_text(json.dumps(dict(stage=stage,updated_utc=datetime.datetime.now(datetime.timezone.utc).isoformat(),**extra),indent=2)+'\n')
+
+try:
+    for seed in range(5):
+        out=root/f'random-seed-{seed}.json'
+        command=[sys.executable,str(root/'limited_eval.py'),'--ckpt',str(export/'ncp.pt'),
+                 '--codec',str(export/'vqvae.pt'),'--validation-source','/home/ubuntu/data/full_owt/valid_gpt2.bin',
+                 '--ref-model','gpt2-large','--ref-dtype','float32','--seed',str(seed),
+                 '--n-samples','128','--batch-size','4','--temperature','1','--top-k','0','--top-p','1',
+                 '--save-all-generations','--out',str(out)]
+        (root/f'command-seed-{seed}.json').write_text(json.dumps(command,indent=2)+'\n')
+        status('running',seed=seed)
+        with (root/f'random-seed-{seed}.log').open('w') as log:
+            subprocess.run(command,cwd=checkout,env=env,stdout=log,stderr=subprocess.STDOUT,check=True,timeout=1800)
+        result=json.loads(out.read_text())
+        assert result['checkpoint_step']==22000 and result['seed']==seed
+        assert result['reference_dtype']=='float32' and not result['tf32_allowed']
+        assert result['protocol']['n_levels']==16
+        assert result['scored_samples']==128 and result['skipped_empty_samples']==0
+        print(json.dumps({k:result[k] for k in ('seed','mean_ppl','median_ppl','mean_entropy_nats','scored_samples')}),flush=True)
+    subprocess.run([sys.executable,'evaluation/summarize_gen_ppl.py','--input-glob',str(root/'random-seed-[0-4].json'),
+                    '--out',str(root/'summary.json')],cwd=checkout,env=dict(env,CUDA_VISIBLE_DEVICES=''),check=True)
+    summary=json.loads((root/'summary.json').read_text())
+    summary.update(reference_dtype='float32',tf32_allowed=False,training_continued_during_evaluation=True)
+    (root/'summary.json').write_text(json.dumps(summary,indent=2)+'\n')
+    status('fp32_complete',seeds=list(range(5)),total_samples=640)
+    (root/'FP32_DONE.json').write_text(json.dumps(dict(step=22000,reference_dtype='float32',sampling='supplied-level0-untruncated',manual=True))+'\n')
+except BaseException as exc:
+    status('failed',error=str(exc))
+    raise

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/summary-fp64.json

@@ -0,0 +1,38 @@
+{
+  "protocol": {
+    "sampling": "random",
+    "temperature": 1.0,
+    "top_k": 0,
+    "top_p": 1.0,
+    "n_provided_levels": 1,
+    "n_levels": 16,
+    "document_aware_validation_sampling": true
+  },
+  "reference_model": "gpt2-large",
+  "checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "checkpoint_step": 22000,
+  "seeds": [
+    0,
+    1,
+    2,
+    3,
+    4
+  ],
+  "samples_per_seed": 128,
+  "total_scored_samples": 640,
+  "total_scored_tokens": 160526,
+  "per_seed_mean_ppl": {
+    "0": 190.84481863176575,
+    "1": 190.74650059895347,
+    "2": 201.55483121276038,
+    "3": 190.96062394498597,
+    "4": 172.3975097731799
+  },
+  "mean_gen_ppl": 189.3008568323291,
+  "gen_ppl_se": 4.706971216847786,
+  "mean_of_seed_median_ppl": 179.9071451840254,
+  "mean_entropy_nats": 4.472143790612192,
+  "reference_dtype": "float64",
+  "tf32_allowed": false,
+  "training_continued_during_evaluation": true
+}

iclr-debug-two-stage-135b-20260916/generator/evaluation/step-22000-fp32-fp64/summary.json

@@ -0,0 +1,38 @@
+{
+  "protocol": {
+    "sampling": "random",
+    "temperature": 1.0,
+    "top_k": 0,
+    "top_p": 1.0,
+    "n_provided_levels": 1,
+    "n_levels": 16,
+    "document_aware_validation_sampling": true
+  },
+  "reference_model": "gpt2-large",
+  "checkpoint": "/home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt",
+  "checkpoint_step": 22000,
+  "seeds": [
+    0,
+    1,
+    2,
+    3,
+    4
+  ],
+  "samples_per_seed": 128,
+  "total_scored_samples": 640,
+  "total_scored_tokens": 160526,
+  "per_seed_mean_ppl": {
+    "0": 190.84469984159884,
+    "1": 190.74649366523607,
+    "2": 201.5548702336531,
+    "3": 190.96056600638894,
+    "4": 172.39751630645821
+  },
+  "mean_gen_ppl": 189.30082921066705,
+  "gen_ppl_se": 4.7069720468787,
+  "mean_of_seed_median_ppl": 179.90716247558595,
+  "mean_entropy_nats": 4.472143790612192,
+  "reference_dtype": "float32",
+  "tf32_allowed": false,
+  "training_continued_during_evaluation": true
+}

iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp-config.yaml

@@ -0,0 +1,179 @@
+task: next-concept
+experiment: iclr-debug-two-stage-135b-20260916-generator
+experiment_dir: /home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000
+model:
+  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
+tokenizer:
+  pre_tokenizer:
+    name: hf
+    tokenizer: hf
+    model_id: gpt2
+    special_tokens:
+      eos_token: <|endoftext|>
+      pad_token: <|pad|>
+  vqvae:
+    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
+    checkpoint_path: /home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt
+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
+wandb:
+  entity: mstok
+  project: owt-repro
+  group: owtsmall-repro-ctx256-sharedhead-ncp
+  enable: false
+  id: owtsmall-repro-ctx256-sharedhead-ncp-5ep
+  resume: allow
+  gradients_and_params:
+    enable: false
+    log: all
+    log_freq: 1000
+optimization:
+  lr: 0.0005
+  reset_lr: true
+  warmup_iters: 300
+  lr_decay_iters: null
+  min_lr: 1.0e-05
+  beta_1: 0.9
+  beta_2: 0.99
+  weight_decay: 0.1
+  max_grad_norm: 1.0
+  grad_accumulation_steps: 4
+  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
+  level_loss_alpha: 1.0
+training:
+  seed: 55
+  log_interval: 10
+  eval_interval: 5000
+  checkpoint_interval: 2500
+  val_iters: 100
+  total_iters: null
+  tqdm_interval: 10
+  n_epochs: 5
+  batch_size: 96
+  eval_batch_size: 96
+  demo_batch_size: 4
+  keep_last: 8
+torch_compile:
+  enable: false
+  mode: default
+  dynamic: false
+  fullgraph: false
+  backend: inductor
+mixed_precision:
+  enable: true

iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/ncp.pt

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

iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/training-config.yaml

@@ -0,0 +1,223 @@
+task: mstok-next-concept
+experiment: iclr-debug-two-stage-135b-20260916-generator
+experiment_dir: /home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/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: 128747
+  generator_lr: 0.0005
+  generator_min_lr: 1.0e-05
+  generator_warmup_iters: 225
+  generator_lr_decay_iters: 128747
+  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-two-stage-135b-20260916-generator
+  resume: allow
+  gradients_and_params:
+    enable: false
+    log: all
+    log_freq: 1000
+export:
+  vqvae_config_template: /home/ubuntu/mstok-runs/iclr-debug-two-stage-20260916/config/repro-ctx256/vqvae.yaml
+  ncp_config_template: /home/ubuntu/mstok-runs/iclr-debug-two-stage-20260916/config/repro-ctx256/ncp-sharedhead.yaml
+iclr_debug:
+  stage: generator
+  target_positions: 135000000000
+  lr_horizon_positions: 135000000000
+  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-two-stage-135b-20260916/generator/exports/step-22000/vqvae-config.yaml

@@ -0,0 +1,117 @@
+task: summarization
+experiment: iclr-debug-two-stage-135b-20260916-generator
+experiment_dir: /home/ubuntu/mstok-results/iclr-debug-two-stage-135b-20260916/generator/exports/step-22000
+training:
+  seed: 42
+  log_interval: 10
+  eval_interval: 10000
+  checkpoint_interval: 10000
+  val_iters: 100
+  n_epochs: 10
+  total_iters: null
+  tqdm_interval: 10
+  batch_size: 256
+  eval_batch_size: 128
+  demo_batch_size: 4
+model:
+  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:
+    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
+wandb:
+  entity: mstok
+  project: owt-repro
+  enable: false
+  group: owtsmall-repro-ctx256-vqvae
+  id: owtsmall-repro-ctx256-vqvae-10ep
+  resume: allow
+  gradients_and_params:
+    enable: false
+    log: all
+    log_freq: 1000
+optimization:
+  lr: 0.001
+  lr_decay_iters: null
+  min_lr: null
+  beta_1: 0.9
+  beta_2: 0.99
+  weight_decay: 0.1
+  max_grad_norm: 1.0
+  grad_accumulation_steps: 2
+torch_compile:
+  enable: false
+  mode: default
+  dynamic: false
+  fullgraph: false
+  backend: inductor
+mixed_precision:
+  enable: true
+  opt_level: O1
+  loss_scale: dynamic
+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|>

iclr-debug-two-stage-135b-20260916/generator/exports/step-22000/vqvae.pt

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