Add files using upload-large-folder tool

pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101/finalize.json

@@ -0,0 +1,29 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101",
+ "w_bits": null,
+ "a_bits": null,
+ "qkv_bits": null,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 4.0493755096973
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 4.0493755096973
+ },
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.7319828271865845,
+ "loss_eval_model": 4.0493755096973,
+ "checks": {
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.03457271680235863,
+  "blocks.0.attention.w_out.weight": 0.03741328418254852,
+  "blocks.0.attention.w_q.weight": 0.036707207560539246,
+  "blocks.0.attention.w_v.weight": 0.03474367409944534,
+  "blocks.0.attention_norm.weight": 0.7402628660202026,
+  "blocks.0.feed_forward.w1.weight": 0.03897462412714958,
+  "blocks.0.feed_forward.w2.weight": 0.03813355043530464,
+  "blocks.0.feed_forward.w3.weight": 0.03718235716223717,
+  "blocks.0.feed_forward_norm.weight": 0.7393606305122375,
+  "blocks.1.attention.w_k.weight": 0.03156030550599098,
+  "blocks.1.attention.w_out.weight": 0.0387837290763855,
+  "blocks.1.attention.w_q.weight": 0.03309893608093262,
+  "blocks.1.attention.w_v.weight": 0.04152149334549904,
+  "blocks.1.attention_norm.weight": 0.6954712271690369,
+  "blocks.1.feed_forward.w1.weight": 0.0488145537674427,
+  "blocks.1.feed_forward.w2.weight": 0.04625376686453819,
+  "blocks.1.feed_forward.w3.weight": 0.04682345688343048,
+  "blocks.1.feed_forward_norm.weight": 0.7512289881706238,
+  "blocks.2.attention.w_k.weight": 0.03609790280461311,
+  "blocks.2.attention.w_out.weight": 0.0422712042927742,
+  "blocks.2.attention.w_q.weight": 0.0382913276553154,
+  "blocks.2.attention.w_v.weight": 0.04444428160786629,
+  "blocks.2.attention_norm.weight": 0.6940568089485168,
+  "blocks.2.feed_forward.w1.weight": 0.051162414252758026,
+  "blocks.2.feed_forward.w2.weight": 0.04817779362201691,
+  "blocks.2.feed_forward.w3.weight": 0.04966016858816147,
+  "blocks.2.feed_forward_norm.weight": 0.7518685460090637,
+  "blocks.3.attention.w_k.weight": 0.03685362637042999,
+  "blocks.3.attention.w_out.weight": 0.042797308415174484,
+  "blocks.3.attention.w_q.weight": 0.040169961750507355,
+  "blocks.3.attention.w_v.weight": 0.041635192930698395,
+  "blocks.3.attention_norm.weight": 0.7284940481185913,
+  "blocks.3.feed_forward.w1.weight": 0.050651222467422485,
+  "blocks.3.feed_forward.w2.weight": 0.04761878401041031,
+  "blocks.3.feed_forward.w3.weight": 0.04890414699912071,
+  "blocks.3.feed_forward_norm.weight": 0.7569454908370972,
+  "blocks.4.attention.w_k.weight": 0.040658239275217056,
+  "blocks.4.attention.w_out.weight": 0.04410743713378906,
+  "blocks.4.attention.w_q.weight": 0.04254842922091484,
+  "blocks.4.attention.w_v.weight": 0.04303383082151413,
+  "blocks.4.attention_norm.weight": 0.703934371471405,
+  "blocks.4.feed_forward.w1.weight": 0.050434891134500504,
+  "blocks.4.feed_forward.w2.weight": 0.04856543242931366,
+  "blocks.4.feed_forward.w3.weight": 0.048799481242895126,
+  "blocks.4.feed_forward_norm.weight": 0.7514355182647705,
+  "blocks.5.attention.w_k.weight": 0.03925139456987381,
+  "blocks.5.attention.w_out.weight": 0.04608171060681343,
+  "blocks.5.attention.w_q.weight": 0.04282260686159134,
+  "blocks.5.attention.w_v.weight": 0.04206324368715286,
+  "blocks.5.attention_norm.weight": 0.7227079272270203,
+  "blocks.5.feed_forward.w1.weight": 0.05082806572318077,
+  "blocks.5.feed_forward.w2.weight": 0.05065442994236946,
+  "blocks.5.feed_forward.w3.weight": 0.05076573044061661,
+  "blocks.5.feed_forward_norm.weight": 0.7479311227798462,
+  "blocks.6.attention.w_k.weight": 0.039561763405799866,
+  "blocks.6.attention.w_out.weight": 0.05980084836483002,
+  "blocks.6.attention.w_q.weight": 0.04352796822786331,
+  "blocks.6.attention.w_v.weight": 0.04301363602280617,
+  "blocks.6.attention_norm.weight": 0.7296965718269348,
+  "blocks.6.feed_forward.w1.weight": 0.051500093191862106,
+  "blocks.6.feed_forward.w2.weight": 0.054001640528440475,
+  "blocks.6.feed_forward.w3.weight": 0.051999837160110474,
+  "blocks.6.feed_forward_norm.weight": 0.7462171912193298,
+  "blocks.7.attention.w_k.weight": 0.042883116751909256,
+  "blocks.7.attention.w_out.weight": 0.06064591929316521,
+  "blocks.7.attention.w_q.weight": 0.04550102725625038,
+  "blocks.7.attention.w_v.weight": 0.043729517608881,
+  "blocks.7.attention_norm.weight": 0.7210500836372375,
+  "blocks.7.feed_forward.w1.weight": 0.05071728676557541,
+  "blocks.7.feed_forward.w2.weight": 0.05820971727371216,
+  "blocks.7.feed_forward.w3.weight": 0.052353113889694214,
+  "blocks.7.feed_forward_norm.weight": 0.741530179977417,
+  "blocks.8.attention.w_k.weight": 0.040963493287563324,
+  "blocks.8.attention.w_out.weight": 0.08356135338544846,
+  "blocks.8.attention.w_q.weight": 0.04410872235894203,
+  "blocks.8.attention.w_v.weight": 0.04232929274439812,
+  "blocks.8.attention_norm.weight": 0.7226449251174927,
+  "blocks.8.feed_forward.w1.weight": 0.049077108502388,
+  "blocks.8.feed_forward.w2.weight": 0.06283104419708252,
+  "blocks.8.feed_forward.w3.weight": 0.05128471553325653,
+  "blocks.8.feed_forward_norm.weight": 0.7345075011253357,
+  "embeddings.weight": 0.05591508001089096,
+  "lm_head.norm.weight": 1.3878247737884521,
+  "lm_head.w_out.weight": 0.06863079965114594
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 4.0493755096973
+ },
+ "base_loss": 4.0493755096973,
+ "loss": {
+  "0.005": [
+   4.0495698425518905,
+   4.049669285950546,
+   4.049879588903228
+  ],
+  "0.01": [
+   4.050270418126432,
+   4.050327205479154,
+   4.0508540152754575
+  ],
+  "0.02": [
+   4.0527773479440095,
+   4.053184565883751,
+   4.05434469844073
+  ],
+  "0.05": [
+   4.070846207219855,
+   4.073331396264503,
+   4.075627577124697
+  ],
+  "0.1": [
+   4.139764054437487,
+   4.152049371249179,
+   4.152491287572062
+  ],
+  "0.15": [
+   4.269614547246076,
+   4.3033789691601365,
+   4.291750445128339
+  ],
+  "0.2": [
+   4.481975479068514,
+   4.55160124920126,
+   4.511998031100518
+  ],
+  "0.3": [
+   5.234732680119546,
+   5.418616669633772,
+   5.279652465939638
+  ],
+  "0.4": [
+   6.4647297840623645,
+   6.775691639048803,
+   6.493282746697875
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    4.0495698425518905,
+    4.049669285950546,
+    4.049879588903228
+   ],
+   "0.01": [
+    4.050270418126432,
+    4.050327205479154,
+    4.0508540152754575
+   ],
+   "0.02": [
+    4.0527773479440095,
+    4.053184565883751,
+    4.05434469844073
+   ],
+   "0.05": [
+    4.070846207219855,
+    4.073331396264503,
+    4.075627577124697
+   ],
+   "0.1": [
+    4.139764054437487,
+    4.152049371249179,
+    4.152491287572062
+   ],
+   "0.15": [
+    4.269614547246076,
+    4.3033789691601365,
+    4.291750445128339
+   ],
+   "0.2": [
+    4.481975479068514,
+    4.55160124920126,
+    4.511998031100518
+   ],
+   "0.3": [
+    5.234732680119546,
+    5.418616669633772,
+    5.279652465939638
+   ],
+   "0.4": [
+    6.4647297840623645,
+    6.775691639048803,
+    6.493282746697875
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.00033072943792156195,
+   "0.01": 0.0011083699297156357,
+   "0.02": 0.004060027725530269,
+   "0.05": 0.02389288383905175,
+   "0.1": 0.09872606138894291,
+   "0.15": 0.23887247748088392,
+   "0.2": 0.4658160767594648,
+   "0.3": 1.2616250955336845,
+   "0.4": 2.528525880239048
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.00033072943792156195,
+  "0.01": 0.0011083699297156357,
+  "0.02": 0.004060027725530269,
+  "0.05": 0.02389288383905175,
+  "0.1": 0.09872606138894291,
+  "0.15": 0.23887247748088392,
+  "0.2": 0.4658160767594648,
+  "0.3": 1.2616250955336845,
+  "0.4": 2.528525880239048
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 56.3
+}

pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.03457271680235863,
+  "blocks.0.attention.w_out.weight": 0.03741328418254852,
+  "blocks.0.attention.w_q.weight": 0.036707207560539246,
+  "blocks.0.attention.w_v.weight": 0.03474367409944534,
+  "blocks.0.feed_forward.w1.weight": 0.03897462412714958,
+  "blocks.0.feed_forward.w2.weight": 0.03813355043530464,
+  "blocks.0.feed_forward.w3.weight": 0.03718235716223717,
+  "blocks.1.attention.w_k.weight": 0.03156030550599098,
+  "blocks.1.attention.w_out.weight": 0.0387837290763855,
+  "blocks.1.attention.w_q.weight": 0.03309893608093262,
+  "blocks.1.attention.w_v.weight": 0.04152149334549904,
+  "blocks.1.feed_forward.w1.weight": 0.0488145537674427,
+  "blocks.1.feed_forward.w2.weight": 0.04625376686453819,
+  "blocks.1.feed_forward.w3.weight": 0.04682345688343048,
+  "blocks.2.attention.w_k.weight": 0.03609790280461311,
+  "blocks.2.attention.w_out.weight": 0.0422712042927742,
+  "blocks.2.attention.w_q.weight": 0.0382913276553154,
+  "blocks.2.attention.w_v.weight": 0.04444428160786629,
+  "blocks.2.feed_forward.w1.weight": 0.051162414252758026,
+  "blocks.2.feed_forward.w2.weight": 0.04817779362201691,
+  "blocks.2.feed_forward.w3.weight": 0.04966016858816147,
+  "blocks.3.attention.w_k.weight": 0.03685362637042999,
+  "blocks.3.attention.w_out.weight": 0.042797308415174484,
+  "blocks.3.attention.w_q.weight": 0.040169961750507355,
+  "blocks.3.attention.w_v.weight": 0.041635192930698395,
+  "blocks.3.feed_forward.w1.weight": 0.050651222467422485,
+  "blocks.3.feed_forward.w2.weight": 0.04761878401041031,
+  "blocks.3.feed_forward.w3.weight": 0.04890414699912071,
+  "blocks.4.attention.w_k.weight": 0.040658239275217056,
+  "blocks.4.attention.w_out.weight": 0.04410743713378906,
+  "blocks.4.attention.w_q.weight": 0.04254842922091484,
+  "blocks.4.attention.w_v.weight": 0.04303383082151413,
+  "blocks.4.feed_forward.w1.weight": 0.050434891134500504,
+  "blocks.4.feed_forward.w2.weight": 0.04856543242931366,
+  "blocks.4.feed_forward.w3.weight": 0.048799481242895126,
+  "blocks.5.attention.w_k.weight": 0.03925139456987381,
+  "blocks.5.attention.w_out.weight": 0.04608171060681343,
+  "blocks.5.attention.w_q.weight": 0.04282260686159134,
+  "blocks.5.attention.w_v.weight": 0.04206324368715286,
+  "blocks.5.feed_forward.w1.weight": 0.05082806572318077,
+  "blocks.5.feed_forward.w2.weight": 0.05065442994236946,
+  "blocks.5.feed_forward.w3.weight": 0.05076573044061661,
+  "blocks.6.attention.w_k.weight": 0.039561763405799866,
+  "blocks.6.attention.w_out.weight": 0.05980084836483002,
+  "blocks.6.attention.w_q.weight": 0.04352796822786331,
+  "blocks.6.attention.w_v.weight": 0.04301363602280617,
+  "blocks.6.feed_forward.w1.weight": 0.051500093191862106,
+  "blocks.6.feed_forward.w2.weight": 0.054001640528440475,
+  "blocks.6.feed_forward.w3.weight": 0.051999837160110474,
+  "blocks.7.attention.w_k.weight": 0.042883116751909256,
+  "blocks.7.attention.w_out.weight": 0.06064591929316521,
+  "blocks.7.attention.w_q.weight": 0.04550102725625038,
+  "blocks.7.attention.w_v.weight": 0.043729517608881,
+  "blocks.7.feed_forward.w1.weight": 0.05071728676557541,
+  "blocks.7.feed_forward.w2.weight": 0.05820971727371216,
+  "blocks.7.feed_forward.w3.weight": 0.052353113889694214,
+  "blocks.8.attention.w_k.weight": 0.040963493287563324,
+  "blocks.8.attention.w_out.weight": 0.08356135338544846,
+  "blocks.8.attention.w_q.weight": 0.04410872235894203,
+  "blocks.8.attention.w_v.weight": 0.04232929274439812,
+  "blocks.8.feed_forward.w1.weight": 0.049077108502388,
+  "blocks.8.feed_forward.w2.weight": 0.06283104419708252,
+  "blocks.8.feed_forward.w3.weight": 0.05128471553325653
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 4.0493755096973
+ },
+ "base_loss": 4.0493755096973,
+ "loss": {
+  "0.005": [
+   4.049424695408524,
+   4.04958523753866,
+   4.049485700244279
+  ],
+  "0.01": [
+   4.049717063225876,
+   4.049925789512824,
+   4.049838928141476
+  ],
+  "0.02": [
+   4.050981975588614,
+   4.051243275657571,
+   4.0511587695040205
+  ],
+  "0.05": [
+   4.060081461271643,
+   4.0611475455219,
+   4.060970325163404
+  ],
+  "0.1": [
+   4.095681630703352,
+   4.097457234096067,
+   4.097584016800324
+  ],
+  "0.15": [
+   4.163234809173837,
+   4.163216593402962,
+   4.165194847246677
+  ],
+  "0.2": [
+   4.273525844141357,
+   4.267101013429092,
+   4.272839509418618
+  ],
+  "0.3": [
+   4.675111350941662,
+   4.636568509013489,
+   4.653812177436882
+  ],
+  "0.4": [
+   5.35069790032881,
+   5.268898275435976,
+   5.295578284372346
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    4.049424695408524,
+    4.04958523753866,
+    4.049485700244279
+   ],
+   "0.01": [
+    4.049717063225876,
+    4.049925789512824,
+    4.049838928141476
+   ],
+   "0.02": [
+    4.050981975588614,
+    4.051243275657571,
+    4.0511587695040205
+   ],
+   "0.05": [
+    4.060081461271643,
+    4.0611475455219,
+    4.060970325163404
+   ],
+   "0.1": [
+    4.095681630703352,
+    4.097457234096067,
+    4.097584016800324
+   ],
+   "0.15": [
+    4.163234809173837,
+    4.163216593402962,
+    4.165194847246677
+   ],
+   "0.2": [
+    4.273525844141357,
+    4.267101013429092,
+    4.272839509418618
+   ],
+   "0.3": [
+    4.675111350941662,
+    4.636568509013489,
+    4.653812177436882
+   ],
+   "0.4": [
+    5.35069790032881,
+    5.268898275435976,
+    5.295578284372346
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.00012303469985397442,
+   "0.01": 0.0004517505960919266,
+   "0.02": 0.0017524972194351562,
+   "0.05": 0.011357600955015812,
+   "0.1": 0.04753211750261421,
+   "0.15": 0.11450657357719152,
+   "0.2": 0.22177994596572148,
+   "0.3": 0.6057885027667105,
+   "0.4": 1.2556826436817445
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.00012303469985397442,
+  "0.01": 0.0004517505960919266,
+  "0.02": 0.0017524972194351562,
+  "0.05": 0.011357600955015812,
+  "0.1": 0.04753211750261421,
+  "0.15": 0.11450657357719152,
+  "0.2": 0.22177994596572148,
+  "0.3": 0.6057885027667105,
+  "0.4": 1.2556826436817445
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 50.2
+}

pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101/ptq.json

@@ -0,0 +1,48 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-adamw-lr2.4e-03-s101",
+ "w_bits": null,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 4.0493755096973
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 4.82802165992871
+  },
+  "4": {
+   "val_dolma_cc": 4.154064397006654
+  },
+  "5": {
+   "val_dolma_cc": 4.072646571698026
+  },
+  "6": {
+   "val_dolma_cc": 4.054745805333628
+  },
+  "8": {
+   "val_dolma_cc": 4.049746444763523
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 0.7786461502314097
+  },
+  "4": {
+   "val_dolma_cc": 0.10468888730935433
+  },
+  "5": {
+   "val_dolma_cc": 0.023271062000725706
+  },
+  "6": {
+   "val_dolma_cc": 0.00537029563632796
+  },
+  "8": {
+   "val_dolma_cc": 0.0003709350662228772
+  }
+ },
+ "checks": {
+  "weights_restored": true
+ },
+ "ok": true
+}

pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": null,
+  "a": null,
+  "qkv": null
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 3.916410182665998,
+  "heldout_cc_en_tail_4M": 4.494598077772183
+ },
+ "loss_bf16_master": 3.916410182665998,
+ "losses_eval_model": {
+  "val_dolma_cc": 3.916410182665998
+ },
+ "bf16_minus_quant": 0.0,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-n5-28",
+ "gpu": "NVIDIA L40S",
+ "t_s": 18.2
+}

pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/finalize.json

@@ -0,0 +1,29 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "w_bits": null,
+ "a_bits": null,
+ "qkv_bits": null,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.916410182665998
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.916410182665998
+ },
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.5616782903671265,
+ "loss_eval_model": 3.916410182665998,
+ "checks": {
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/launch.yaml

@@ -0,0 +1,126 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 101
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: null
+  fq_act_bits: null
+  fq_qkv_bits: null
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: precq.muon26.MuonConfig
+    lr: 0.01
+    weight_decay: 0.1
+    momentum: 0.95
+    nesterov: true
+    ns_steps: 5
+    eps: 1.0e-07
+    adjust_lr_fn: null
+    adamw_lr: 0.0024
+    adamw_betas:
+    - 0.9
+    - 0.95
+    adamw_eps: 1.0e-15
+    adamw_weight_decay: 0.1
+    adamw_fused: true
+    muon_on_lm_head: false
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 232
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1002327
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 101
+n_tokens: 610009088
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.033792030066251755,
+  "blocks.0.attention.w_out.weight": 0.026513146236538887,
+  "blocks.0.attention.w_q.weight": 0.033974722027778625,
+  "blocks.0.attention.w_v.weight": 0.030098984017968178,
+  "blocks.0.attention_norm.weight": 0.7212387323379517,
+  "blocks.0.feed_forward.w1.weight": 0.03262019529938698,
+  "blocks.0.feed_forward.w2.weight": 0.018550856038928032,
+  "blocks.0.feed_forward.w3.weight": 0.03300027176737785,
+  "blocks.0.feed_forward_norm.weight": 0.7314680218696594,
+  "blocks.1.attention.w_k.weight": 0.028769828379154205,
+  "blocks.1.attention.w_out.weight": 0.024300970137119293,
+  "blocks.1.attention.w_q.weight": 0.0290773157030344,
+  "blocks.1.attention.w_v.weight": 0.02681984379887581,
+  "blocks.1.attention_norm.weight": 0.7247515916824341,
+  "blocks.1.feed_forward.w1.weight": 0.035662610083818436,
+  "blocks.1.feed_forward.w2.weight": 0.019625242799520493,
+  "blocks.1.feed_forward.w3.weight": 0.036090947687625885,
+  "blocks.1.feed_forward_norm.weight": 0.7420072555541992,
+  "blocks.2.attention.w_k.weight": 0.031684815883636475,
+  "blocks.2.attention.w_out.weight": 0.0261231679469347,
+  "blocks.2.attention.w_q.weight": 0.03192731738090515,
+  "blocks.2.attention.w_v.weight": 0.029027849435806274,
+  "blocks.2.attention_norm.weight": 0.7313526272773743,
+  "blocks.2.feed_forward.w1.weight": 0.035158202052116394,
+  "blocks.2.feed_forward.w2.weight": 0.017811624333262444,
+  "blocks.2.feed_forward.w3.weight": 0.035606566816568375,
+  "blocks.2.feed_forward_norm.weight": 0.7424768209457397,
+  "blocks.3.attention.w_k.weight": 0.03212268650531769,
+  "blocks.3.attention.w_out.weight": 0.02810906246304512,
+  "blocks.3.attention.w_q.weight": 0.03197857737541199,
+  "blocks.3.attention.w_v.weight": 0.03150033950805664,
+  "blocks.3.attention_norm.weight": 0.7306333184242249,
+  "blocks.3.feed_forward.w1.weight": 0.03466830775141716,
+  "blocks.3.feed_forward.w2.weight": 0.018097171559929848,
+  "blocks.3.feed_forward.w3.weight": 0.03538702800869942,
+  "blocks.3.feed_forward_norm.weight": 0.7410686016082764,
+  "blocks.4.attention.w_k.weight": 0.03302117437124252,
+  "blocks.4.attention.w_out.weight": 0.02875415049493313,
+  "blocks.4.attention.w_q.weight": 0.032805606722831726,
+  "blocks.4.attention.w_v.weight": 0.031778451055288315,
+  "blocks.4.attention_norm.weight": 0.7310084104537964,
+  "blocks.4.feed_forward.w1.weight": 0.0345308855175972,
+  "blocks.4.feed_forward.w2.weight": 0.019194750115275383,
+  "blocks.4.feed_forward.w3.weight": 0.03569754585623741,
+  "blocks.4.feed_forward_norm.weight": 0.7415932416915894,
+  "blocks.5.attention.w_k.weight": 0.03010912798345089,
+  "blocks.5.attention.w_out.weight": 0.03453635796904564,
+  "blocks.5.attention.w_q.weight": 0.030799981206655502,
+  "blocks.5.attention.w_v.weight": 0.03295043483376503,
+  "blocks.5.attention_norm.weight": 0.7344293594360352,
+  "blocks.5.feed_forward.w1.weight": 0.03447722643613815,
+  "blocks.5.feed_forward.w2.weight": 0.020819514989852905,
+  "blocks.5.feed_forward.w3.weight": 0.03604985773563385,
+  "blocks.5.feed_forward_norm.weight": 0.742188036441803,
+  "blocks.6.attention.w_k.weight": 0.03257782757282257,
+  "blocks.6.attention.w_out.weight": 0.036244895309209824,
+  "blocks.6.attention.w_q.weight": 0.0317949503660202,
+  "blocks.6.attention.w_v.weight": 0.034245118498802185,
+  "blocks.6.attention_norm.weight": 0.7350813746452332,
+  "blocks.6.feed_forward.w1.weight": 0.03427771106362343,
+  "blocks.6.feed_forward.w2.weight": 0.02320433221757412,
+  "blocks.6.feed_forward.w3.weight": 0.036269668489694595,
+  "blocks.6.feed_forward_norm.weight": 0.7430891990661621,
+  "blocks.7.attention.w_k.weight": 0.033718906342983246,
+  "blocks.7.attention.w_out.weight": 0.040113773196935654,
+  "blocks.7.attention.w_q.weight": 0.0327695831656456,
+  "blocks.7.attention.w_v.weight": 0.034865327179431915,
+  "blocks.7.attention_norm.weight": 0.7363488674163818,
+  "blocks.7.feed_forward.w1.weight": 0.033236388117074966,
+  "blocks.7.feed_forward.w2.weight": 0.02615497075021267,
+  "blocks.7.feed_forward.w3.weight": 0.03580132871866226,
+  "blocks.7.feed_forward_norm.weight": 0.7429885268211365,
+  "blocks.8.attention.w_k.weight": 0.03233204782009125,
+  "blocks.8.attention.w_out.weight": 0.04406481236219406,
+  "blocks.8.attention.w_q.weight": 0.031116899102926254,
+  "blocks.8.attention.w_v.weight": 0.03517533093690872,
+  "blocks.8.attention_norm.weight": 0.7402186989784241,
+  "blocks.8.feed_forward.w1.weight": 0.03163241967558861,
+  "blocks.8.feed_forward.w2.weight": 0.030139151960611343,
+  "blocks.8.feed_forward.w3.weight": 0.03480131924152374,
+  "blocks.8.feed_forward_norm.weight": 0.741178035736084,
+  "embeddings.weight": 0.055930908769369125,
+  "lm_head.norm.weight": 1.363660454750061,
+  "lm_head.w_out.weight": 0.06779012084007263
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.916410182665998
+ },
+ "base_loss": 3.916410182665998,
+ "loss": {
+  "0.005": [
+   3.916753880146988,
+   3.916703325123091,
+   3.916709217079908
+  ],
+  "0.01": [
+   3.9176285853457067,
+   3.9174146400213132,
+   3.917594133438847
+  ],
+  "0.02": [
+   3.9209808701229307,
+   3.92055286992913,
+   3.920792448786767
+  ],
+  "0.05": [
+   3.94461695175913,
+   3.9424623265307757,
+   3.944583939586739
+  ],
+  "0.1": [
+   4.036349081743616,
+   4.027967018093331,
+   4.03692925952328
+  ],
+  "0.15": [
+   4.225827218217674,
+   4.199057738391233,
+   4.224654679687876
+  ],
+  "0.2": [
+   4.583495436479917,
+   4.510621069837889,
+   4.570480080397276
+  ],
+  "0.3": [
+   6.054807350069808,
+   5.9142904866477775,
+   6.009911476361313
+  ],
+  "0.4": [
+   7.959366780994102,
+   7.961219625358557,
+   7.720636466859065
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.916753880146988,
+    3.916703325123091,
+    3.916709217079908
+   ],
+   "0.01": [
+    3.9176285853457067,
+    3.9174146400213132,
+    3.917594133438847
+   ],
+   "0.02": [
+    3.9209808701229307,
+    3.92055286992913,
+    3.920792448786767
+   ],
+   "0.05": [
+    3.94461695175913,
+    3.9424623265307757,
+    3.944583939586739
+   ],
+   "0.1": [
+    4.036349081743616,
+    4.027967018093331,
+    4.03692925952328
+   ],
+   "0.15": [
+    4.225827218217674,
+    4.199057738391233,
+    4.224654679687876
+   ],
+   "0.2": [
+    4.583495436479917,
+    4.510621069837889,
+    4.570480080397276
+   ],
+   "0.3": [
+    6.054807350069808,
+    5.9142904866477775,
+    6.009911476361313
+   ],
+   "0.4": [
+    7.959366780994102,
+    7.961219625358557,
+    7.720636466859065
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.00031195811733164547,
+   "0.01": 0.0011356036026244354,
+   "0.02": 0.004365213613611196,
+   "0.05": 0.027477556626217314,
+   "0.1": 0.11733827045407796,
+   "0.15": 0.3001030294329299,
+   "0.2": 0.6384553462390299,
+   "0.3": 2.076592921693635,
+   "0.4": 3.9639974417379102
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.00031195811733164547,
+  "0.01": 0.0011356036026244354,
+  "0.02": 0.004365213613611196,
+  "0.05": 0.027477556626217314,
+  "0.1": 0.11733827045407796,
+  "0.15": 0.3001030294329299,
+  "0.2": 0.6384553462390299,
+  "0.3": 2.076592921693635,
+  "0.4": 3.9639974417379102
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 61.2
+}

pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.033792030066251755,
+  "blocks.0.attention.w_out.weight": 0.026513146236538887,
+  "blocks.0.attention.w_q.weight": 0.033974722027778625,
+  "blocks.0.attention.w_v.weight": 0.030098984017968178,
+  "blocks.0.feed_forward.w1.weight": 0.03262019529938698,
+  "blocks.0.feed_forward.w2.weight": 0.018550856038928032,
+  "blocks.0.feed_forward.w3.weight": 0.03300027176737785,
+  "blocks.1.attention.w_k.weight": 0.028769828379154205,
+  "blocks.1.attention.w_out.weight": 0.024300970137119293,
+  "blocks.1.attention.w_q.weight": 0.0290773157030344,
+  "blocks.1.attention.w_v.weight": 0.02681984379887581,
+  "blocks.1.feed_forward.w1.weight": 0.035662610083818436,
+  "blocks.1.feed_forward.w2.weight": 0.019625242799520493,
+  "blocks.1.feed_forward.w3.weight": 0.036090947687625885,
+  "blocks.2.attention.w_k.weight": 0.031684815883636475,
+  "blocks.2.attention.w_out.weight": 0.0261231679469347,
+  "blocks.2.attention.w_q.weight": 0.03192731738090515,
+  "blocks.2.attention.w_v.weight": 0.029027849435806274,
+  "blocks.2.feed_forward.w1.weight": 0.035158202052116394,
+  "blocks.2.feed_forward.w2.weight": 0.017811624333262444,
+  "blocks.2.feed_forward.w3.weight": 0.035606566816568375,
+  "blocks.3.attention.w_k.weight": 0.03212268650531769,
+  "blocks.3.attention.w_out.weight": 0.02810906246304512,
+  "blocks.3.attention.w_q.weight": 0.03197857737541199,
+  "blocks.3.attention.w_v.weight": 0.03150033950805664,
+  "blocks.3.feed_forward.w1.weight": 0.03466830775141716,
+  "blocks.3.feed_forward.w2.weight": 0.018097171559929848,
+  "blocks.3.feed_forward.w3.weight": 0.03538702800869942,
+  "blocks.4.attention.w_k.weight": 0.03302117437124252,
+  "blocks.4.attention.w_out.weight": 0.02875415049493313,
+  "blocks.4.attention.w_q.weight": 0.032805606722831726,
+  "blocks.4.attention.w_v.weight": 0.031778451055288315,
+  "blocks.4.feed_forward.w1.weight": 0.0345308855175972,
+  "blocks.4.feed_forward.w2.weight": 0.019194750115275383,
+  "blocks.4.feed_forward.w3.weight": 0.03569754585623741,
+  "blocks.5.attention.w_k.weight": 0.03010912798345089,
+  "blocks.5.attention.w_out.weight": 0.03453635796904564,
+  "blocks.5.attention.w_q.weight": 0.030799981206655502,
+  "blocks.5.attention.w_v.weight": 0.03295043483376503,
+  "blocks.5.feed_forward.w1.weight": 0.03447722643613815,
+  "blocks.5.feed_forward.w2.weight": 0.020819514989852905,
+  "blocks.5.feed_forward.w3.weight": 0.03604985773563385,
+  "blocks.6.attention.w_k.weight": 0.03257782757282257,
+  "blocks.6.attention.w_out.weight": 0.036244895309209824,
+  "blocks.6.attention.w_q.weight": 0.0317949503660202,
+  "blocks.6.attention.w_v.weight": 0.034245118498802185,
+  "blocks.6.feed_forward.w1.weight": 0.03427771106362343,
+  "blocks.6.feed_forward.w2.weight": 0.02320433221757412,
+  "blocks.6.feed_forward.w3.weight": 0.036269668489694595,
+  "blocks.7.attention.w_k.weight": 0.033718906342983246,
+  "blocks.7.attention.w_out.weight": 0.040113773196935654,
+  "blocks.7.attention.w_q.weight": 0.0327695831656456,
+  "blocks.7.attention.w_v.weight": 0.034865327179431915,
+  "blocks.7.feed_forward.w1.weight": 0.033236388117074966,
+  "blocks.7.feed_forward.w2.weight": 0.02615497075021267,
+  "blocks.7.feed_forward.w3.weight": 0.03580132871866226,
+  "blocks.8.attention.w_k.weight": 0.03233204782009125,
+  "blocks.8.attention.w_out.weight": 0.04406481236219406,
+  "blocks.8.attention.w_q.weight": 0.031116899102926254,
+  "blocks.8.attention.w_v.weight": 0.03517533093690872,
+  "blocks.8.feed_forward.w1.weight": 0.03163241967558861,
+  "blocks.8.feed_forward.w2.weight": 0.030139151960611343,
+  "blocks.8.feed_forward.w3.weight": 0.03480131924152374
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.916410182665998
+ },
+ "base_loss": 3.916410182665998,
+ "loss": {
+  "0.005": [
+   3.9166031188216612,
+   3.9166111077830292,
+   3.9166282124544796
+  ],
+  "0.01": [
+   3.9170695145428205,
+   3.917076439352625,
+   3.9171663953706153
+  ],
+  "0.02": [
+   3.9188807005034314,
+   3.918924248558763,
+   3.919003539587189
+  ],
+  "0.05": [
+   3.9309205856238885,
+   3.9306903689526145,
+   3.9308996703508847
+  ],
+  "0.1": [
+   3.9758266704926046,
+   3.973638188030323,
+   3.974508139757904
+  ],
+  "0.15": [
+   4.059841444233951,
+   4.0517076565329395,
+   4.055604896455549
+  ],
+  "0.2": [
+   4.1971037213818345,
+   4.177021170450288,
+   4.188559100755623
+  ],
+  "0.3": [
+   4.710997663686541,
+   4.645984176817421,
+   4.6946044558029065
+  ],
+  "0.4": [
+   5.60428553083135,
+   5.503434611767463,
+   5.5852172258845725
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.9166031188216612,
+    3.9166111077830292,
+    3.9166282124544796
+   ],
+   "0.01": [
+    3.9170695145428205,
+    3.917076439352625,
+    3.9171663953706153
+   ],
+   "0.02": [
+    3.9188807005034314,
+    3.918924248558763,
+    3.919003539587189
+   ],
+   "0.05": [
+    3.9309205856238885,
+    3.9306903689526145,
+    3.9308996703508847
+   ],
+   "0.1": [
+    3.9758266704926046,
+    3.973638188030323,
+    3.974508139757904
+   ],
+   "0.15": [
+    4.059841444233951,
+    4.0517076565329395,
+    4.055604896455549
+   ],
+   "0.2": [
+    4.1971037213818345,
+    4.177021170450288,
+    4.188559100755623
+   ],
+   "0.3": [
+    4.710997663686541,
+    4.645984176817421,
+    4.6946044558029065
+   ],
+   "0.4": [
+    5.60428553083135,
+    5.503434611767463,
+    5.5852172258845725
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.00020396368705899448,
+   "0.01": 0.0006939337560227798,
+   "0.02": 0.0025259802171300194,
+   "0.05": 0.014426692309798206,
+   "0.1": 0.05824748342761277,
+   "0.15": 0.1393078164081487,
+   "0.2": 0.2711511481965845,
+   "0.3": 0.7674519161029583,
+   "0.4": 1.6479022734951307
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.00020396368705899448,
+  "0.01": 0.0006939337560227798,
+  "0.02": 0.0025259802171300194,
+  "0.05": 0.014426692309798206,
+  "0.1": 0.05824748342761277,
+  "0.15": 0.1393078164081487,
+  "0.2": 0.2711511481965845,
+  "0.3": 0.7674519161029583,
+  "0.4": 1.6479022734951307
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 52.6
+}

pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/ptq.json

@@ -0,0 +1,48 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "w_bits": null,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.916410182665998
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 4.996126214447322
+  },
+  "4": {
+   "val_dolma_cc": 4.042027073674593
+  },
+  "5": {
+   "val_dolma_cc": 3.942661643683497
+  },
+  "6": {
+   "val_dolma_cc": 3.9220034376164308
+  },
+  "8": {
+   "val_dolma_cc": 3.9167275423957983
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 1.0797160317813246
+  },
+  "4": {
+   "val_dolma_cc": 0.12561689100859486
+  },
+  "5": {
+   "val_dolma_cc": 0.02625146101749909
+  },
+  "6": {
+   "val_dolma_cc": 0.0055932549504329465
+  },
+  "8": {
+   "val_dolma_cc": 0.00031735972980051486
+  }
+ },
+ "checks": {
+  "weights_restored": true
+ },
+ "ok": true
+}

pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": 6,
+  "a": 6,
+  "qkv": 6
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 4.116169639029774,
+  "heldout_cc_en_tail_4M": 4.678772911880955
+ },
+ "loss_bf16_master": 4.116169639029774,
+ "losses_eval_model": {
+  "val_dolma_cc": 4.133484409991807
+ },
+ "bf16_minus_quant": -0.01731477096203271,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-p9-24",
+ "gpu": "NVIDIA L40S",
+ "t_s": 17.5
+}

pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/finalize.json

@@ -0,0 +1,102 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102",
+ "w_bits": 6,
+ "a_bits": 6,
+ "qkv_bits": 6,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 4.1323714209140086
+ },
+ "losses_fakequant_master": {
+  "val_dolma_cc": 4.133484409991807
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 4.133484409991807
+ },
+ "max_distinct_values_per_row": 63,
+ "rel_quant_noise_per_layer": {
+  "blocks.0.attention.w_q": 0.031120842322707176,
+  "blocks.0.attention.w_k": 0.03135550022125244,
+  "blocks.0.attention.w_v": 0.03434007987380028,
+  "blocks.0.attention.w_out": 0.03542300686240196,
+  "blocks.0.feed_forward.w1": 0.031323470175266266,
+  "blocks.0.feed_forward.w2": 0.035301655530929565,
+  "blocks.0.feed_forward.w3": 0.03124716505408287,
+  "blocks.1.attention.w_q": 0.030890315771102905,
+  "blocks.1.attention.w_k": 0.03229247033596039,
+  "blocks.1.attention.w_v": 0.033395953476428986,
+  "blocks.1.attention.w_out": 0.032106492668390274,
+  "blocks.1.feed_forward.w1": 0.031076418235898018,
+  "blocks.1.feed_forward.w2": 0.034724216908216476,
+  "blocks.1.feed_forward.w3": 0.03132103011012077,
+  "blocks.2.attention.w_q": 0.031059447675943375,
+  "blocks.2.attention.w_k": 0.032405510544776917,
+  "blocks.2.attention.w_v": 0.033812735229730606,
+  "blocks.2.attention.w_out": 0.03228583559393883,
+  "blocks.2.feed_forward.w1": 0.03101040981709957,
+  "blocks.2.feed_forward.w2": 0.03550992161035538,
+  "blocks.2.feed_forward.w3": 0.03121618926525116,
+  "blocks.3.attention.w_q": 0.03107432834804058,
+  "blocks.3.attention.w_k": 0.03350537270307541,
+  "blocks.3.attention.w_v": 0.03315924108028412,
+  "blocks.3.attention.w_out": 0.03204391151666641,
+  "blocks.3.feed_forward.w1": 0.03128368780016899,
+  "blocks.3.feed_forward.w2": 0.03541150316596031,
+  "blocks.3.feed_forward.w3": 0.03160535916686058,
+  "blocks.4.attention.w_q": 0.031151646748185158,
+  "blocks.4.attention.w_k": 0.032543595880270004,
+  "blocks.4.attention.w_v": 0.033398281782865524,
+  "blocks.4.attention.w_out": 0.03176643326878548,
+  "blocks.4.feed_forward.w1": 0.031281664967536926,
+  "blocks.4.feed_forward.w2": 0.03540990874171257,
+  "blocks.4.feed_forward.w3": 0.031895920634269714,
+  "blocks.5.attention.w_q": 0.03117757849395275,
+  "blocks.5.attention.w_k": 0.033542435616254807,
+  "blocks.5.attention.w_v": 0.033953651785850525,
+  "blocks.5.attention.w_out": 0.032777994871139526,
+  "blocks.5.feed_forward.w1": 0.03133861720561981,
+  "blocks.5.feed_forward.w2": 0.03533758595585823,
+  "blocks.5.feed_forward.w3": 0.0314180962741375,
+  "blocks.6.attention.w_q": 0.031007658690214157,
+  "blocks.6.attention.w_k": 0.03384607285261154,
+  "blocks.6.attention.w_v": 0.03334997221827507,
+  "blocks.6.attention.w_out": 0.033462539315223694,
+  "blocks.6.feed_forward.w1": 0.03140699863433838,
+  "blocks.6.feed_forward.w2": 0.03598141297698021,
+  "blocks.6.feed_forward.w3": 0.03141026571393013,
+  "blocks.7.attention.w_q": 0.03154123201966286,
+  "blocks.7.attention.w_k": 0.031651951372623444,
+  "blocks.7.attention.w_v": 0.035259682685136795,
+  "blocks.7.attention.w_out": 0.03749556094408035,
+  "blocks.7.feed_forward.w1": 0.031442634761333466,
+  "blocks.7.feed_forward.w2": 0.03741942718625069,
+  "blocks.7.feed_forward.w3": 0.03139324113726616,
+  "blocks.8.attention.w_q": 0.030701495707035065,
+  "blocks.8.attention.w_k": 0.031224187463521957,
+  "blocks.8.attention.w_v": 0.03455720096826553,
+  "blocks.8.attention.w_out": 0.03540285304188728,
+  "blocks.8.feed_forward.w1": 0.031160781159996986,
+  "blocks.8.feed_forward.w2": 0.0417548343539238,
+  "blocks.8.feed_forward.w3": 0.03098127990961075
+ },
+ "rel_quant_noise_mean": 0.03286899629211615,
+ "n_quantized_params": 30670848,
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.7896115779876709,
+ "loss_eval_model": 4.133484409991807,
+ "checks": {
+  "snap_reproduces_fakequant": true,
+  "grid_levels_ok": true,
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/launch.yaml

@@ -0,0 +1,118 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 102
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: 6
+  fq_act_bits: 6
+  fq_qkv_bits: 6
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: olmo_core.optim.AdamWConfig
+    lr: 0.0024
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-15
+    weight_decay: 0.1
+    fused: true
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 232
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1002327
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 102
+n_tokens: 610009088
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.03452238067984581,
+  "blocks.0.attention.w_out.weight": 0.039284028112888336,
+  "blocks.0.attention.w_q.weight": 0.037009578198194504,
+  "blocks.0.attention.w_v.weight": 0.045134395360946655,
+  "blocks.0.attention_norm.weight": 0.739876389503479,
+  "blocks.0.feed_forward.w1.weight": 0.03798346593976021,
+  "blocks.0.feed_forward.w2.weight": 0.0359223298728466,
+  "blocks.0.feed_forward.w3.weight": 0.040152616798877716,
+  "blocks.0.feed_forward_norm.weight": 0.7368800044059753,
+  "blocks.1.attention.w_k.weight": 0.02928074263036251,
+  "blocks.1.attention.w_out.weight": 0.03693598881363869,
+  "blocks.1.attention.w_q.weight": 0.029001660645008087,
+  "blocks.1.attention.w_v.weight": 0.05250340327620506,
+  "blocks.1.attention_norm.weight": 0.7101855874061584,
+  "blocks.1.feed_forward.w1.weight": 0.047833625227212906,
+  "blocks.1.feed_forward.w2.weight": 0.04353787750005722,
+  "blocks.1.feed_forward.w3.weight": 0.048920538276433945,
+  "blocks.1.feed_forward_norm.weight": 0.7491304874420166,
+  "blocks.2.attention.w_k.weight": 0.03276506066322327,
+  "blocks.2.attention.w_out.weight": 0.03439931944012642,
+  "blocks.2.attention.w_q.weight": 0.03276437148451805,
+  "blocks.2.attention.w_v.weight": 0.0495532788336277,
+  "blocks.2.attention_norm.weight": 0.7206140756607056,
+  "blocks.2.feed_forward.w1.weight": 0.04991117864847183,
+  "blocks.2.feed_forward.w2.weight": 0.04640195146203041,
+  "blocks.2.feed_forward.w3.weight": 0.050042614340782166,
+  "blocks.2.feed_forward_norm.weight": 0.7519209980964661,
+  "blocks.3.attention.w_k.weight": 0.03559667617082596,
+  "blocks.3.attention.w_out.weight": 0.03937024250626564,
+  "blocks.3.attention.w_q.weight": 0.03479573875665665,
+  "blocks.3.attention.w_v.weight": 0.054586127400398254,
+  "blocks.3.attention_norm.weight": 0.7257862687110901,
+  "blocks.3.feed_forward.w1.weight": 0.048882536590099335,
+  "blocks.3.feed_forward.w2.weight": 0.04395680129528046,
+  "blocks.3.feed_forward.w3.weight": 0.04914790019392967,
+  "blocks.3.feed_forward_norm.weight": 0.754949152469635,
+  "blocks.4.attention.w_k.weight": 0.03553074598312378,
+  "blocks.4.attention.w_out.weight": 0.041324224323034286,
+  "blocks.4.attention.w_q.weight": 0.03667189180850983,
+  "blocks.4.attention.w_v.weight": 0.05485207587480545,
+  "blocks.4.attention_norm.weight": 0.7308045625686646,
+  "blocks.4.feed_forward.w1.weight": 0.04888234660029411,
+  "blocks.4.feed_forward.w2.weight": 0.04575841501355171,
+  "blocks.4.feed_forward.w3.weight": 0.049392424523830414,
+  "blocks.4.feed_forward_norm.weight": 0.752315878868103,
+  "blocks.5.attention.w_k.weight": 0.036597561091184616,
+  "blocks.5.attention.w_out.weight": 0.04749889299273491,
+  "blocks.5.attention.w_q.weight": 0.0354626402258873,
+  "blocks.5.attention.w_v.weight": 0.04642144963145256,
+  "blocks.5.attention_norm.weight": 0.7264080047607422,
+  "blocks.5.feed_forward.w1.weight": 0.04904254153370857,
+  "blocks.5.feed_forward.w2.weight": 0.04696989804506302,
+  "blocks.5.feed_forward.w3.weight": 0.05032774806022644,
+  "blocks.5.feed_forward_norm.weight": 0.7507573366165161,
+  "blocks.6.attention.w_k.weight": 0.03965156897902489,
+  "blocks.6.attention.w_out.weight": 0.05916905403137207,
+  "blocks.6.attention.w_q.weight": 0.03646329045295715,
+  "blocks.6.attention.w_v.weight": 0.04659885913133621,
+  "blocks.6.attention_norm.weight": 0.7252739667892456,
+  "blocks.6.feed_forward.w1.weight": 0.048663850873708725,
+  "blocks.6.feed_forward.w2.weight": 0.0510810911655426,
+  "blocks.6.feed_forward.w3.weight": 0.050724003463983536,
+  "blocks.6.feed_forward_norm.weight": 0.7489078640937805,
+  "blocks.7.attention.w_k.weight": 0.03603087738156319,
+  "blocks.7.attention.w_out.weight": 0.08853258937597275,
+  "blocks.7.attention.w_q.weight": 0.03710123151540756,
+  "blocks.7.attention.w_v.weight": 0.05319391191005707,
+  "blocks.7.attention_norm.weight": 0.7369636297225952,
+  "blocks.7.feed_forward.w1.weight": 0.049033112823963165,
+  "blocks.7.feed_forward.w2.weight": 0.06086074933409691,
+  "blocks.7.feed_forward.w3.weight": 0.05256909877061844,
+  "blocks.7.feed_forward_norm.weight": 0.7491270899772644,
+  "blocks.8.attention.w_k.weight": 0.04031990468502045,
+  "blocks.8.attention.w_out.weight": 0.09576456993818283,
+  "blocks.8.attention.w_q.weight": 0.03576497733592987,
+  "blocks.8.attention.w_v.weight": 0.042381834238767624,
+  "blocks.8.attention_norm.weight": 0.7316818237304688,
+  "blocks.8.feed_forward.w1.weight": 0.047424208372831345,
+  "blocks.8.feed_forward.w2.weight": 0.06759675592184067,
+  "blocks.8.feed_forward.w3.weight": 0.05208192765712738,
+  "blocks.8.feed_forward_norm.weight": 0.7457408905029297,
+  "embeddings.weight": 0.05595573037862778,
+  "lm_head.norm.weight": 1.3673397302627563,
+  "lm_head.w_out.weight": 0.07094171643257141
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 4.133484409991807
+ },
+ "base_loss": 4.133484409991807,
+ "loss": {
+  "0.005": [
+   4.134091336316463,
+   4.13444897731279,
+   4.133487833273491
+  ],
+  "0.01": [
+   4.13518646564468,
+   4.134274726406633,
+   4.133956944146065
+  ],
+  "0.02": [
+   4.135846947744149,
+   4.136405912914343,
+   4.1364804895948515
+  ],
+  "0.05": [
+   4.147796513000677,
+   4.148100367395197,
+   4.1488774366879895
+  ],
+  "0.1": [
+   4.191778945611088,
+   4.190513446400063,
+   4.193277110606912
+  ],
+  "0.15": [
+   4.268052713123267,
+   4.266817600916623,
+   4.270579134128668
+  ],
+  "0.2": [
+   4.384039275048625,
+   4.381098965594603,
+   4.388253307414357
+  ],
+  "0.3": [
+   4.788946718744993,
+   4.775071755113937,
+   4.790860837895414
+  ],
+  "0.4": [
+   5.53872043004493,
+   5.510531493315125,
+   5.542292638174446
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    4.134091336316463,
+    4.13444897731279,
+    4.133487833273491
+   ],
+   "0.01": [
+    4.13518646564468,
+    4.134274726406633,
+    4.133956944146065
+   ],
+   "0.02": [
+    4.135846947744149,
+    4.136405912914343,
+    4.1364804895948515
+   ],
+   "0.05": [
+    4.147796513000677,
+    4.148100367395197,
+    4.1488774366879895
+   ],
+   "0.1": [
+    4.191778945611088,
+    4.190513446400063,
+    4.193277110606912
+   ],
+   "0.15": [
+    4.268052713123267,
+    4.266817600916623,
+    4.270579134128668
+   ],
+   "0.2": [
+    4.384039275048625,
+    4.381098965594603,
+    4.388253307414357
+   ],
+   "0.3": [
+    4.788946718744993,
+    4.775071755113937,
+    4.790860837895414
+   ],
+   "0.4": [
+    5.53872043004493,
+    5.510531493315125,
+    5.542292638174446
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0005249723091074188,
+   "0.01": 0.000988302073985814,
+   "0.02": 0.0027600400926406365,
+   "0.05": 0.014773695702814393,
+   "0.1": 0.05837209088088091,
+   "0.15": 0.13499873939771145,
+   "0.2": 0.2509794393607212,
+   "0.3": 0.6514753605929746,
+   "0.4": 1.3970304438530272
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0005249723091074188,
+  "0.01": 0.000988302073985814,
+  "0.02": 0.0027600400926406365,
+  "0.05": 0.014773695702814393,
+  "0.1": 0.05837209088088091,
+  "0.15": 0.13499873939771145,
+  "0.2": 0.2509794393607212,
+  "0.3": 0.6514753605929746,
+  "0.4": 1.3970304438530272
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 77.9
+}

pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.03452238067984581,
+  "blocks.0.attention.w_out.weight": 0.039284028112888336,
+  "blocks.0.attention.w_q.weight": 0.037009578198194504,
+  "blocks.0.attention.w_v.weight": 0.045134395360946655,
+  "blocks.0.feed_forward.w1.weight": 0.03798346593976021,
+  "blocks.0.feed_forward.w2.weight": 0.0359223298728466,
+  "blocks.0.feed_forward.w3.weight": 0.040152616798877716,
+  "blocks.1.attention.w_k.weight": 0.02928074263036251,
+  "blocks.1.attention.w_out.weight": 0.03693598881363869,
+  "blocks.1.attention.w_q.weight": 0.029001660645008087,
+  "blocks.1.attention.w_v.weight": 0.05250340327620506,
+  "blocks.1.feed_forward.w1.weight": 0.047833625227212906,
+  "blocks.1.feed_forward.w2.weight": 0.04353787750005722,
+  "blocks.1.feed_forward.w3.weight": 0.048920538276433945,
+  "blocks.2.attention.w_k.weight": 0.03276506066322327,
+  "blocks.2.attention.w_out.weight": 0.03439931944012642,
+  "blocks.2.attention.w_q.weight": 0.03276437148451805,
+  "blocks.2.attention.w_v.weight": 0.0495532788336277,
+  "blocks.2.feed_forward.w1.weight": 0.04991117864847183,
+  "blocks.2.feed_forward.w2.weight": 0.04640195146203041,
+  "blocks.2.feed_forward.w3.weight": 0.050042614340782166,
+  "blocks.3.attention.w_k.weight": 0.03559667617082596,
+  "blocks.3.attention.w_out.weight": 0.03937024250626564,
+  "blocks.3.attention.w_q.weight": 0.03479573875665665,
+  "blocks.3.attention.w_v.weight": 0.054586127400398254,
+  "blocks.3.feed_forward.w1.weight": 0.048882536590099335,
+  "blocks.3.feed_forward.w2.weight": 0.04395680129528046,
+  "blocks.3.feed_forward.w3.weight": 0.04914790019392967,
+  "blocks.4.attention.w_k.weight": 0.03553074598312378,
+  "blocks.4.attention.w_out.weight": 0.041324224323034286,
+  "blocks.4.attention.w_q.weight": 0.03667189180850983,
+  "blocks.4.attention.w_v.weight": 0.05485207587480545,
+  "blocks.4.feed_forward.w1.weight": 0.04888234660029411,
+  "blocks.4.feed_forward.w2.weight": 0.04575841501355171,
+  "blocks.4.feed_forward.w3.weight": 0.049392424523830414,
+  "blocks.5.attention.w_k.weight": 0.036597561091184616,
+  "blocks.5.attention.w_out.weight": 0.04749889299273491,
+  "blocks.5.attention.w_q.weight": 0.0354626402258873,
+  "blocks.5.attention.w_v.weight": 0.04642144963145256,
+  "blocks.5.feed_forward.w1.weight": 0.04904254153370857,
+  "blocks.5.feed_forward.w2.weight": 0.04696989804506302,
+  "blocks.5.feed_forward.w3.weight": 0.05032774806022644,
+  "blocks.6.attention.w_k.weight": 0.03965156897902489,
+  "blocks.6.attention.w_out.weight": 0.05916905403137207,
+  "blocks.6.attention.w_q.weight": 0.03646329045295715,
+  "blocks.6.attention.w_v.weight": 0.04659885913133621,
+  "blocks.6.feed_forward.w1.weight": 0.048663850873708725,
+  "blocks.6.feed_forward.w2.weight": 0.0510810911655426,
+  "blocks.6.feed_forward.w3.weight": 0.050724003463983536,
+  "blocks.7.attention.w_k.weight": 0.03603087738156319,
+  "blocks.7.attention.w_out.weight": 0.08853258937597275,
+  "blocks.7.attention.w_q.weight": 0.03710123151540756,
+  "blocks.7.attention.w_v.weight": 0.05319391191005707,
+  "blocks.7.feed_forward.w1.weight": 0.049033112823963165,
+  "blocks.7.feed_forward.w2.weight": 0.06086074933409691,
+  "blocks.7.feed_forward.w3.weight": 0.05256909877061844,
+  "blocks.8.attention.w_k.weight": 0.04031990468502045,
+  "blocks.8.attention.w_out.weight": 0.09576456993818283,
+  "blocks.8.attention.w_q.weight": 0.03576497733592987,
+  "blocks.8.attention.w_v.weight": 0.042381834238767624,
+  "blocks.8.feed_forward.w1.weight": 0.047424208372831345,
+  "blocks.8.feed_forward.w2.weight": 0.06759675592184067,
+  "blocks.8.feed_forward.w3.weight": 0.05208192765712738
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 4.133484409991807
+ },
+ "base_loss": 4.133484409991807,
+ "loss": {
+  "0.005": [
+   4.133970637272798,
+   4.133880180477601,
+   4.133772411657893
+  ],
+  "0.01": [
+   4.133476006324315,
+   4.134386121948764,
+   4.133500517021167
+  ],
+  "0.02": [
+   4.134247136712427,
+   4.134187356433435,
+   4.134905647001647
+  ],
+  "0.05": [
+   4.1378708036122065,
+   4.138091039950253,
+   4.138317696408576
+  ],
+  "0.1": [
+   4.151370927679759,
+   4.150401763380517,
+   4.15023772363458
+  ],
+  "0.15": [
+   4.172548873977413,
+   4.172452908632948,
+   4.172826934433491
+  ],
+  "0.2": [
+   4.205900721983018,
+   4.205130561850742,
+   4.20530034097294
+  ],
+  "0.3": [
+   4.312236227343587,
+   4.312655600867546,
+   4.315095485224318
+  ],
+  "0.4": [
+   4.502774966904886,
+   4.495064891492847,
+   4.509433621449239
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    4.133970637272798,
+    4.133880180477601,
+    4.133772411657893
+   ],
+   "0.01": [
+    4.133476006324315,
+    4.134386121948764,
+    4.133500517021167
+   ],
+   "0.02": [
+    4.134247136712427,
+    4.134187356433435,
+    4.134905647001647
+   ],
+   "0.05": [
+    4.1378708036122065,
+    4.138091039950253,
+    4.138317696408576
+   ],
+   "0.1": [
+    4.151370927679759,
+    4.150401763380517,
+    4.15023772363458
+   ],
+   "0.15": [
+    4.172548873977413,
+    4.172452908632948,
+    4.172826934433491
+   ],
+   "0.2": [
+    4.205900721983018,
+    4.205130561850742,
+    4.20530034097294
+   ],
+   "0.3": [
+    4.312236227343587,
+    4.312655600867546,
+    4.315095485224318
+   ],
+   "0.4": [
+    4.502774966904886,
+    4.495064891492847,
+    4.509433621449239
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.00038999981095777514,
+   "0.01": 0.0003031384396088299,
+   "0.02": 0.0009623033906960288,
+   "0.05": 0.004608769998537987,
+   "0.1": 0.017185728239812015,
+   "0.15": 0.03912516235614394,
+   "0.2": 0.07195946494376049,
+   "0.3": 0.17984469448667717,
+   "0.4": 0.368940083290517
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.00038999981095777514,
+  "0.01": 0.0003031384396088299,
+  "0.02": 0.0009623033906960288,
+  "0.05": 0.004608769998537987,
+  "0.1": 0.017185728239812015,
+  "0.15": 0.03912516235614394,
+  "0.2": 0.07195946494376049,
+  "0.3": 0.17984469448667717,
+  "0.4": 0.368940083290517
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 72.2
+}

pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102/ptq.json

@@ -0,0 +1,52 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D0.61B-f6-adamw-lr2.4e-03-s102",
+ "w_bits": 6,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 4.133484409991807
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 4.402245652440087
+  },
+  "4": {
+   "val_dolma_cc": 4.17207222014817
+  },
+  "5": {
+   "val_dolma_cc": 4.144966006773763
+  },
+  "6": {
+   "val_dolma_cc": 4.133246677750118
+  },
+  "8": {
+   "val_dolma_cc": 4.134107232080439
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 0.2687612424482797
+  },
+  "4": {
+   "val_dolma_cc": 0.038587810156363034
+  },
+  "5": {
+   "val_dolma_cc": 0.01148159678195615
+  },
+  "6": {
+   "val_dolma_cc": -0.00023773224168888163
+  },
+  "8": {
+   "val_dolma_cc": 0.0006228220886317359
+  }
+ },
+ "own_grid_max_rel_weight_change": 2.2988398029610835e-07,
+ "own_grid_delta_loss": -0.00023773224168888163,
+ "checks": {
+  "weights_restored": true,
+  "own_grid_weights_unchanged": true,
+  "own_grid_is_fixed_point": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": null,
+  "a": null,
+  "qkv": null
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 3.831054171073506,
+  "heldout_cc_en_tail_4M": 4.413598668423799
+ },
+ "loss_bf16_master": 3.831054171073506,
+ "losses_eval_model": {
+  "val_dolma_cc": 3.831054171073506
+ },
+ "bf16_minus_quant": 0.0,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-m5-16",
+ "gpu": "NVIDIA L40S",
+ "t_s": 17.8
+}

pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/finalize.json

@@ -0,0 +1,29 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101",
+ "w_bits": null,
+ "a_bits": null,
+ "qkv_bits": null,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.831054171073506
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.831054171073506
+ },
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 1.0510008335113525,
+ "loss_eval_model": 3.831054171073506,
+ "checks": {
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/launch.yaml

@@ -0,0 +1,118 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 101
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: null
+  fq_act_bits: null
+  fq_qkv_bits: null
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: olmo_core.optim.AdamWConfig
+    lr: 0.0012
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-15
+    weight_decay: 0.1
+    fused: true
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 701
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1007019
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 101
+n_tokens: 1839988736
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.03268071636557579,
+  "blocks.0.attention.w_out.weight": 0.03668426722288132,
+  "blocks.0.attention.w_q.weight": 0.03406849503517151,
+  "blocks.0.attention.w_v.weight": 0.033527445048093796,
+  "blocks.0.attention_norm.weight": 0.6293824911117554,
+  "blocks.0.feed_forward.w1.weight": 0.038742586970329285,
+  "blocks.0.feed_forward.w2.weight": 0.039205584675073624,
+  "blocks.0.feed_forward.w3.weight": 0.04028899595141411,
+  "blocks.0.feed_forward_norm.weight": 0.6330354809761047,
+  "blocks.1.attention.w_k.weight": 0.03544138744473457,
+  "blocks.1.attention.w_out.weight": 0.03620398789644241,
+  "blocks.1.attention.w_q.weight": 0.0365966372191906,
+  "blocks.1.attention.w_v.weight": 0.03734150528907776,
+  "blocks.1.attention_norm.weight": 0.6099925637245178,
+  "blocks.1.feed_forward.w1.weight": 0.04364994540810585,
+  "blocks.1.feed_forward.w2.weight": 0.04219818487763405,
+  "blocks.1.feed_forward.w3.weight": 0.04411840811371803,
+  "blocks.1.feed_forward_norm.weight": 0.643448531627655,
+  "blocks.2.attention.w_k.weight": 0.03736107051372528,
+  "blocks.2.attention.w_out.weight": 0.03443260118365288,
+  "blocks.2.attention.w_q.weight": 0.03904962167143822,
+  "blocks.2.attention.w_v.weight": 0.03684396296739578,
+  "blocks.2.attention_norm.weight": 0.6134346723556519,
+  "blocks.2.feed_forward.w1.weight": 0.04324464499950409,
+  "blocks.2.feed_forward.w2.weight": 0.040622010827064514,
+  "blocks.2.feed_forward.w3.weight": 0.04401577636599541,
+  "blocks.2.feed_forward_norm.weight": 0.6434590220451355,
+  "blocks.3.attention.w_k.weight": 0.03984769433736801,
+  "blocks.3.attention.w_out.weight": 0.03913891315460205,
+  "blocks.3.attention.w_q.weight": 0.042075786739587784,
+  "blocks.3.attention.w_v.weight": 0.039377178996801376,
+  "blocks.3.attention_norm.weight": 0.6263216137886047,
+  "blocks.3.feed_forward.w1.weight": 0.04268060624599457,
+  "blocks.3.feed_forward.w2.weight": 0.0402805358171463,
+  "blocks.3.feed_forward.w3.weight": 0.04456881061196327,
+  "blocks.3.feed_forward_norm.weight": 0.6389124989509583,
+  "blocks.4.attention.w_k.weight": 0.04043733328580856,
+  "blocks.4.attention.w_out.weight": 0.04046373814344406,
+  "blocks.4.attention.w_q.weight": 0.04414784163236618,
+  "blocks.4.attention.w_v.weight": 0.0379972942173481,
+  "blocks.4.attention_norm.weight": 0.6256927251815796,
+  "blocks.4.feed_forward.w1.weight": 0.04161463677883148,
+  "blocks.4.feed_forward.w2.weight": 0.040296878665685654,
+  "blocks.4.feed_forward.w3.weight": 0.04442773759365082,
+  "blocks.4.feed_forward_norm.weight": 0.6358168721199036,
+  "blocks.5.attention.w_k.weight": 0.03878093510866165,
+  "blocks.5.attention.w_out.weight": 0.04189108684659004,
+  "blocks.5.attention.w_q.weight": 0.041319891810417175,
+  "blocks.5.attention.w_v.weight": 0.0382239855825901,
+  "blocks.5.attention_norm.weight": 0.6185463070869446,
+  "blocks.5.feed_forward.w1.weight": 0.041318804025650024,
+  "blocks.5.feed_forward.w2.weight": 0.04324207454919815,
+  "blocks.5.feed_forward.w3.weight": 0.04525141790509224,
+  "blocks.5.feed_forward_norm.weight": 0.6307767033576965,
+  "blocks.6.attention.w_k.weight": 0.04055728018283844,
+  "blocks.6.attention.w_out.weight": 0.055124517530202866,
+  "blocks.6.attention.w_q.weight": 0.04400692880153656,
+  "blocks.6.attention.w_v.weight": 0.04032757133245468,
+  "blocks.6.attention_norm.weight": 0.6209803819656372,
+  "blocks.6.feed_forward.w1.weight": 0.04073719307780266,
+  "blocks.6.feed_forward.w2.weight": 0.04550210386514664,
+  "blocks.6.feed_forward.w3.weight": 0.04522354528307915,
+  "blocks.6.feed_forward_norm.weight": 0.6327573657035828,
+  "blocks.7.attention.w_k.weight": 0.039778005331754684,
+  "blocks.7.attention.w_out.weight": 0.057591959834098816,
+  "blocks.7.attention.w_q.weight": 0.04154594987630844,
+  "blocks.7.attention.w_v.weight": 0.03982347622513771,
+  "blocks.7.attention_norm.weight": 0.6144234538078308,
+  "blocks.7.feed_forward.w1.weight": 0.040671318769454956,
+  "blocks.7.feed_forward.w2.weight": 0.0517294704914093,
+  "blocks.7.feed_forward.w3.weight": 0.04585808888077736,
+  "blocks.7.feed_forward_norm.weight": 0.6315277814865112,
+  "blocks.8.attention.w_k.weight": 0.03729722276329994,
+  "blocks.8.attention.w_out.weight": 0.0786689817905426,
+  "blocks.8.attention.w_q.weight": 0.04076381400227547,
+  "blocks.8.attention.w_v.weight": 0.03802540898323059,
+  "blocks.8.attention_norm.weight": 0.6195935606956482,
+  "blocks.8.feed_forward.w1.weight": 0.038343265652656555,
+  "blocks.8.feed_forward.w2.weight": 0.05855748802423477,
+  "blocks.8.feed_forward.w3.weight": 0.044858478009700775,
+  "blocks.8.feed_forward_norm.weight": 0.6258807182312012,
+  "embeddings.weight": 0.045032817870378494,
+  "lm_head.norm.weight": 1.8185430765151978,
+  "lm_head.w_out.weight": 0.05665469169616699
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.831054171073506
+ },
+ "base_loss": 3.831054171073506,
+ "loss": {
+  "0.005": [
+   3.8317018403186744,
+   3.8312902014551584,
+   3.8314713380479004
+  ],
+  "0.01": [
+   3.8330465366624002,
+   3.8324081983050196,
+   3.8327434412583328
+  ],
+  "0.02": [
+   3.838123752438785,
+   3.836907145514839,
+   3.8374857075055484
+  ],
+  "0.05": [
+   3.872347300979871,
+   3.8698361849778355,
+   3.871757166340911
+  ],
+  "0.1": [
+   4.003700027599084,
+   4.0011724876696135,
+   4.013152979830948
+  ],
+  "0.15": [
+   4.276705020410427,
+   4.271452419154528,
+   4.3165530950910815
+  ],
+  "0.2": [
+   4.786053736649279,
+   4.758407978870707,
+   4.870210199041412
+  ],
+  "0.3": [
+   6.652791131836611,
+   6.553454359603605,
+   6.76614946240315
+  ],
+  "0.4": [
+   8.568021647722265,
+   8.522721748528717,
+   8.586911589912544
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.8317018403186744,
+    3.8312902014551584,
+    3.8314713380479004
+   ],
+   "0.01": [
+    3.8330465366624002,
+    3.8324081983050196,
+    3.8327434412583328
+   ],
+   "0.02": [
+    3.838123752438785,
+    3.836907145514839,
+    3.8374857075055484
+   ],
+   "0.05": [
+    3.872347300979871,
+    3.8698361849778355,
+    3.871757166340911
+   ],
+   "0.1": [
+    4.003700027599084,
+    4.0011724876696135,
+    4.013152979830948
+   ],
+   "0.15": [
+    4.276705020410427,
+    4.271452419154528,
+    4.3165530950910815
+   ],
+   "0.2": [
+    4.786053736649279,
+    4.758407978870707,
+    4.870210199041412
+   ],
+   "0.3": [
+    6.652791131836611,
+    6.553454359603605,
+    6.76614946240315
+   ],
+   "0.4": [
+    8.568021647722265,
+    8.522721748528717,
+    8.586911589912544
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0004336222004055479,
+   "0.01": 0.0016785543350783527,
+   "0.02": 0.0064513640795516025,
+   "0.05": 0.040259379692700215,
+   "0.1": 0.1749543272930425,
+   "0.15": 0.4571826738118401,
+   "0.2": 0.9738364671136268,
+   "0.3": 2.82641081354095,
+   "0.4": 4.728164157647669
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0004336222004055479,
+  "0.01": 0.0016785543350783527,
+  "0.02": 0.0064513640795516025,
+  "0.05": 0.040259379692700215,
+  "0.1": 0.1749543272930425,
+  "0.15": 0.4571826738118401,
+  "0.2": 0.9738364671136268,
+  "0.3": 2.82641081354095,
+  "0.4": 4.728164157647669
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 56.0
+}

pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.03268071636557579,
+  "blocks.0.attention.w_out.weight": 0.03668426722288132,
+  "blocks.0.attention.w_q.weight": 0.03406849503517151,
+  "blocks.0.attention.w_v.weight": 0.033527445048093796,
+  "blocks.0.feed_forward.w1.weight": 0.038742586970329285,
+  "blocks.0.feed_forward.w2.weight": 0.039205584675073624,
+  "blocks.0.feed_forward.w3.weight": 0.04028899595141411,
+  "blocks.1.attention.w_k.weight": 0.03544138744473457,
+  "blocks.1.attention.w_out.weight": 0.03620398789644241,
+  "blocks.1.attention.w_q.weight": 0.0365966372191906,
+  "blocks.1.attention.w_v.weight": 0.03734150528907776,
+  "blocks.1.feed_forward.w1.weight": 0.04364994540810585,
+  "blocks.1.feed_forward.w2.weight": 0.04219818487763405,
+  "blocks.1.feed_forward.w3.weight": 0.04411840811371803,
+  "blocks.2.attention.w_k.weight": 0.03736107051372528,
+  "blocks.2.attention.w_out.weight": 0.03443260118365288,
+  "blocks.2.attention.w_q.weight": 0.03904962167143822,
+  "blocks.2.attention.w_v.weight": 0.03684396296739578,
+  "blocks.2.feed_forward.w1.weight": 0.04324464499950409,
+  "blocks.2.feed_forward.w2.weight": 0.040622010827064514,
+  "blocks.2.feed_forward.w3.weight": 0.04401577636599541,
+  "blocks.3.attention.w_k.weight": 0.03984769433736801,
+  "blocks.3.attention.w_out.weight": 0.03913891315460205,
+  "blocks.3.attention.w_q.weight": 0.042075786739587784,
+  "blocks.3.attention.w_v.weight": 0.039377178996801376,
+  "blocks.3.feed_forward.w1.weight": 0.04268060624599457,
+  "blocks.3.feed_forward.w2.weight": 0.0402805358171463,
+  "blocks.3.feed_forward.w3.weight": 0.04456881061196327,
+  "blocks.4.attention.w_k.weight": 0.04043733328580856,
+  "blocks.4.attention.w_out.weight": 0.04046373814344406,
+  "blocks.4.attention.w_q.weight": 0.04414784163236618,
+  "blocks.4.attention.w_v.weight": 0.0379972942173481,
+  "blocks.4.feed_forward.w1.weight": 0.04161463677883148,
+  "blocks.4.feed_forward.w2.weight": 0.040296878665685654,
+  "blocks.4.feed_forward.w3.weight": 0.04442773759365082,
+  "blocks.5.attention.w_k.weight": 0.03878093510866165,
+  "blocks.5.attention.w_out.weight": 0.04189108684659004,
+  "blocks.5.attention.w_q.weight": 0.041319891810417175,
+  "blocks.5.attention.w_v.weight": 0.0382239855825901,
+  "blocks.5.feed_forward.w1.weight": 0.041318804025650024,
+  "blocks.5.feed_forward.w2.weight": 0.04324207454919815,
+  "blocks.5.feed_forward.w3.weight": 0.04525141790509224,
+  "blocks.6.attention.w_k.weight": 0.04055728018283844,
+  "blocks.6.attention.w_out.weight": 0.055124517530202866,
+  "blocks.6.attention.w_q.weight": 0.04400692880153656,
+  "blocks.6.attention.w_v.weight": 0.04032757133245468,
+  "blocks.6.feed_forward.w1.weight": 0.04073719307780266,
+  "blocks.6.feed_forward.w2.weight": 0.04550210386514664,
+  "blocks.6.feed_forward.w3.weight": 0.04522354528307915,
+  "blocks.7.attention.w_k.weight": 0.039778005331754684,
+  "blocks.7.attention.w_out.weight": 0.057591959834098816,
+  "blocks.7.attention.w_q.weight": 0.04154594987630844,
+  "blocks.7.attention.w_v.weight": 0.03982347622513771,
+  "blocks.7.feed_forward.w1.weight": 0.040671318769454956,
+  "blocks.7.feed_forward.w2.weight": 0.0517294704914093,
+  "blocks.7.feed_forward.w3.weight": 0.04585808888077736,
+  "blocks.8.attention.w_k.weight": 0.03729722276329994,
+  "blocks.8.attention.w_out.weight": 0.0786689817905426,
+  "blocks.8.attention.w_q.weight": 0.04076381400227547,
+  "blocks.8.attention.w_v.weight": 0.03802540898323059,
+  "blocks.8.feed_forward.w1.weight": 0.038343265652656555,
+  "blocks.8.feed_forward.w2.weight": 0.05855748802423477,
+  "blocks.8.feed_forward.w3.weight": 0.044858478009700775
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.831054171073506
+ },
+ "base_loss": 3.831054171073506,
+ "loss": {
+  "0.005": [
+   3.83126655146913,
+   3.831416064763763,
+   3.8311541387232984
+  ],
+  "0.01": [
+   3.8320505338347615,
+   3.8322378597207867,
+   3.8318906489747087
+  ],
+  "0.02": [
+   3.835079523889827,
+   3.8354311594720008,
+   3.8344159588806828
+  ],
+  "0.05": [
+   3.8555864952406136,
+   3.857775685833163,
+   3.8541448554719424
+  ],
+  "0.1": [
+   3.9345608681548763,
+   3.946372234551774,
+   3.9354457101796987
+  ],
+  "0.15": [
+   4.090792315924315,
+   4.127439078418456,
+   4.105478190432802
+  ],
+  "0.2": [
+   4.360869604782873,
+   4.4485717888851,
+   4.4065429116301
+  ],
+  "0.3": [
+   5.350963750428164,
+   5.582059019661014,
+   5.417273765133487
+  ],
+  "0.4": [
+   6.654905765310262,
+   6.920882303084387,
+   6.549217831386812
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.83126655146913,
+    3.831416064763763,
+    3.8311541387232984
+   ],
+   "0.01": [
+    3.8320505338347615,
+    3.8322378597207867,
+    3.8318906489747087
+   ],
+   "0.02": [
+    3.835079523889827,
+    3.8354311594720008,
+    3.8344159588806828
+   ],
+   "0.05": [
+    3.8555864952406136,
+    3.857775685833163,
+    3.8541448554719424
+   ],
+   "0.1": [
+    3.9345608681548763,
+    3.946372234551774,
+    3.9354457101796987
+   ],
+   "0.15": [
+    4.090792315924315,
+    4.127439078418456,
+    4.105478190432802
+   ],
+   "0.2": [
+    4.360869604782873,
+    4.4485717888851,
+    4.4065429116301
+   ],
+   "0.3": [
+    5.350963750428164,
+    5.582059019661014,
+    5.417273765133487
+   ],
+   "0.4": [
+    6.654905765310262,
+    6.920882303084387,
+    6.549217831386812
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.000224747245224588,
+   "0.01": 0.0010055097699126492,
+   "0.02": 0.003921376340664473,
+   "0.05": 0.024781507775066824,
+   "0.1": 0.10773876655527692,
+   "0.15": 0.2768490238516841,
+   "0.2": 0.574273930692518,
+   "0.3": 1.6190446740007158,
+   "0.4": 2.8772811288536473
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.000224747245224588,
+  "0.01": 0.0010055097699126492,
+  "0.02": 0.003921376340664473,
+  "0.05": 0.024781507775066824,
+  "0.1": 0.10773876655527692,
+  "0.15": 0.2768490238516841,
+  "0.2": 0.574273930692518,
+  "0.3": 1.6190446740007158,
+  "0.4": 2.8772811288536473
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 49.8
+}

pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101/ptq.json

@@ -0,0 +1,48 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-adamw-lr1.2e-03-s101",
+ "w_bits": null,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.831054171073506
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 6.164653690965879
+  },
+  "4": {
+   "val_dolma_cc": 4.078234822788718
+  },
+  "5": {
+   "val_dolma_cc": 3.8779693310580523
+  },
+  "6": {
+   "val_dolma_cc": 3.839041501003307
+  },
+  "8": {
+   "val_dolma_cc": 3.8320119113927316
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 2.3335995198923727
+  },
+  "4": {
+   "val_dolma_cc": 0.2471806517152122
+  },
+  "5": {
+   "val_dolma_cc": 0.04691515998454632
+  },
+  "6": {
+   "val_dolma_cc": 0.007987329929801223
+  },
+  "8": {
+   "val_dolma_cc": 0.0009577403192255751
+  }
+ },
+ "checks": {
+  "weights_restored": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": null,
+  "a": null,
+  "qkv": null
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 3.741282098723469,
+  "heldout_cc_en_tail_4M": 4.324322303835364
+ },
+ "loss_bf16_master": 3.741282098723469,
+ "losses_eval_model": {
+  "val_dolma_cc": 3.741282098723469
+ },
+ "bf16_minus_quant": 0.0,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-t5-28",
+ "gpu": "NVIDIA L40S",
+ "t_s": 17.7
+}

pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/finalize.json

@@ -0,0 +1,29 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "w_bits": null,
+ "a_bits": null,
+ "qkv_bits": null,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.741282098723469
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.741282098723469
+ },
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.7421518564224243,
+ "loss_eval_model": 3.741282098723469,
+ "checks": {
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/launch.yaml

@@ -0,0 +1,126 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 101
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: null
+  fq_act_bits: null
+  fq_qkv_bits: null
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: precq.muon26.MuonConfig
+    lr: 0.01
+    weight_decay: 0.1
+    momentum: 0.95
+    nesterov: true
+    ns_steps: 5
+    eps: 1.0e-07
+    adjust_lr_fn: null
+    adamw_lr: 0.0024
+    adamw_betas:
+    - 0.9
+    - 0.95
+    adamw_eps: 1.0e-15
+    adamw_weight_decay: 0.1
+    adamw_fused: true
+    muon_on_lm_head: false
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 701
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1007019
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 101
+n_tokens: 1839988736
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.035397958010435104,
+  "blocks.0.attention.w_out.weight": 0.02164395898580551,
+  "blocks.0.attention.w_q.weight": 0.035328783094882965,
+  "blocks.0.attention.w_v.weight": 0.025519702583551407,
+  "blocks.0.attention_norm.weight": 0.38304951786994934,
+  "blocks.0.feed_forward.w1.weight": 0.030326496809720993,
+  "blocks.0.feed_forward.w2.weight": 0.01632818393409252,
+  "blocks.0.feed_forward.w3.weight": 0.030313409864902496,
+  "blocks.0.feed_forward_norm.weight": 0.39804232120513916,
+  "blocks.1.attention.w_k.weight": 0.027827803045511246,
+  "blocks.1.attention.w_out.weight": 0.01929481327533722,
+  "blocks.1.attention.w_q.weight": 0.028433021157979965,
+  "blocks.1.attention.w_v.weight": 0.021994709968566895,
+  "blocks.1.attention_norm.weight": 0.3855082094669342,
+  "blocks.1.feed_forward.w1.weight": 0.030362948775291443,
+  "blocks.1.feed_forward.w2.weight": 0.01613524742424488,
+  "blocks.1.feed_forward.w3.weight": 0.029716597869992256,
+  "blocks.1.feed_forward_norm.weight": 0.40586504340171814,
+  "blocks.2.attention.w_k.weight": 0.03229168802499771,
+  "blocks.2.attention.w_out.weight": 0.0226726196706295,
+  "blocks.2.attention.w_q.weight": 0.03159981593489647,
+  "blocks.2.attention.w_v.weight": 0.0259097870439291,
+  "blocks.2.attention_norm.weight": 0.3970188796520233,
+  "blocks.2.feed_forward.w1.weight": 0.029754923656582832,
+  "blocks.2.feed_forward.w2.weight": 0.01480945385992527,
+  "blocks.2.feed_forward.w3.weight": 0.02957199513912201,
+  "blocks.2.feed_forward_norm.weight": 0.40933650732040405,
+  "blocks.3.attention.w_k.weight": 0.032474976032972336,
+  "blocks.3.attention.w_out.weight": 0.02344447560608387,
+  "blocks.3.attention.w_q.weight": 0.03168802335858345,
+  "blocks.3.attention.w_v.weight": 0.026256276294589043,
+  "blocks.3.attention_norm.weight": 0.3984137773513794,
+  "blocks.3.feed_forward.w1.weight": 0.02918386459350586,
+  "blocks.3.feed_forward.w2.weight": 0.014518745243549347,
+  "blocks.3.feed_forward.w3.weight": 0.029343916103243828,
+  "blocks.3.feed_forward_norm.weight": 0.40580126643180847,
+  "blocks.4.attention.w_k.weight": 0.03286627680063248,
+  "blocks.4.attention.w_out.weight": 0.024900637567043304,
+  "blocks.4.attention.w_q.weight": 0.03240432217717171,
+  "blocks.4.attention.w_v.weight": 0.027380263432860374,
+  "blocks.4.attention_norm.weight": 0.3966514468193054,
+  "blocks.4.feed_forward.w1.weight": 0.02874669060111046,
+  "blocks.4.feed_forward.w2.weight": 0.01469415333122015,
+  "blocks.4.feed_forward.w3.weight": 0.029311303049325943,
+  "blocks.4.feed_forward_norm.weight": 0.40535956621170044,
+  "blocks.5.attention.w_k.weight": 0.03163629397749901,
+  "blocks.5.attention.w_out.weight": 0.026556408032774925,
+  "blocks.5.attention.w_q.weight": 0.0323525033891201,
+  "blocks.5.attention.w_v.weight": 0.027816232293844223,
+  "blocks.5.attention_norm.weight": 0.4023590683937073,
+  "blocks.5.feed_forward.w1.weight": 0.028784623369574547,
+  "blocks.5.feed_forward.w2.weight": 0.015200715512037277,
+  "blocks.5.feed_forward.w3.weight": 0.029435208067297935,
+  "blocks.5.feed_forward_norm.weight": 0.40602362155914307,
+  "blocks.6.attention.w_k.weight": 0.03339715674519539,
+  "blocks.6.attention.w_out.weight": 0.029964886605739594,
+  "blocks.6.attention.w_q.weight": 0.0327397957444191,
+  "blocks.6.attention.w_v.weight": 0.029403777793049812,
+  "blocks.6.attention_norm.weight": 0.40344297885894775,
+  "blocks.6.feed_forward.w1.weight": 0.02904530242085457,
+  "blocks.6.feed_forward.w2.weight": 0.017573032528162003,
+  "blocks.6.feed_forward.w3.weight": 0.02994123101234436,
+  "blocks.6.feed_forward_norm.weight": 0.408926784992218,
+  "blocks.7.attention.w_k.weight": 0.032194577157497406,
+  "blocks.7.attention.w_out.weight": 0.03255749121308327,
+  "blocks.7.attention.w_q.weight": 0.03139836713671684,
+  "blocks.7.attention.w_v.weight": 0.029072798788547516,
+  "blocks.7.attention_norm.weight": 0.40457212924957275,
+  "blocks.7.feed_forward.w1.weight": 0.028424231335520744,
+  "blocks.7.feed_forward.w2.weight": 0.02051093615591526,
+  "blocks.7.feed_forward.w3.weight": 0.029400411993265152,
+  "blocks.7.feed_forward_norm.weight": 0.4079188406467438,
+  "blocks.8.attention.w_k.weight": 0.030158329755067825,
+  "blocks.8.attention.w_out.weight": 0.038990020751953125,
+  "blocks.8.attention.w_q.weight": 0.029624801129102707,
+  "blocks.8.attention.w_v.weight": 0.02889910154044628,
+  "blocks.8.attention_norm.weight": 0.4107125699520111,
+  "blocks.8.feed_forward.w1.weight": 0.026936113834381104,
+  "blocks.8.feed_forward.w2.weight": 0.025364842265844345,
+  "blocks.8.feed_forward.w3.weight": 0.028375402092933655,
+  "blocks.8.feed_forward_norm.weight": 0.4011843800544739,
+  "embeddings.weight": 0.06977110356092453,
+  "lm_head.norm.weight": 1.3798266649246216,
+  "lm_head.w_out.weight": 0.07359108328819275
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.741282098723469
+ },
+ "base_loss": 3.741282098723469,
+ "loss": {
+  "0.005": [
+   3.7416292781709712,
+   3.741794487701162,
+   3.742089562845286
+  ],
+  "0.01": [
+   3.74275158248727,
+   3.7429588612151057,
+   3.7436641863484628
+  ],
+  "0.02": [
+   3.7473024344720507,
+   3.747717246228647,
+   3.749264851236619
+  ],
+  "0.05": [
+   3.780543043428401,
+   3.781887751509092,
+   3.785293220390546
+  ],
+  "0.1": [
+   3.916245888619137,
+   3.919793292627658,
+   3.924784930274286
+  ],
+  "0.15": [
+   4.211604527414047,
+   4.2234978373936345,
+   4.22084451631432
+  ],
+  "0.2": [
+   4.7966840943394145,
+   4.855740092999335,
+   4.809550477027847
+  ],
+  "0.3": [
+   7.265636298785715,
+   7.430777746750457,
+   7.120484006803559
+  ],
+  "0.4": [
+   9.54127430772559,
+   9.48668455317653,
+   9.123880935086254
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.7416292781709712,
+    3.741794487701162,
+    3.742089562845286
+   ],
+   "0.01": [
+    3.74275158248727,
+    3.7429588612151057,
+    3.7436641863484628
+   ],
+   "0.02": [
+    3.7473024344720507,
+    3.747717246228647,
+    3.749264851236619
+   ],
+   "0.05": [
+    3.780543043428401,
+    3.781887751509092,
+    3.785293220390546
+   ],
+   "0.1": [
+    3.916245888619137,
+    3.919793292627658,
+    3.924784930274286
+   ],
+   "0.15": [
+    4.211604527414047,
+    4.2234978373936345,
+    4.22084451631432
+   ],
+   "0.2": [
+    4.7966840943394145,
+    4.855740092999335,
+    4.809550477027847
+   ],
+   "0.3": [
+    7.265636298785715,
+    7.430777746750457,
+    7.120484006803559
+   ],
+   "0.4": [
+    9.54127430772559,
+    9.48668455317653,
+    9.123880935086254
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0005556775156705385,
+   "0.01": 0.00184277796014376,
+   "0.02": 0.0068127452556367984,
+   "0.05": 0.04129257305254397,
+   "0.1": 0.17899260511689175,
+   "0.15": 0.47736686165053177,
+   "0.2": 1.079376122732063,
+   "0.3": 3.5310172520564413,
+   "0.4": 5.642664499939322
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0005556775156705385,
+  "0.01": 0.00184277796014376,
+  "0.02": 0.0068127452556367984,
+  "0.05": 0.04129257305254397,
+  "0.1": 0.17899260511689175,
+  "0.15": 0.47736686165053177,
+  "0.2": 1.079376122732063,
+  "0.3": 3.5310172520564413,
+  "0.4": 5.642664499939322
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 56.5
+}

pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.035397958010435104,
+  "blocks.0.attention.w_out.weight": 0.02164395898580551,
+  "blocks.0.attention.w_q.weight": 0.035328783094882965,
+  "blocks.0.attention.w_v.weight": 0.025519702583551407,
+  "blocks.0.feed_forward.w1.weight": 0.030326496809720993,
+  "blocks.0.feed_forward.w2.weight": 0.01632818393409252,
+  "blocks.0.feed_forward.w3.weight": 0.030313409864902496,
+  "blocks.1.attention.w_k.weight": 0.027827803045511246,
+  "blocks.1.attention.w_out.weight": 0.01929481327533722,
+  "blocks.1.attention.w_q.weight": 0.028433021157979965,
+  "blocks.1.attention.w_v.weight": 0.021994709968566895,
+  "blocks.1.feed_forward.w1.weight": 0.030362948775291443,
+  "blocks.1.feed_forward.w2.weight": 0.01613524742424488,
+  "blocks.1.feed_forward.w3.weight": 0.029716597869992256,
+  "blocks.2.attention.w_k.weight": 0.03229168802499771,
+  "blocks.2.attention.w_out.weight": 0.0226726196706295,
+  "blocks.2.attention.w_q.weight": 0.03159981593489647,
+  "blocks.2.attention.w_v.weight": 0.0259097870439291,
+  "blocks.2.feed_forward.w1.weight": 0.029754923656582832,
+  "blocks.2.feed_forward.w2.weight": 0.01480945385992527,
+  "blocks.2.feed_forward.w3.weight": 0.02957199513912201,
+  "blocks.3.attention.w_k.weight": 0.032474976032972336,
+  "blocks.3.attention.w_out.weight": 0.02344447560608387,
+  "blocks.3.attention.w_q.weight": 0.03168802335858345,
+  "blocks.3.attention.w_v.weight": 0.026256276294589043,
+  "blocks.3.feed_forward.w1.weight": 0.02918386459350586,
+  "blocks.3.feed_forward.w2.weight": 0.014518745243549347,
+  "blocks.3.feed_forward.w3.weight": 0.029343916103243828,
+  "blocks.4.attention.w_k.weight": 0.03286627680063248,
+  "blocks.4.attention.w_out.weight": 0.024900637567043304,
+  "blocks.4.attention.w_q.weight": 0.03240432217717171,
+  "blocks.4.attention.w_v.weight": 0.027380263432860374,
+  "blocks.4.feed_forward.w1.weight": 0.02874669060111046,
+  "blocks.4.feed_forward.w2.weight": 0.01469415333122015,
+  "blocks.4.feed_forward.w3.weight": 0.029311303049325943,
+  "blocks.5.attention.w_k.weight": 0.03163629397749901,
+  "blocks.5.attention.w_out.weight": 0.026556408032774925,
+  "blocks.5.attention.w_q.weight": 0.0323525033891201,
+  "blocks.5.attention.w_v.weight": 0.027816232293844223,
+  "blocks.5.feed_forward.w1.weight": 0.028784623369574547,
+  "blocks.5.feed_forward.w2.weight": 0.015200715512037277,
+  "blocks.5.feed_forward.w3.weight": 0.029435208067297935,
+  "blocks.6.attention.w_k.weight": 0.03339715674519539,
+  "blocks.6.attention.w_out.weight": 0.029964886605739594,
+  "blocks.6.attention.w_q.weight": 0.0327397957444191,
+  "blocks.6.attention.w_v.weight": 0.029403777793049812,
+  "blocks.6.feed_forward.w1.weight": 0.02904530242085457,
+  "blocks.6.feed_forward.w2.weight": 0.017573032528162003,
+  "blocks.6.feed_forward.w3.weight": 0.02994123101234436,
+  "blocks.7.attention.w_k.weight": 0.032194577157497406,
+  "blocks.7.attention.w_out.weight": 0.03255749121308327,
+  "blocks.7.attention.w_q.weight": 0.03139836713671684,
+  "blocks.7.attention.w_v.weight": 0.029072798788547516,
+  "blocks.7.feed_forward.w1.weight": 0.028424231335520744,
+  "blocks.7.feed_forward.w2.weight": 0.02051093615591526,
+  "blocks.7.feed_forward.w3.weight": 0.029400411993265152,
+  "blocks.8.attention.w_k.weight": 0.030158329755067825,
+  "blocks.8.attention.w_out.weight": 0.038990020751953125,
+  "blocks.8.attention.w_q.weight": 0.029624801129102707,
+  "blocks.8.attention.w_v.weight": 0.02889910154044628,
+  "blocks.8.feed_forward.w1.weight": 0.026936113834381104,
+  "blocks.8.feed_forward.w2.weight": 0.025364842265844345,
+  "blocks.8.feed_forward.w3.weight": 0.028375402092933655
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.741282098723469
+ },
+ "base_loss": 3.741282098723469,
+ "loss": {
+  "0.005": [
+   3.741544666384729,
+   3.741646793636223,
+   3.7415041934144293
+  ],
+  "0.01": [
+   3.742238622310576,
+   3.7425577856207837,
+   3.7421255522946613
+  ],
+  "0.02": [
+   3.7449042593793065,
+   3.74560978012876,
+   3.7447914123657418
+  ],
+  "0.05": [
+   3.764693616993574,
+   3.7656580669448987,
+   3.7637789552510497
+  ],
+  "0.1": [
+   3.8448218197141966,
+   3.842812259470448,
+   3.8400778085714027
+  ],
+  "0.15": [
+   4.014115305391468,
+   3.9956352023966195,
+   3.998548066912348
+  ],
+  "0.2": [
+   4.336877310147849,
+   4.268974573538131,
+   4.296562694990125
+  ],
+  "0.3": [
+   5.680166579468835,
+   5.429263905080685,
+   5.506133190586791
+  ],
+  "0.4": [
+   7.38970396117274,
+   7.121290946674239,
+   7.185266885769214
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.741544666384729,
+    3.741646793636223,
+    3.7415041934144293
+   ],
+   "0.01": [
+    3.742238622310576,
+    3.7425577856207837,
+    3.7421255522946613
+   ],
+   "0.02": [
+    3.7449042593793065,
+    3.74560978012876,
+    3.7447914123657418
+   ],
+   "0.05": [
+    3.764693616993574,
+    3.7656580669448987,
+    3.7637789552510497
+   ],
+   "0.1": [
+    3.8448218197141966,
+    3.842812259470448,
+    3.8400778085714027
+   ],
+   "0.15": [
+    4.014115305391468,
+    3.9956352023966195,
+    3.998548066912348
+   ],
+   "0.2": [
+    4.336877310147849,
+    4.268974573538131,
+    4.296562694990125
+   ],
+   "0.3": [
+    5.680166579468835,
+    5.429263905080685,
+    5.506133190586791
+   ],
+   "0.4": [
+    7.38970396117274,
+    7.121290946674239,
+    7.185266885769214
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0002831190883254564,
+   "0.01": 0.0010252213518713305,
+   "0.02": 0.0038197185678003898,
+   "0.05": 0.02342811433970482,
+   "0.1": 0.10128853052854714,
+   "0.15": 0.2614840928433435,
+   "0.2": 0.5595227608352329,
+   "0.3": 1.7972391263219678,
+   "0.4": 3.4908051658152615
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0002831190883254564,
+  "0.01": 0.0010252213518713305,
+  "0.02": 0.0038197185678003898,
+  "0.05": 0.02342811433970482,
+  "0.1": 0.10128853052854714,
+  "0.15": 0.2614840928433435,
+  "0.2": 0.5595227608352329,
+  "0.3": 1.7972391263219678,
+  "0.4": 3.4908051658152615
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 49.9
+}

pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101/ptq.json

@@ -0,0 +1,48 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr1.0e-02-alr2.4e-03-s101",
+ "w_bits": null,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.741282098723469
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 5.78544710766109
+  },
+  "4": {
+   "val_dolma_cc": 3.9773186148965696
+  },
+  "5": {
+   "val_dolma_cc": 3.7802833630607204
+  },
+  "6": {
+   "val_dolma_cc": 3.751075317329996
+  },
+  "8": {
+   "val_dolma_cc": 3.7417057523276123
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 2.0441650089376213
+  },
+  "4": {
+   "val_dolma_cc": 0.23603651617310062
+  },
+  "5": {
+   "val_dolma_cc": 0.039001264337251484
+  },
+  "6": {
+   "val_dolma_cc": 0.009793218606527265
+  },
+  "8": {
+   "val_dolma_cc": 0.00042365360414331477
+  }
+ },
+ "checks": {
+  "weights_restored": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": null,
+  "a": null,
+  "qkv": null
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 3.7980115686420324,
+  "heldout_cc_en_tail_4M": 4.380696555852657
+ },
+ "loss_bf16_master": 3.7980115686420324,
+ "losses_eval_model": {
+  "val_dolma_cc": 3.7980115686420324
+ },
+ "bf16_minus_quant": 0.0,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-q9-28",
+ "gpu": "NVIDIA L40S",
+ "t_s": 17.9
+}

pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/finalize.json

@@ -0,0 +1,29 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101",
+ "w_bits": null,
+ "a_bits": null,
+ "qkv_bits": null,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.7980115686420324
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.7980115686420324
+ },
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.8346741199493408,
+ "loss_eval_model": 3.7980115686420324,
+ "checks": {
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/launch.yaml

@@ -0,0 +1,126 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 101
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: null
+  fq_act_bits: null
+  fq_qkv_bits: null
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: precq.muon26.MuonConfig
+    lr: 0.0025
+    weight_decay: 0.1
+    momentum: 0.95
+    nesterov: true
+    ns_steps: 5
+    eps: 1.0e-07
+    adjust_lr_fn: null
+    adamw_lr: 0.0024
+    adamw_betas:
+    - 0.9
+    - 0.95
+    adamw_eps: 1.0e-15
+    adamw_weight_decay: 0.1
+    adamw_fused: true
+    muon_on_lm_head: false
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 701
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1007019
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 101
+n_tokens: 1839988736
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.02823960967361927,
+  "blocks.0.attention.w_out.weight": 0.012285185977816582,
+  "blocks.0.attention.w_q.weight": 0.027905184775590897,
+  "blocks.0.attention.w_v.weight": 0.015026269480586052,
+  "blocks.0.attention_norm.weight": 0.36402034759521484,
+  "blocks.0.feed_forward.w1.weight": 0.018776489421725273,
+  "blocks.0.feed_forward.w2.weight": 0.011734351515769958,
+  "blocks.0.feed_forward.w3.weight": 0.01938738487660885,
+  "blocks.0.feed_forward_norm.weight": 0.40471863746643066,
+  "blocks.1.attention.w_k.weight": 0.021475251764059067,
+  "blocks.1.attention.w_out.weight": 0.013774644583463669,
+  "blocks.1.attention.w_q.weight": 0.02127489447593689,
+  "blocks.1.attention.w_v.weight": 0.01492380816489458,
+  "blocks.1.attention_norm.weight": 0.3935880959033966,
+  "blocks.1.feed_forward.w1.weight": 0.019992079585790634,
+  "blocks.1.feed_forward.w2.weight": 0.012382809072732925,
+  "blocks.1.feed_forward.w3.weight": 0.020290696993470192,
+  "blocks.1.feed_forward_norm.weight": 0.4071131944656372,
+  "blocks.2.attention.w_k.weight": 0.025532985106110573,
+  "blocks.2.attention.w_out.weight": 0.015362584963440895,
+  "blocks.2.attention.w_q.weight": 0.024703780189156532,
+  "blocks.2.attention.w_v.weight": 0.017002055421471596,
+  "blocks.2.attention_norm.weight": 0.404275506734848,
+  "blocks.2.feed_forward.w1.weight": 0.01942611113190651,
+  "blocks.2.feed_forward.w2.weight": 0.012278757989406586,
+  "blocks.2.feed_forward.w3.weight": 0.019887568429112434,
+  "blocks.2.feed_forward_norm.weight": 0.40656808018684387,
+  "blocks.3.attention.w_k.weight": 0.02614908292889595,
+  "blocks.3.attention.w_out.weight": 0.015408954583108425,
+  "blocks.3.attention.w_q.weight": 0.025211982429027557,
+  "blocks.3.attention.w_v.weight": 0.01707778498530388,
+  "blocks.3.attention_norm.weight": 0.40307778120040894,
+  "blocks.3.feed_forward.w1.weight": 0.018848182633519173,
+  "blocks.3.feed_forward.w2.weight": 0.011943705379962921,
+  "blocks.3.feed_forward.w3.weight": 0.01950244978070259,
+  "blocks.3.feed_forward_norm.weight": 0.4055529236793518,
+  "blocks.4.attention.w_k.weight": 0.026735495775938034,
+  "blocks.4.attention.w_out.weight": 0.016208117827773094,
+  "blocks.4.attention.w_q.weight": 0.025898175314068794,
+  "blocks.4.attention.w_v.weight": 0.017795512452721596,
+  "blocks.4.attention_norm.weight": 0.39835861325263977,
+  "blocks.4.feed_forward.w1.weight": 0.018612002953886986,
+  "blocks.4.feed_forward.w2.weight": 0.011872970499098301,
+  "blocks.4.feed_forward.w3.weight": 0.019427184015512466,
+  "blocks.4.feed_forward_norm.weight": 0.40461936593055725,
+  "blocks.5.attention.w_k.weight": 0.026231853291392326,
+  "blocks.5.attention.w_out.weight": 0.01922769658267498,
+  "blocks.5.attention.w_q.weight": 0.025986338034272194,
+  "blocks.5.attention.w_v.weight": 0.019273919984698296,
+  "blocks.5.attention_norm.weight": 0.4032162129878998,
+  "blocks.5.feed_forward.w1.weight": 0.018274476751685143,
+  "blocks.5.feed_forward.w2.weight": 0.011750449426472187,
+  "blocks.5.feed_forward.w3.weight": 0.019116736948490143,
+  "blocks.5.feed_forward_norm.weight": 0.40511777997016907,
+  "blocks.6.attention.w_k.weight": 0.027327269315719604,
+  "blocks.6.attention.w_out.weight": 0.017924083396792412,
+  "blocks.6.attention.w_q.weight": 0.026071004569530487,
+  "blocks.6.attention.w_v.weight": 0.019368091598153114,
+  "blocks.6.attention_norm.weight": 0.4013967216014862,
+  "blocks.6.feed_forward.w1.weight": 0.018336961045861244,
+  "blocks.6.feed_forward.w2.weight": 0.012955289334058762,
+  "blocks.6.feed_forward.w3.weight": 0.01925334520637989,
+  "blocks.6.feed_forward_norm.weight": 0.4079210162162781,
+  "blocks.7.attention.w_k.weight": 0.02898509055376053,
+  "blocks.7.attention.w_out.weight": 0.020589452236890793,
+  "blocks.7.attention.w_q.weight": 0.027352629229426384,
+  "blocks.7.attention.w_v.weight": 0.020123351365327835,
+  "blocks.7.attention_norm.weight": 0.4034270644187927,
+  "blocks.7.feed_forward.w1.weight": 0.01795170083642006,
+  "blocks.7.feed_forward.w2.weight": 0.01468669530004263,
+  "blocks.7.feed_forward.w3.weight": 0.019164539873600006,
+  "blocks.7.feed_forward_norm.weight": 0.40853628516197205,
+  "blocks.8.attention.w_k.weight": 0.02681003138422966,
+  "blocks.8.attention.w_out.weight": 0.025461845099925995,
+  "blocks.8.attention.w_q.weight": 0.025192417204380035,
+  "blocks.8.attention.w_v.weight": 0.0201339740306139,
+  "blocks.8.attention_norm.weight": 0.40670523047447205,
+  "blocks.8.feed_forward.w1.weight": 0.016974635422229767,
+  "blocks.8.feed_forward.w2.weight": 0.017902221530675888,
+  "blocks.8.feed_forward.w3.weight": 0.01865983195602894,
+  "blocks.8.feed_forward_norm.weight": 0.4039677381515503,
+  "embeddings.weight": 0.07120760530233383,
+  "lm_head.norm.weight": 1.4134982824325562,
+  "lm_head.w_out.weight": 0.0739491730928421
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.7980115686420324
+ },
+ "base_loss": 3.7980115686420324,
+ "loss": {
+  "0.005": [
+   3.798565494741589,
+   3.798943397653158,
+   3.7988173856763194
+  ],
+  "0.01": [
+   3.8004938782177438,
+   3.801210975527769,
+   3.800666517247268
+  ],
+  "0.02": [
+   3.808603229556021,
+   3.8101174585456627,
+   3.808464905502089
+  ],
+  "0.05": [
+   3.869149998923329,
+   3.873003010046592,
+   3.8638678557245845
+  ],
+  "0.1": [
+   4.129703283871569,
+   4.14031559402067,
+   4.10362852287375
+  ],
+  "0.15": [
+   4.732539682837764,
+   4.720558259543812,
+   4.678776374996494
+  ],
+  "0.2": [
+   5.877093123527715,
+   5.69483182430382,
+   5.735007019957474
+  ],
+  "0.3": [
+   8.204661368819409,
+   8.114015678896147,
+   8.297994659586678
+  ],
+  "0.4": [
+   9.525585677721825,
+   9.636142527577801,
+   10.106714495897632
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.798565494741589,
+    3.798943397653158,
+    3.7988173856763194
+   ],
+   "0.01": [
+    3.8004938782177438,
+    3.801210975527769,
+    3.800666517247268
+   ],
+   "0.02": [
+    3.808603229556021,
+    3.8101174585456627,
+    3.808464905502089
+   ],
+   "0.05": [
+    3.869149998923329,
+    3.873003010046592,
+    3.8638678557245845
+   ],
+   "0.1": [
+    4.129703283871569,
+    4.14031559402067,
+    4.10362852287375
+   ],
+   "0.15": [
+    4.732539682837764,
+    4.720558259543812,
+    4.678776374996494
+   ],
+   "0.2": [
+    5.877093123527715,
+    5.69483182430382,
+    5.735007019957474
+   ],
+   "0.3": [
+    8.204661368819409,
+    8.114015678896147,
+    8.297994659586678
+   ],
+   "0.4": [
+    9.525585677721825,
+    9.636142527577801,
+    10.106714495897632
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0007638573816564076,
+   "0.01": 0.002778888355561726,
+   "0.02": 0.011050295892558903,
+   "0.05": 0.07066205292280348,
+   "0.1": 0.32653756494663044,
+   "0.15": 0.9126132038173247,
+   "0.2": 1.9709657539543035,
+   "0.3": 4.407545667125378,
+   "0.4": 5.958135998423721
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0007638573816564076,
+  "0.01": 0.002778888355561726,
+  "0.02": 0.011050295892558903,
+  "0.05": 0.07066205292280348,
+  "0.1": 0.32653756494663044,
+  "0.15": 0.9126132038173247,
+  "0.2": 1.9709657539543035,
+  "0.3": 4.407545667125378,
+  "0.4": 5.958135998423721
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 57.6
+}

pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.02823960967361927,
+  "blocks.0.attention.w_out.weight": 0.012285185977816582,
+  "blocks.0.attention.w_q.weight": 0.027905184775590897,
+  "blocks.0.attention.w_v.weight": 0.015026269480586052,
+  "blocks.0.feed_forward.w1.weight": 0.018776489421725273,
+  "blocks.0.feed_forward.w2.weight": 0.011734351515769958,
+  "blocks.0.feed_forward.w3.weight": 0.01938738487660885,
+  "blocks.1.attention.w_k.weight": 0.021475251764059067,
+  "blocks.1.attention.w_out.weight": 0.013774644583463669,
+  "blocks.1.attention.w_q.weight": 0.02127489447593689,
+  "blocks.1.attention.w_v.weight": 0.01492380816489458,
+  "blocks.1.feed_forward.w1.weight": 0.019992079585790634,
+  "blocks.1.feed_forward.w2.weight": 0.012382809072732925,
+  "blocks.1.feed_forward.w3.weight": 0.020290696993470192,
+  "blocks.2.attention.w_k.weight": 0.025532985106110573,
+  "blocks.2.attention.w_out.weight": 0.015362584963440895,
+  "blocks.2.attention.w_q.weight": 0.024703780189156532,
+  "blocks.2.attention.w_v.weight": 0.017002055421471596,
+  "blocks.2.feed_forward.w1.weight": 0.01942611113190651,
+  "blocks.2.feed_forward.w2.weight": 0.012278757989406586,
+  "blocks.2.feed_forward.w3.weight": 0.019887568429112434,
+  "blocks.3.attention.w_k.weight": 0.02614908292889595,
+  "blocks.3.attention.w_out.weight": 0.015408954583108425,
+  "blocks.3.attention.w_q.weight": 0.025211982429027557,
+  "blocks.3.attention.w_v.weight": 0.01707778498530388,
+  "blocks.3.feed_forward.w1.weight": 0.018848182633519173,
+  "blocks.3.feed_forward.w2.weight": 0.011943705379962921,
+  "blocks.3.feed_forward.w3.weight": 0.01950244978070259,
+  "blocks.4.attention.w_k.weight": 0.026735495775938034,
+  "blocks.4.attention.w_out.weight": 0.016208117827773094,
+  "blocks.4.attention.w_q.weight": 0.025898175314068794,
+  "blocks.4.attention.w_v.weight": 0.017795512452721596,
+  "blocks.4.feed_forward.w1.weight": 0.018612002953886986,
+  "blocks.4.feed_forward.w2.weight": 0.011872970499098301,
+  "blocks.4.feed_forward.w3.weight": 0.019427184015512466,
+  "blocks.5.attention.w_k.weight": 0.026231853291392326,
+  "blocks.5.attention.w_out.weight": 0.01922769658267498,
+  "blocks.5.attention.w_q.weight": 0.025986338034272194,
+  "blocks.5.attention.w_v.weight": 0.019273919984698296,
+  "blocks.5.feed_forward.w1.weight": 0.018274476751685143,
+  "blocks.5.feed_forward.w2.weight": 0.011750449426472187,
+  "blocks.5.feed_forward.w3.weight": 0.019116736948490143,
+  "blocks.6.attention.w_k.weight": 0.027327269315719604,
+  "blocks.6.attention.w_out.weight": 0.017924083396792412,
+  "blocks.6.attention.w_q.weight": 0.026071004569530487,
+  "blocks.6.attention.w_v.weight": 0.019368091598153114,
+  "blocks.6.feed_forward.w1.weight": 0.018336961045861244,
+  "blocks.6.feed_forward.w2.weight": 0.012955289334058762,
+  "blocks.6.feed_forward.w3.weight": 0.01925334520637989,
+  "blocks.7.attention.w_k.weight": 0.02898509055376053,
+  "blocks.7.attention.w_out.weight": 0.020589452236890793,
+  "blocks.7.attention.w_q.weight": 0.027352629229426384,
+  "blocks.7.attention.w_v.weight": 0.020123351365327835,
+  "blocks.7.feed_forward.w1.weight": 0.01795170083642006,
+  "blocks.7.feed_forward.w2.weight": 0.01468669530004263,
+  "blocks.7.feed_forward.w3.weight": 0.019164539873600006,
+  "blocks.8.attention.w_k.weight": 0.02681003138422966,
+  "blocks.8.attention.w_out.weight": 0.025461845099925995,
+  "blocks.8.attention.w_q.weight": 0.025192417204380035,
+  "blocks.8.attention.w_v.weight": 0.0201339740306139,
+  "blocks.8.feed_forward.w1.weight": 0.016974635422229767,
+  "blocks.8.feed_forward.w2.weight": 0.017902221530675888,
+  "blocks.8.feed_forward.w3.weight": 0.01865983195602894
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": null,
+ "base_losses": {
+  "val_dolma_cc": 3.7980115686420324
+ },
+ "base_loss": 3.7980115686420324,
+ "loss": {
+  "0.005": [
+   3.798558264770674,
+   3.7982190821091835,
+   3.798555240545827
+  ],
+  "0.01": [
+   3.7999863209576263,
+   3.7992850724637135,
+   3.7998101021535424
+  ],
+  "0.02": [
+   3.8054275391657852,
+   3.8040605073761977,
+   3.8053264134687033
+  ],
+  "0.05": [
+   3.8446402723276525,
+   3.8436552592122974,
+   3.846435186941383
+  ],
+  "0.1": [
+   4.01551527501933,
+   4.025724475166709,
+   4.028125819240211
+  ],
+  "0.15": [
+   4.412840005031559,
+   4.434557264906385,
+   4.445230928087892
+  ],
+  "0.2": [
+   5.227415800885474,
+   5.085496949077468,
+   5.149504562330737
+  ],
+  "0.3": [
+   7.651915984545514,
+   6.825729054930752,
+   6.897205173366565
+  ],
+  "0.4": [
+   8.87147847208868,
+   8.323283079088746,
+   8.34723364098581
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.798558264770674,
+    3.7982190821091835,
+    3.798555240545827
+   ],
+   "0.01": [
+    3.7999863209576263,
+    3.7992850724637135,
+    3.7998101021535424
+   ],
+   "0.02": [
+    3.8054275391657852,
+    3.8040605073761977,
+    3.8053264134687033
+   ],
+   "0.05": [
+    3.8446402723276525,
+    3.8436552592122974,
+    3.846435186941383
+   ],
+   "0.1": [
+    4.01551527501933,
+    4.025724475166709,
+    4.028125819240211
+   ],
+   "0.15": [
+    4.412840005031559,
+    4.434557264906385,
+    4.445230928087892
+   ],
+   "0.2": [
+    5.227415800885474,
+    5.085496949077468,
+    5.149504562330737
+   ],
+   "0.3": [
+    7.651915984545514,
+    6.825729054930752,
+    6.897205173366565
+   ],
+   "0.4": [
+    8.87147847208868,
+    8.323283079088746,
+    8.34723364098581
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0004326271665293646,
+   "0.01": 0.001682263216261859,
+   "0.02": 0.00692658469486318,
+   "0.05": 0.04689867085174537,
+   "0.1": 0.22511028783338416,
+   "0.15": 0.6328644973665791,
+   "0.2": 1.3561275354558604,
+   "0.3": 3.3269385023055773,
+   "0.4": 4.715986828745715
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0004326271665293646,
+  "0.01": 0.001682263216261859,
+  "0.02": 0.00692658469486318,
+  "0.05": 0.04689867085174537,
+  "0.1": 0.22511028783338416,
+  "0.15": 0.6328644973665791,
+  "0.2": 1.3561275354558604,
+  "0.3": 3.3269385023055773,
+  "0.4": 4.715986828745715
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 51.6
+}

pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101/ptq.json

@@ -0,0 +1,48 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-bf16-muon-lr2.5e-03-alr2.4e-03-s101",
+ "w_bits": null,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.7980115686420324
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 7.490522650528621
+  },
+  "4": {
+   "val_dolma_cc": 4.375064463328336
+  },
+  "5": {
+   "val_dolma_cc": 3.8832731915557948
+  },
+  "6": {
+   "val_dolma_cc": 3.8167199790950557
+  },
+  "8": {
+   "val_dolma_cc": 3.7995107586662
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 3.6925110818865883
+  },
+  "4": {
+   "val_dolma_cc": 0.5770528946863034
+  },
+  "5": {
+   "val_dolma_cc": 0.08526162291376238
+  },
+  "6": {
+   "val_dolma_cc": 0.018708410453023294
+  },
+  "8": {
+   "val_dolma_cc": 0.0014991900241674294
+  }
+ },
+ "checks": {
+  "weights_restored": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": 6,
+  "a": 6,
+  "qkv": 6
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 3.872419299440231,
+  "heldout_cc_en_tail_4M": 4.45661268858849
+ },
+ "loss_bf16_master": 3.872419299440231,
+ "losses_eval_model": {
+  "val_dolma_cc": 3.904266601452955
+ },
+ "bf16_minus_quant": -0.03184730201272368,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-q9-24",
+ "gpu": "NVIDIA L40S",
+ "t_s": 18.6
+}

pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/finalize.json

@@ -0,0 +1,102 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101",
+ "w_bits": 6,
+ "a_bits": 6,
+ "qkv_bits": 6,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.902572777325409
+ },
+ "losses_fakequant_master": {
+  "val_dolma_cc": 3.904266601452955
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.904266601452955
+ },
+ "max_distinct_values_per_row": 62,
+ "rel_quant_noise_per_layer": {
+  "blocks.0.attention.w_q": 0.03132954612374306,
+  "blocks.0.attention.w_k": 0.032565560191869736,
+  "blocks.0.attention.w_v": 0.03376292809844017,
+  "blocks.0.attention.w_out": 0.035142719745635986,
+  "blocks.0.feed_forward.w1": 0.030101409181952477,
+  "blocks.0.feed_forward.w2": 0.03631531447172165,
+  "blocks.0.feed_forward.w3": 0.030498016625642776,
+  "blocks.1.attention.w_q": 0.03087511472404003,
+  "blocks.1.attention.w_k": 0.03758403658866882,
+  "blocks.1.attention.w_v": 0.031966060400009155,
+  "blocks.1.attention.w_out": 0.03447512537240982,
+  "blocks.1.feed_forward.w1": 0.03064761683344841,
+  "blocks.1.feed_forward.w2": 0.036250922828912735,
+  "blocks.1.feed_forward.w3": 0.03089948371052742,
+  "blocks.2.attention.w_q": 0.030825156718492508,
+  "blocks.2.attention.w_k": 0.03836624696850777,
+  "blocks.2.attention.w_v": 0.03155285865068436,
+  "blocks.2.attention.w_out": 0.03362602740526199,
+  "blocks.2.feed_forward.w1": 0.030612805858254433,
+  "blocks.2.feed_forward.w2": 0.036123599857091904,
+  "blocks.2.feed_forward.w3": 0.030936317518353462,
+  "blocks.3.attention.w_q": 0.03071513958275318,
+  "blocks.3.attention.w_k": 0.03686222434043884,
+  "blocks.3.attention.w_v": 0.03170269727706909,
+  "blocks.3.attention.w_out": 0.03152976930141449,
+  "blocks.3.feed_forward.w1": 0.030767356976866722,
+  "blocks.3.feed_forward.w2": 0.036432571709156036,
+  "blocks.3.feed_forward.w3": 0.031014615669846535,
+  "blocks.4.attention.w_q": 0.031317088752985,
+  "blocks.4.attention.w_k": 0.04460396245121956,
+  "blocks.4.attention.w_v": 0.032061345875263214,
+  "blocks.4.attention.w_out": 0.03140087053179741,
+  "blocks.4.feed_forward.w1": 0.03076915629208088,
+  "blocks.4.feed_forward.w2": 0.0370493121445179,
+  "blocks.4.feed_forward.w3": 0.031176608055830002,
+  "blocks.5.attention.w_q": 0.03135261312127113,
+  "blocks.5.attention.w_k": 0.04125480353832245,
+  "blocks.5.attention.w_v": 0.032418251037597656,
+  "blocks.5.attention.w_out": 0.03230922296643257,
+  "blocks.5.feed_forward.w1": 0.031060615554451942,
+  "blocks.5.feed_forward.w2": 0.037558216601610184,
+  "blocks.5.feed_forward.w3": 0.031182829290628433,
+  "blocks.6.attention.w_q": 0.031420376151800156,
+  "blocks.6.attention.w_k": 0.040330782532691956,
+  "blocks.6.attention.w_v": 0.0323045440018177,
+  "blocks.6.attention.w_out": 0.037702176719903946,
+  "blocks.6.feed_forward.w1": 0.030968982726335526,
+  "blocks.6.feed_forward.w2": 0.03673739731311798,
+  "blocks.6.feed_forward.w3": 0.031195146963000298,
+  "blocks.7.attention.w_q": 0.03101927787065506,
+  "blocks.7.attention.w_k": 0.04204888641834259,
+  "blocks.7.attention.w_v": 0.032408490777015686,
+  "blocks.7.attention.w_out": 0.03582685440778732,
+  "blocks.7.feed_forward.w1": 0.03105190210044384,
+  "blocks.7.feed_forward.w2": 0.03869321197271347,
+  "blocks.7.feed_forward.w3": 0.03149225935339928,
+  "blocks.8.attention.w_q": 0.03125115856528282,
+  "blocks.8.attention.w_k": 0.033950477838516235,
+  "blocks.8.attention.w_v": 0.03320341184735298,
+  "blocks.8.attention.w_out": 0.03774721175432205,
+  "blocks.8.feed_forward.w1": 0.031206922605633736,
+  "blocks.8.feed_forward.w2": 0.04580705240368843,
+  "blocks.8.feed_forward.w3": 0.03194732218980789
+ },
+ "rel_quant_noise_mean": 0.03376682516601351,
+ "n_quantized_params": 30670848,
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.8170008659362793,
+ "loss_eval_model": 3.904266601452955,
+ "checks": {
+  "snap_reproduces_fakequant": true,
+  "grid_levels_ok": true,
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/launch.yaml

@@ -0,0 +1,118 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 101
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: 6
+  fq_act_bits: 6
+  fq_qkv_bits: 6
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: olmo_core.optim.AdamWConfig
+    lr: 0.0024
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-15
+    weight_decay: 0.1
+    fused: true
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 701
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1007019
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 101
+n_tokens: 1839988736
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.04447223246097565,
+  "blocks.0.attention.w_out.weight": 0.06222694367170334,
+  "blocks.0.attention.w_q.weight": 0.04480505734682083,
+  "blocks.0.attention.w_v.weight": 0.11886989325284958,
+  "blocks.0.attention_norm.weight": 0.404301255941391,
+  "blocks.0.feed_forward.w1.weight": 0.057579103857278824,
+  "blocks.0.feed_forward.w2.weight": 0.04649105295538902,
+  "blocks.0.feed_forward.w3.weight": 0.06854056566953659,
+  "blocks.0.feed_forward_norm.weight": 0.398822158575058,
+  "blocks.1.attention.w_k.weight": 0.04462267458438873,
+  "blocks.1.attention.w_out.weight": 0.0488797165453434,
+  "blocks.1.attention.w_q.weight": 0.03434634581208229,
+  "blocks.1.attention.w_v.weight": 0.08584452420473099,
+  "blocks.1.attention_norm.weight": 0.36681878566741943,
+  "blocks.1.feed_forward.w1.weight": 0.062314797192811966,
+  "blocks.1.feed_forward.w2.weight": 0.05205081030726433,
+  "blocks.1.feed_forward.w3.weight": 0.07317005097866058,
+  "blocks.1.feed_forward_norm.weight": 0.4142328202724457,
+  "blocks.2.attention.w_k.weight": 0.044960957020521164,
+  "blocks.2.attention.w_out.weight": 0.05148233845829964,
+  "blocks.2.attention.w_q.weight": 0.03788594901561737,
+  "blocks.2.attention.w_v.weight": 0.10192523151636124,
+  "blocks.2.attention_norm.weight": 0.3903007209300995,
+  "blocks.2.feed_forward.w1.weight": 0.06271107494831085,
+  "blocks.2.feed_forward.w2.weight": 0.054983798414468765,
+  "blocks.2.feed_forward.w3.weight": 0.073421411216259,
+  "blocks.2.feed_forward_norm.weight": 0.4138107895851135,
+  "blocks.3.attention.w_k.weight": 0.0436176098883152,
+  "blocks.3.attention.w_out.weight": 0.05381327494978905,
+  "blocks.3.attention.w_q.weight": 0.03841737285256386,
+  "blocks.3.attention.w_v.weight": 0.09413135051727295,
+  "blocks.3.attention_norm.weight": 0.38852155208587646,
+  "blocks.3.feed_forward.w1.weight": 0.06022169068455696,
+  "blocks.3.feed_forward.w2.weight": 0.057319436222314835,
+  "blocks.3.feed_forward.w3.weight": 0.07215607166290283,
+  "blocks.3.feed_forward_norm.weight": 0.41180163621902466,
+  "blocks.4.attention.w_k.weight": 0.04765772819519043,
+  "blocks.4.attention.w_out.weight": 0.06548422574996948,
+  "blocks.4.attention.w_q.weight": 0.033922698348760605,
+  "blocks.4.attention.w_v.weight": 0.12088527530431747,
+  "blocks.4.attention_norm.weight": 0.40293774008750916,
+  "blocks.4.feed_forward.w1.weight": 0.05975082889199257,
+  "blocks.4.feed_forward.w2.weight": 0.06048689782619476,
+  "blocks.4.feed_forward.w3.weight": 0.07176288217306137,
+  "blocks.4.feed_forward_norm.weight": 0.41228407621383667,
+  "blocks.5.attention.w_k.weight": 0.045243971049785614,
+  "blocks.5.attention.w_out.weight": 0.08075490593910217,
+  "blocks.5.attention.w_q.weight": 0.034430962055921555,
+  "blocks.5.attention.w_v.weight": 0.07127781957387924,
+  "blocks.5.attention_norm.weight": 0.39218249917030334,
+  "blocks.5.feed_forward.w1.weight": 0.059570491313934326,
+  "blocks.5.feed_forward.w2.weight": 0.06447386741638184,
+  "blocks.5.feed_forward.w3.weight": 0.0718831717967987,
+  "blocks.5.feed_forward_norm.weight": 0.41259849071502686,
+  "blocks.6.attention.w_k.weight": 0.050092969089746475,
+  "blocks.6.attention.w_out.weight": 0.09478561580181122,
+  "blocks.6.attention.w_q.weight": 0.03283941373229027,
+  "blocks.6.attention.w_v.weight": 0.0808403417468071,
+  "blocks.6.attention_norm.weight": 0.40019938349723816,
+  "blocks.6.feed_forward.w1.weight": 0.0594840906560421,
+  "blocks.6.feed_forward.w2.weight": 0.075862817466259,
+  "blocks.6.feed_forward.w3.weight": 0.07260943949222565,
+  "blocks.6.feed_forward_norm.weight": 0.41187795996665955,
+  "blocks.7.attention.w_k.weight": 0.050687216222286224,
+  "blocks.7.attention.w_out.weight": 0.10306242853403091,
+  "blocks.7.attention.w_q.weight": 0.03203458711504936,
+  "blocks.7.attention.w_v.weight": 0.09186305850744247,
+  "blocks.7.attention_norm.weight": 0.40369895100593567,
+  "blocks.7.feed_forward.w1.weight": 0.057336632162332535,
+  "blocks.7.feed_forward.w2.weight": 0.08863422274589539,
+  "blocks.7.feed_forward.w3.weight": 0.07244992256164551,
+  "blocks.7.feed_forward_norm.weight": 0.4105563461780548,
+  "blocks.8.attention.w_k.weight": 0.0413111075758934,
+  "blocks.8.attention.w_out.weight": 0.17310740053653717,
+  "blocks.8.attention.w_q.weight": 0.03296327218413353,
+  "blocks.8.attention.w_v.weight": 0.09154683351516724,
+  "blocks.8.attention_norm.weight": 0.4069346785545349,
+  "blocks.8.feed_forward.w1.weight": 0.05462755262851715,
+  "blocks.8.feed_forward.w2.weight": 0.10933306068181992,
+  "blocks.8.feed_forward.w3.weight": 0.07172390073537827,
+  "blocks.8.feed_forward_norm.weight": 0.40633541345596313,
+  "embeddings.weight": 0.06951554119586945,
+  "lm_head.norm.weight": 1.401055932044983,
+  "lm_head.w_out.weight": 0.07445036619901657
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 3.904266601452955
+ },
+ "base_loss": 3.904266601452955,
+ "loss": {
+  "0.005": [
+   3.904142573023332,
+   3.905856580819682,
+   3.904488195366709
+  ],
+  "0.01": [
+   3.906111492066938,
+   3.905960675969104,
+   3.9051861183869865
+  ],
+  "0.02": [
+   3.909523917056772,
+   3.909952218499361,
+   3.9093763325368474
+  ],
+  "0.05": [
+   3.925472618660411,
+   3.9244579266709434,
+   3.9253922126202365
+  ],
+  "0.1": [
+   3.984940330517964,
+   3.9796524988031425,
+   3.978508801910565
+  ],
+  "0.15": [
+   4.09187802438134,
+   4.074926657265826,
+   4.0745152218430505
+  ],
+  "0.2": [
+   4.262114352255951,
+   4.221008223512892,
+   4.219886568642033
+  ],
+  "0.3": [
+   4.900139414809284,
+   4.741884334174138,
+   4.7508327298968
+  ],
+  "0.4": [
+   6.209186188214108,
+   5.761083701905357,
+   5.801262238681934
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.904142573023332,
+    3.905856580819682,
+    3.904488195366709
+   ],
+   "0.01": [
+    3.906111492066938,
+    3.905960675969104,
+    3.9051861183869865
+   ],
+   "0.02": [
+    3.909523917056772,
+    3.909952218499361,
+    3.9093763325368474
+   ],
+   "0.05": [
+    3.925472618660411,
+    3.9244579266709434,
+    3.9253922126202365
+   ],
+   "0.1": [
+    3.984940330517964,
+    3.9796524988031425,
+    3.978508801910565
+   ],
+   "0.15": [
+    4.09187802438134,
+    4.074926657265826,
+    4.0745152218430505
+   ],
+   "0.2": [
+    4.262114352255951,
+    4.221008223512892,
+    4.219886568642033
+   ],
+   "0.3": [
+    4.900139414809284,
+    4.741884334174138,
+    4.7508327298968
+   ],
+   "0.4": [
+    6.209186188214108,
+    5.761083701905357,
+    5.801262238681934
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0005625149502859195,
+   "0.01": 0.0014861606880547384,
+   "0.02": 0.005350887911372393,
+   "0.05": 0.020840984530908546,
+   "0.1": 0.076767275624269,
+   "0.15": 0.17617336637711745,
+   "0.2": 0.3300697800173369,
+   "0.3": 0.8933522248404531,
+   "0.4": 2.019577441480844
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0005625149502859195,
+  "0.01": 0.0014861606880547384,
+  "0.02": 0.005350887911372393,
+  "0.05": 0.020840984530908546,
+  "0.1": 0.076767275624269,
+  "0.15": 0.17617336637711745,
+  "0.2": 0.3300697800173369,
+  "0.3": 0.8933522248404531,
+  "0.4": 2.019577441480844
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 78.5
+}

pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.04447223246097565,
+  "blocks.0.attention.w_out.weight": 0.06222694367170334,
+  "blocks.0.attention.w_q.weight": 0.04480505734682083,
+  "blocks.0.attention.w_v.weight": 0.11886989325284958,
+  "blocks.0.feed_forward.w1.weight": 0.057579103857278824,
+  "blocks.0.feed_forward.w2.weight": 0.04649105295538902,
+  "blocks.0.feed_forward.w3.weight": 0.06854056566953659,
+  "blocks.1.attention.w_k.weight": 0.04462267458438873,
+  "blocks.1.attention.w_out.weight": 0.0488797165453434,
+  "blocks.1.attention.w_q.weight": 0.03434634581208229,
+  "blocks.1.attention.w_v.weight": 0.08584452420473099,
+  "blocks.1.feed_forward.w1.weight": 0.062314797192811966,
+  "blocks.1.feed_forward.w2.weight": 0.05205081030726433,
+  "blocks.1.feed_forward.w3.weight": 0.07317005097866058,
+  "blocks.2.attention.w_k.weight": 0.044960957020521164,
+  "blocks.2.attention.w_out.weight": 0.05148233845829964,
+  "blocks.2.attention.w_q.weight": 0.03788594901561737,
+  "blocks.2.attention.w_v.weight": 0.10192523151636124,
+  "blocks.2.feed_forward.w1.weight": 0.06271107494831085,
+  "blocks.2.feed_forward.w2.weight": 0.054983798414468765,
+  "blocks.2.feed_forward.w3.weight": 0.073421411216259,
+  "blocks.3.attention.w_k.weight": 0.0436176098883152,
+  "blocks.3.attention.w_out.weight": 0.05381327494978905,
+  "blocks.3.attention.w_q.weight": 0.03841737285256386,
+  "blocks.3.attention.w_v.weight": 0.09413135051727295,
+  "blocks.3.feed_forward.w1.weight": 0.06022169068455696,
+  "blocks.3.feed_forward.w2.weight": 0.057319436222314835,
+  "blocks.3.feed_forward.w3.weight": 0.07215607166290283,
+  "blocks.4.attention.w_k.weight": 0.04765772819519043,
+  "blocks.4.attention.w_out.weight": 0.06548422574996948,
+  "blocks.4.attention.w_q.weight": 0.033922698348760605,
+  "blocks.4.attention.w_v.weight": 0.12088527530431747,
+  "blocks.4.feed_forward.w1.weight": 0.05975082889199257,
+  "blocks.4.feed_forward.w2.weight": 0.06048689782619476,
+  "blocks.4.feed_forward.w3.weight": 0.07176288217306137,
+  "blocks.5.attention.w_k.weight": 0.045243971049785614,
+  "blocks.5.attention.w_out.weight": 0.08075490593910217,
+  "blocks.5.attention.w_q.weight": 0.034430962055921555,
+  "blocks.5.attention.w_v.weight": 0.07127781957387924,
+  "blocks.5.feed_forward.w1.weight": 0.059570491313934326,
+  "blocks.5.feed_forward.w2.weight": 0.06447386741638184,
+  "blocks.5.feed_forward.w3.weight": 0.0718831717967987,
+  "blocks.6.attention.w_k.weight": 0.050092969089746475,
+  "blocks.6.attention.w_out.weight": 0.09478561580181122,
+  "blocks.6.attention.w_q.weight": 0.03283941373229027,
+  "blocks.6.attention.w_v.weight": 0.0808403417468071,
+  "blocks.6.feed_forward.w1.weight": 0.0594840906560421,
+  "blocks.6.feed_forward.w2.weight": 0.075862817466259,
+  "blocks.6.feed_forward.w3.weight": 0.07260943949222565,
+  "blocks.7.attention.w_k.weight": 0.050687216222286224,
+  "blocks.7.attention.w_out.weight": 0.10306242853403091,
+  "blocks.7.attention.w_q.weight": 0.03203458711504936,
+  "blocks.7.attention.w_v.weight": 0.09186305850744247,
+  "blocks.7.feed_forward.w1.weight": 0.057336632162332535,
+  "blocks.7.feed_forward.w2.weight": 0.08863422274589539,
+  "blocks.7.feed_forward.w3.weight": 0.07244992256164551,
+  "blocks.8.attention.w_k.weight": 0.0413111075758934,
+  "blocks.8.attention.w_out.weight": 0.17310740053653717,
+  "blocks.8.attention.w_q.weight": 0.03296327218413353,
+  "blocks.8.attention.w_v.weight": 0.09154683351516724,
+  "blocks.8.feed_forward.w1.weight": 0.05462755262851715,
+  "blocks.8.feed_forward.w2.weight": 0.10933306068181992,
+  "blocks.8.feed_forward.w3.weight": 0.07172390073537827
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 3.904266601452955
+ },
+ "base_loss": 3.904266601452955,
+ "loss": {
+  "0.005": [
+   3.9049811009812228,
+   3.9050922656963607,
+   3.9045786482495832
+  ],
+  "0.01": [
+   3.9052768881679367,
+   3.9059409226557116,
+   3.9047522105870813
+  ],
+  "0.02": [
+   3.9056773221414636,
+   3.906056570891774,
+   3.906350495813252
+  ],
+  "0.05": [
+   3.9104463643199243,
+   3.909633035627544,
+   3.9087067855874817
+  ],
+  "0.1": [
+   3.927618077318498,
+   3.925788155248195,
+   3.9249304138773375
+  ],
+  "0.15": [
+   3.955862181699748,
+   3.952781071692016,
+   3.952633987949297
+  ],
+  "0.2": [
+   4.0011487398755365,
+   3.9967307581679266,
+   3.995774160246703
+  ],
+  "0.3": [
+   4.152392294090428,
+   4.146487391400203,
+   4.144532861876632
+  ],
+  "0.4": [
+   4.442498969337913,
+   4.438886389958466,
+   4.428868185455077
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.9049811009812228,
+    3.9050922656963607,
+    3.9045786482495832
+   ],
+   "0.01": [
+    3.9052768881679367,
+    3.9059409226557116,
+    3.9047522105870813
+   ],
+   "0.02": [
+    3.9056773221414636,
+    3.906056570891774,
+    3.906350495813252
+   ],
+   "0.05": [
+    3.9104463643199243,
+    3.909633035627544,
+    3.9087067855874817
+   ],
+   "0.1": [
+    3.927618077318498,
+    3.925788155248195,
+    3.9249304138773375
+   ],
+   "0.15": [
+    3.955862181699748,
+    3.952781071692016,
+    3.952633987949297
+   ],
+   "0.2": [
+    4.0011487398755365,
+    3.9967307581679266,
+    3.995774160246703
+   ],
+   "0.3": [
+    4.152392294090428,
+    4.146487391400203,
+    4.144532861876632
+   ],
+   "0.4": [
+    4.442498969337913,
+    4.438886389958466,
+    4.428868185455077
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0006174035227672903,
+   "0.01": 0.0010567390172884217,
+   "0.02": 0.0017615281625413814,
+   "0.05": 0.00532879372536188,
+   "0.1": 0.02184561402838847,
+   "0.15": 0.04949247899406517,
+   "0.2": 0.09361795131043449,
+   "0.3": 0.24353758100280043,
+   "0.4": 0.5324845801308635
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0006174035227672903,
+  "0.01": 0.0010567390172884217,
+  "0.02": 0.0017615281625413814,
+  "0.05": 0.00532879372536188,
+  "0.1": 0.02184561402838847,
+  "0.15": 0.04949247899406517,
+  "0.2": 0.09361795131043449,
+  "0.3": 0.24353758100280043,
+  "0.4": 0.5324845801308635
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 72.9
+}

pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101/ptq.json

@@ -0,0 +1,52 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr2.4e-03-s101",
+ "w_bits": 6,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.904266601452955
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 4.34406060706216
+  },
+  "4": {
+   "val_dolma_cc": 3.9607311680540698
+  },
+  "5": {
+   "val_dolma_cc": 3.923247943308014
+  },
+  "6": {
+   "val_dolma_cc": 3.904266601452955
+  },
+  "8": {
+   "val_dolma_cc": 3.9050877665261488
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 0.43979400560920556
+  },
+  "4": {
+   "val_dolma_cc": 0.056464566601114985
+  },
+  "5": {
+   "val_dolma_cc": 0.018981341855059064
+  },
+  "6": {
+   "val_dolma_cc": 0.0
+  },
+  "8": {
+   "val_dolma_cc": 0.0008211650731939635
+  }
+ },
+ "own_grid_max_rel_weight_change": 2.2933392074264702e-07,
+ "own_grid_delta_loss": 0.0,
+ "checks": {
+  "weights_restored": true,
+  "own_grid_weights_unchanged": true,
+  "own_grid_is_fixed_point": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6-adamw-lr4.8e-03-s101/finalize.json

@@ -0,0 +1,102 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr4.8e-03-s101",
+ "w_bits": 6,
+ "a_bits": 6,
+ "qkv_bits": 6,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.9273053028299922
+ },
+ "losses_fakequant_master": {
+  "val_dolma_cc": 3.9291688200086536
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.9291688200086536
+ },
+ "max_distinct_values_per_row": 62,
+ "rel_quant_noise_per_layer": {
+  "blocks.0.attention.w_q": 0.029958514496684074,
+  "blocks.0.attention.w_k": 0.030053334310650826,
+  "blocks.0.attention.w_v": 0.03284899517893791,
+  "blocks.0.attention.w_out": 0.03345663100481033,
+  "blocks.0.feed_forward.w1": 0.029894156381487846,
+  "blocks.0.feed_forward.w2": 0.03646356984972954,
+  "blocks.0.feed_forward.w3": 0.029902366921305656,
+  "blocks.1.attention.w_q": 0.03054300881922245,
+  "blocks.1.attention.w_k": 0.03423832729458809,
+  "blocks.1.attention.w_v": 0.03337271139025688,
+  "blocks.1.attention.w_out": 0.03670801594853401,
+  "blocks.1.feed_forward.w1": 0.03054652363061905,
+  "blocks.1.feed_forward.w2": 0.03641407936811447,
+  "blocks.1.feed_forward.w3": 0.0308720413595438,
+  "blocks.2.attention.w_q": 0.03073035180568695,
+  "blocks.2.attention.w_k": 0.0352826826274395,
+  "blocks.2.attention.w_v": 0.03186545521020889,
+  "blocks.2.attention.w_out": 0.03409954160451889,
+  "blocks.2.feed_forward.w1": 0.030620494857430458,
+  "blocks.2.feed_forward.w2": 0.039687350392341614,
+  "blocks.2.feed_forward.w3": 0.030994972214102745,
+  "blocks.3.attention.w_q": 0.03141950070858002,
+  "blocks.3.attention.w_k": 0.04081984609365463,
+  "blocks.3.attention.w_v": 0.032464489340782166,
+  "blocks.3.attention.w_out": 0.03311280161142349,
+  "blocks.3.feed_forward.w1": 0.030588535591959953,
+  "blocks.3.feed_forward.w2": 0.0373603031039238,
+  "blocks.3.feed_forward.w3": 0.031179826706647873,
+  "blocks.4.attention.w_q": 0.032007645815610886,
+  "blocks.4.attention.w_k": 0.04560428112745285,
+  "blocks.4.attention.w_v": 0.0322083942592144,
+  "blocks.4.attention.w_out": 0.03341170400381088,
+  "blocks.4.feed_forward.w1": 0.030829299241304398,
+  "blocks.4.feed_forward.w2": 0.03711560368537903,
+  "blocks.4.feed_forward.w3": 0.03147368133068085,
+  "blocks.5.attention.w_q": 0.03180655464529991,
+  "blocks.5.attention.w_k": 0.043339189141988754,
+  "blocks.5.attention.w_v": 0.032215967774391174,
+  "blocks.5.attention.w_out": 0.03398735448718071,
+  "blocks.5.feed_forward.w1": 0.03098096139729023,
+  "blocks.5.feed_forward.w2": 0.03750373050570488,
+  "blocks.5.feed_forward.w3": 0.031507816165685654,
+  "blocks.6.attention.w_q": 0.03224027529358864,
+  "blocks.6.attention.w_k": 0.038728248327970505,
+  "blocks.6.attention.w_v": 0.03304196894168854,
+  "blocks.6.attention.w_out": 0.036059584468603134,
+  "blocks.6.feed_forward.w1": 0.030856654047966003,
+  "blocks.6.feed_forward.w2": 0.038444433361291885,
+  "blocks.6.feed_forward.w3": 0.031284697353839874,
+  "blocks.7.attention.w_q": 0.031709592789411545,
+  "blocks.7.attention.w_k": 0.04747031256556511,
+  "blocks.7.attention.w_v": 0.032274387776851654,
+  "blocks.7.attention.w_out": 0.0332445427775383,
+  "blocks.7.feed_forward.w1": 0.030889054760336876,
+  "blocks.7.feed_forward.w2": 0.040012676268815994,
+  "blocks.7.feed_forward.w3": 0.03180297836661339,
+  "blocks.8.attention.w_q": 0.031886614859104156,
+  "blocks.8.attention.w_k": 0.03702918440103531,
+  "blocks.8.attention.w_v": 0.03320465236902237,
+  "blocks.8.attention.w_out": 0.03827323764562607,
+  "blocks.8.feed_forward.w1": 0.03127485141158104,
+  "blocks.8.feed_forward.w2": 0.04559102654457092,
+  "blocks.8.feed_forward.w3": 0.03272668644785881
+ },
+ "rel_quant_noise_mean": 0.03408787733623906,
+ "n_quantized_params": 30670848,
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.9966011643409729,
+ "loss_eval_model": 3.9291688200086536,
+ "checks": {
+  "snap_reproduces_fakequant": true,
+  "grid_levels_ok": true,
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6-adamw-lr4.8e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr4.8e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.10295666009187698,
+  "blocks.0.attention.w_out.weight": 0.07433991879224777,
+  "blocks.0.attention.w_q.weight": 0.10253647714853287,
+  "blocks.0.attention.w_v.weight": 0.2449776977300644,
+  "blocks.0.feed_forward.w1.weight": 0.09247293323278427,
+  "blocks.0.feed_forward.w2.weight": 0.06845320016145706,
+  "blocks.0.feed_forward.w3.weight": 0.09594220668077469,
+  "blocks.1.attention.w_k.weight": 0.0940186083316803,
+  "blocks.1.attention.w_out.weight": 0.059166502207517624,
+  "blocks.1.attention.w_q.weight": 0.09678329527378082,
+  "blocks.1.attention.w_v.weight": 0.12398216873407364,
+  "blocks.1.feed_forward.w1.weight": 0.09601962566375732,
+  "blocks.1.feed_forward.w2.weight": 0.07874780148267746,
+  "blocks.1.feed_forward.w3.weight": 0.1040811538696289,
+  "blocks.2.attention.w_k.weight": 0.09405351430177689,
+  "blocks.2.attention.w_out.weight": 0.06419295817613602,
+  "blocks.2.attention.w_q.weight": 0.09307216107845306,
+  "blocks.2.attention.w_v.weight": 0.1580667495727539,
+  "blocks.2.feed_forward.w1.weight": 0.09151057153940201,
+  "blocks.2.feed_forward.w2.weight": 0.08589061349630356,
+  "blocks.2.feed_forward.w3.weight": 0.10345649719238281,
+  "blocks.3.attention.w_k.weight": 0.05170898512005806,
+  "blocks.3.attention.w_out.weight": 0.07905052602291107,
+  "blocks.3.attention.w_q.weight": 0.039911024272441864,
+  "blocks.3.attention.w_v.weight": 0.13418684899806976,
+  "blocks.3.feed_forward.w1.weight": 0.08949227631092072,
+  "blocks.3.feed_forward.w2.weight": 0.08585387468338013,
+  "blocks.3.feed_forward.w3.weight": 0.10214531421661377,
+  "blocks.4.attention.w_k.weight": 0.09089221805334091,
+  "blocks.4.attention.w_out.weight": 0.08829694986343384,
+  "blocks.4.attention.w_q.weight": 0.08905211836099625,
+  "blocks.4.attention.w_v.weight": 0.10798521339893341,
+  "blocks.4.feed_forward.w1.weight": 0.08801241219043732,
+  "blocks.4.feed_forward.w2.weight": 0.086526058614254,
+  "blocks.4.feed_forward.w3.weight": 0.09784262627363205,
+  "blocks.5.attention.w_k.weight": 0.05476023629307747,
+  "blocks.5.attention.w_out.weight": 0.09816227108240128,
+  "blocks.5.attention.w_q.weight": 0.03678838536143303,
+  "blocks.5.attention.w_v.weight": 0.14953851699829102,
+  "blocks.5.feed_forward.w1.weight": 0.08689950406551361,
+  "blocks.5.feed_forward.w2.weight": 0.09530370682477951,
+  "blocks.5.feed_forward.w3.weight": 0.09842880815267563,
+  "blocks.6.attention.w_k.weight": 0.046556148678064346,
+  "blocks.6.attention.w_out.weight": 0.14122274518013,
+  "blocks.6.attention.w_q.weight": 0.03567918762564659,
+  "blocks.6.attention.w_v.weight": 0.17810435593128204,
+  "blocks.6.feed_forward.w1.weight": 0.08768820762634277,
+  "blocks.6.feed_forward.w2.weight": 0.11142566800117493,
+  "blocks.6.feed_forward.w3.weight": 0.10013778507709503,
+  "blocks.7.attention.w_k.weight": 0.06066993623971939,
+  "blocks.7.attention.w_out.weight": 0.14124080538749695,
+  "blocks.7.attention.w_q.weight": 0.032068297266960144,
+  "blocks.7.attention.w_v.weight": 0.13702520728111267,
+  "blocks.7.feed_forward.w1.weight": 0.08460602909326553,
+  "blocks.7.feed_forward.w2.weight": 0.13135100901126862,
+  "blocks.7.feed_forward.w3.weight": 0.09995918720960617,
+  "blocks.8.attention.w_k.weight": 0.04496841877698898,
+  "blocks.8.attention.w_out.weight": 0.20446938276290894,
+  "blocks.8.attention.w_q.weight": 0.03546059504151344,
+  "blocks.8.attention.w_v.weight": 0.1397651731967926,
+  "blocks.8.feed_forward.w1.weight": 0.08004491031169891,
+  "blocks.8.feed_forward.w2.weight": 0.14966775476932526,
+  "blocks.8.feed_forward.w3.weight": 0.09858717769384384
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 3.9291688200086536
+ },
+ "base_loss": 3.9291688200086536,
+ "loss": {
+  "0.005": [
+   3.929990462385122,
+   3.9300811226210577,
+   3.9286797484696248
+  ],
+  "0.01": [
+   3.9299742810216456,
+   3.929023837182808,
+   3.9296316007443552
+  ],
+  "0.02": [
+   3.9296030525312786,
+   3.929805515190826,
+   3.930205507071976
+  ],
+  "0.05": [
+   3.9350386917683746,
+   3.934420873539483,
+   3.935254033791695
+  ],
+  "0.1": [
+   3.951707530556173,
+   3.952103301041967,
+   3.953444804931003
+  ],
+  "0.15": [
+   3.984590967994828,
+   3.9838400564864167,
+   3.9871740785449017
+  ],
+  "0.2": [
+   4.031742510469999,
+   4.0354458752796845,
+   4.036791694459802
+  ],
+  "0.3": [
+   4.207225285518124,
+   4.217448648270866,
+   4.214980704741259
+  ],
+  "0.4": [
+   4.549954341638222,
+   4.537920336737617,
+   4.557709361804521
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.929990462385122,
+    3.9300811226210577,
+    3.9286797484696248
+   ],
+   "0.01": [
+    3.9299742810216456,
+    3.929023837182808,
+    3.9296316007443552
+   ],
+   "0.02": [
+    3.9296030525312786,
+    3.929805515190826,
+    3.930205507071976
+   ],
+   "0.05": [
+    3.9350386917683746,
+    3.934420873539483,
+    3.935254033791695
+   ],
+   "0.1": [
+    3.951707530556173,
+    3.952103301041967,
+    3.953444804931003
+   ],
+   "0.15": [
+    3.984590967994828,
+    3.9838400564864167,
+    3.9871740785449017
+   ],
+   "0.2": [
+    4.031742510469999,
+    4.0354458752796845,
+    4.036791694459802
+   ],
+   "0.3": [
+    4.207225285518124,
+    4.217448648270866,
+    4.214980704741259
+   ],
+   "0.4": [
+    4.549954341638222,
+    4.537920336737617,
+    4.557709361804521
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.00041495781661460285,
+   "0.01": 0.0003744196409489753,
+   "0.02": 0.0007025382560397198,
+   "0.05": 0.0057357130245310195,
+   "0.1": 0.023249725501060414,
+   "0.15": 0.056032881000061874,
+   "0.2": 0.10549120672784174,
+   "0.3": 0.2840493928347625,
+   "0.4": 0.6193591933847999
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.00041495781661460285,
+  "0.01": 0.0003744196409489753,
+  "0.02": 0.0007025382560397198,
+  "0.05": 0.0057357130245310195,
+  "0.1": 0.023249725501060414,
+  "0.15": 0.056032881000061874,
+  "0.2": 0.10549120672784174,
+  "0.3": 0.2840493928347625,
+  "0.4": 0.6193591933847999
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 71.7
+}

pom-30M-D1.84B-f6-adamw-lr4.8e-03-s101/ptq.json

@@ -0,0 +1,52 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-adamw-lr4.8e-03-s101",
+ "w_bits": 6,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.9291688200086536
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 4.449217951886328
+  },
+  "4": {
+   "val_dolma_cc": 3.9960233321176966
+  },
+  "5": {
+   "val_dolma_cc": 3.953879643863526
+  },
+  "6": {
+   "val_dolma_cc": 3.929098875517303
+  },
+  "8": {
+   "val_dolma_cc": 3.9300679185345655
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 0.5200491318776743
+  },
+  "4": {
+   "val_dolma_cc": 0.06685451210904292
+  },
+  "5": {
+   "val_dolma_cc": 0.024710823854872555
+  },
+  "6": {
+   "val_dolma_cc": -6.994449135078185e-05
+  },
+  "8": {
+   "val_dolma_cc": 0.0008990985259118567
+  }
+ },
+ "own_grid_max_rel_weight_change": 2.2763482832033333e-07,
+ "own_grid_delta_loss": -6.994449135078185e-05,
+ "checks": {
+  "weights_restored": true,
+  "own_grid_weights_unchanged": true,
+  "own_grid_is_fixed_point": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": 6,
+  "a": 6,
+  "qkv": 6
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 3.7739331285905724,
+  "heldout_cc_en_tail_4M": 4.356792957085902
+ },
+ "loss_bf16_master": 3.7739331285905724,
+ "losses_eval_model": {
+  "val_dolma_cc": 3.8165860759648975
+ },
+ "bf16_minus_quant": -0.04265294737432512,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-m5-20",
+ "gpu": "NVIDIA L40S",
+ "t_s": 18.2
+}

pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/finalize.json

@@ -0,0 +1,102 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101",
+ "w_bits": 6,
+ "a_bits": 6,
+ "qkv_bits": 6,
+ "pre_norm": true,
+ "primary": "val_dolma_cc",
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "losses_master_plain": {
+  "val_dolma_cc": 3.816500349166907
+ },
+ "losses_fakequant_master": {
+  "val_dolma_cc": 3.8165860759648975
+ },
+ "losses_eval_model": {
+  "val_dolma_cc": 3.8165860759648975
+ },
+ "max_distinct_values_per_row": 62,
+ "rel_quant_noise_per_layer": {
+  "blocks.0.attention.w_q": 0.030932195484638214,
+  "blocks.0.attention.w_k": 0.031067954376339912,
+  "blocks.0.attention.w_v": 0.031908489763736725,
+  "blocks.0.attention.w_out": 0.03268344700336456,
+  "blocks.0.feed_forward.w1": 0.030635656788945198,
+  "blocks.0.feed_forward.w2": 0.03498230129480362,
+  "blocks.0.feed_forward.w3": 0.030625954270362854,
+  "blocks.1.attention.w_q": 0.030392387881875038,
+  "blocks.1.attention.w_k": 0.030469106510281563,
+  "blocks.1.attention.w_v": 0.03046298213303089,
+  "blocks.1.attention.w_out": 0.03174380585551262,
+  "blocks.1.feed_forward.w1": 0.030810555443167686,
+  "blocks.1.feed_forward.w2": 0.03493629768490791,
+  "blocks.1.feed_forward.w3": 0.030692772939801216,
+  "blocks.2.attention.w_q": 0.03051871620118618,
+  "blocks.2.attention.w_k": 0.030241983011364937,
+  "blocks.2.attention.w_v": 0.030938543379306793,
+  "blocks.2.attention.w_out": 0.033308736979961395,
+  "blocks.2.feed_forward.w1": 0.030609318986535072,
+  "blocks.2.feed_forward.w2": 0.03414329141378403,
+  "blocks.2.feed_forward.w3": 0.030738672241568565,
+  "blocks.3.attention.w_q": 0.030858557671308517,
+  "blocks.3.attention.w_k": 0.030566640198230743,
+  "blocks.3.attention.w_v": 0.030706897377967834,
+  "blocks.3.attention.w_out": 0.032133810222148895,
+  "blocks.3.feed_forward.w1": 0.030591895803809166,
+  "blocks.3.feed_forward.w2": 0.03722628951072693,
+  "blocks.3.feed_forward.w3": 0.030861755833029747,
+  "blocks.4.attention.w_q": 0.030701402574777603,
+  "blocks.4.attention.w_k": 0.03091822750866413,
+  "blocks.4.attention.w_v": 0.031039241701364517,
+  "blocks.4.attention.w_out": 0.032929953187704086,
+  "blocks.4.feed_forward.w1": 0.030669309198856354,
+  "blocks.4.feed_forward.w2": 0.034518878906965256,
+  "blocks.4.feed_forward.w3": 0.030812233686447144,
+  "blocks.5.attention.w_q": 0.030762795358896255,
+  "blocks.5.attention.w_k": 0.0308926273137331,
+  "blocks.5.attention.w_v": 0.030903514474630356,
+  "blocks.5.attention.w_out": 0.033717043697834015,
+  "blocks.5.feed_forward.w1": 0.03070686198771,
+  "blocks.5.feed_forward.w2": 0.03584417328238487,
+  "blocks.5.feed_forward.w3": 0.030798662453889847,
+  "blocks.6.attention.w_q": 0.03058581054210663,
+  "blocks.6.attention.w_k": 0.030940260738134384,
+  "blocks.6.attention.w_v": 0.031143827363848686,
+  "blocks.6.attention.w_out": 0.03557548299431801,
+  "blocks.6.feed_forward.w1": 0.030763618648052216,
+  "blocks.6.feed_forward.w2": 0.0363113172352314,
+  "blocks.6.feed_forward.w3": 0.030922498553991318,
+  "blocks.7.attention.w_q": 0.031031783670186996,
+  "blocks.7.attention.w_k": 0.031099854037165642,
+  "blocks.7.attention.w_v": 0.03084469400346279,
+  "blocks.7.attention.w_out": 0.033200811594724655,
+  "blocks.7.feed_forward.w1": 0.03116525150835514,
+  "blocks.7.feed_forward.w2": 0.041475873440504074,
+  "blocks.7.feed_forward.w3": 0.031072761863470078,
+  "blocks.8.attention.w_q": 0.03149639070034027,
+  "blocks.8.attention.w_k": 0.03132655471563339,
+  "blocks.8.attention.w_v": 0.03121422417461872,
+  "blocks.8.attention.w_out": 0.037734467536211014,
+  "blocks.8.feed_forward.w1": 0.031303565949201584,
+  "blocks.8.feed_forward.w2": 0.04681382700800896,
+  "blocks.8.feed_forward.w3": 0.03129962459206581
+ },
+ "rel_quant_noise_mean": 0.03219565781690772,
+ "n_quantized_params": 30670848,
+ "n_params": 82191872,
+ "param_dtypes": [
+  "torch.float32"
+ ],
+ "norm_gain_max_abs_dev_from_1": 0.8291723132133484,
+ "loss_eval_model": 3.8165860759648975,
+ "checks": {
+  "snap_reproduces_fakequant": true,
+  "grid_levels_ok": true,
+  "master_weights_fp32": true,
+  "norm_gains_trained": true,
+  "loss_finite": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/launch.yaml

@@ -0,0 +1,126 @@
+model:
+  d_model: 512
+  vocab_size: 50304
+  n_layers: 9
+  block:
+    attention:
+      name: default
+      n_heads: 8
+      bias: false
+      rope:
+        name: default
+        theta: 10000
+        full_precision: true
+        _CLASS_: olmo_core.nn.rope.RoPEConfig
+      backend: flash_2
+      dtype: float32
+      _CLASS_: olmo_core.nn.attention.AttentionConfig
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    feed_forward:
+      hidden_size: 1536
+      name: default
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.feed_forward.FeedForwardConfig
+    name: default
+    _CLASS_: olmo_core.nn.transformer.config.TransformerBlockConfig
+  lm_head:
+    name: default
+    layer_norm:
+      name: rms
+      eps: 1.0e-06
+      bias: false
+      dtype: float32
+      _CLASS_: olmo_core.nn.layer_norm.LayerNormConfig
+    bias: false
+    dtype: float32
+    loss_implementation: default
+    _CLASS_: olmo_core.nn.lm_head.LMHeadConfig
+  name: default
+  dtype: float32
+  init_method: normal
+  init_seed: 101
+  init_std: 0.02
+  _CLASS_: olmo_core.nn.transformer.config.TransformerConfig
+dataset:
+  _CLASS_: olmo_core.data.NumpyFSLDatasetConfig
+  tokenizer:
+    _CLASS_: olmo_core.data.TokenizerConfig
+    identifier: allenai/gpt-neox-olmo-dolma-v1_5
+    vocab_size: 50280
+    eos_token_id: 50279
+    pad_token_id: 1
+  sequence_length: 1024
+  work_dir: /data/user_data/jgai/precision-optimizer/work/dataset-cache
+  paths:
+  - /data/user_data/jgai/precision-perturb/dolma/cc_en_head/part-000-00000.npy
+data_loader:
+  _CLASS_: olmo_core.data.NumpyDataLoaderConfig
+  global_batch_size: 262144
+  seed: 0
+  num_workers: 4
+train_module_type: normal
+train_module:
+  _CLASS_: precq.train_module.FakeQuantTransformerTrainModuleConfig
+  fq_weight_bits: 6
+  fq_act_bits: 6
+  fq_qkv_bits: 6
+  fq_pre_norm: true
+  rank_microbatch_size: 32768
+  max_sequence_length: 1024
+  optim:
+    _CLASS_: precq.muon26.MuonConfig
+    lr: 0.005
+    weight_decay: 0.1
+    momentum: 0.95
+    nesterov: true
+    ns_steps: 5
+    eps: 1.0e-07
+    adjust_lr_fn: null
+    adamw_lr: 0.0024
+    adamw_betas:
+    - 0.9
+    - 0.95
+    adamw_eps: 1.0e-15
+    adamw_weight_decay: 0.1
+    adamw_fused: true
+    muon_on_lm_head: false
+  scheduler:
+    _CLASS_: olmo_core.optim.CosWithWarmup
+    warmup: 701
+    alpha_f: 0.1
+  compile_model: true
+  dp_config:
+    _CLASS_: olmo_core.train.train_module.transformer.config.TransformerDataParallelConfig
+    name: ddp
+    param_dtype: float32
+    reduce_dtype: float32
+  autocast_precision: bfloat16
+  max_grad_norm: 1.0
+trainer:
+  _CLASS_: olmo_core.train.config.TrainerConfig
+  save_folder: /data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/ckpt
+  save_overwrite: true
+  metrics_collect_interval: 10
+  cancel_check_interval: 5
+  callbacks:
+    config_saver:
+      _CLASS_: olmo_core.train.callbacks.config_saver.ConfigSaverCallback
+    checkpointer:
+      _CLASS_: olmo_core.train.callbacks.checkpointer.CheckpointerCallback
+      save_interval: 1007019
+      save_async: false
+wandb:
+  enabled: false
+init_seed: 101
+n_tokens: 1839988736
+load_path: null
+validation_datasets:
+  dolma_cc_val:
+  - /data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy
+validation_eval_interval: 500

pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/perturb_all.json

@@ -0,0 +1,284 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101",
+ "scope": "all",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.033884573727846146,
+  "blocks.0.attention.w_out.weight": 0.021545495837926865,
+  "blocks.0.attention.w_q.weight": 0.032501790672540665,
+  "blocks.0.attention.w_v.weight": 0.026910170912742615,
+  "blocks.0.attention_norm.weight": 0.3879252076148987,
+  "blocks.0.feed_forward.w1.weight": 0.02422112412750721,
+  "blocks.0.feed_forward.w2.weight": 0.008776023052632809,
+  "blocks.0.feed_forward.w3.weight": 0.024775676429271698,
+  "blocks.0.feed_forward_norm.weight": 0.4009745121002197,
+  "blocks.1.attention.w_k.weight": 0.027325544506311417,
+  "blocks.1.attention.w_out.weight": 0.017743699252605438,
+  "blocks.1.attention.w_q.weight": 0.026878027245402336,
+  "blocks.1.attention.w_v.weight": 0.02216923050582409,
+  "blocks.1.attention_norm.weight": 0.3890649080276489,
+  "blocks.1.feed_forward.w1.weight": 0.0247330479323864,
+  "blocks.1.feed_forward.w2.weight": 0.009560699574649334,
+  "blocks.1.feed_forward.w3.weight": 0.024706684052944183,
+  "blocks.1.feed_forward_norm.weight": 0.4063830077648163,
+  "blocks.2.attention.w_k.weight": 0.02852986566722393,
+  "blocks.2.attention.w_out.weight": 0.018322842195630074,
+  "blocks.2.attention.w_q.weight": 0.02715517021715641,
+  "blocks.2.attention.w_v.weight": 0.023582633584737778,
+  "blocks.2.attention_norm.weight": 0.3945653438568115,
+  "blocks.2.feed_forward.w1.weight": 0.02440229430794716,
+  "blocks.2.feed_forward.w2.weight": 0.010609530843794346,
+  "blocks.2.feed_forward.w3.weight": 0.024790987372398376,
+  "blocks.2.feed_forward_norm.weight": 0.40593409538269043,
+  "blocks.3.attention.w_k.weight": 0.029783764854073524,
+  "blocks.3.attention.w_out.weight": 0.02107260376214981,
+  "blocks.3.attention.w_q.weight": 0.027262361720204353,
+  "blocks.3.attention.w_v.weight": 0.024727392941713333,
+  "blocks.3.attention_norm.weight": 0.39631086587905884,
+  "blocks.3.feed_forward.w1.weight": 0.023838182911276817,
+  "blocks.3.feed_forward.w2.weight": 0.011221504770219326,
+  "blocks.3.feed_forward.w3.weight": 0.024804741144180298,
+  "blocks.3.feed_forward_norm.weight": 0.4052497446537018,
+  "blocks.4.attention.w_k.weight": 0.029391638934612274,
+  "blocks.4.attention.w_out.weight": 0.02368864417076111,
+  "blocks.4.attention.w_q.weight": 0.030474061146378517,
+  "blocks.4.attention.w_v.weight": 0.0250234492123127,
+  "blocks.4.attention_norm.weight": 0.40151700377464294,
+  "blocks.4.feed_forward.w1.weight": 0.023672061040997505,
+  "blocks.4.feed_forward.w2.weight": 0.012774117290973663,
+  "blocks.4.feed_forward.w3.weight": 0.024888446554541588,
+  "blocks.4.feed_forward_norm.weight": 0.40640202164649963,
+  "blocks.5.attention.w_k.weight": 0.0299989040941,
+  "blocks.5.attention.w_out.weight": 0.02527996338903904,
+  "blocks.5.attention.w_q.weight": 0.03006657026708126,
+  "blocks.5.attention.w_v.weight": 0.02713628113269806,
+  "blocks.5.attention_norm.weight": 0.40378764271736145,
+  "blocks.5.feed_forward.w1.weight": 0.023739339783787727,
+  "blocks.5.feed_forward.w2.weight": 0.01463079173117876,
+  "blocks.5.feed_forward.w3.weight": 0.025103239342570305,
+  "blocks.5.feed_forward_norm.weight": 0.4088801145553589,
+  "blocks.6.attention.w_k.weight": 0.029504111036658287,
+  "blocks.6.attention.w_out.weight": 0.02975505404174328,
+  "blocks.6.attention.w_q.weight": 0.02926630713045597,
+  "blocks.6.attention.w_v.weight": 0.027235006913542747,
+  "blocks.6.attention_norm.weight": 0.40453341603279114,
+  "blocks.6.feed_forward.w1.weight": 0.02357933484017849,
+  "blocks.6.feed_forward.w2.weight": 0.016710400581359863,
+  "blocks.6.feed_forward.w3.weight": 0.024904726073145866,
+  "blocks.6.feed_forward_norm.weight": 0.4096214175224304,
+  "blocks.7.attention.w_k.weight": 0.030117973685264587,
+  "blocks.7.attention.w_out.weight": 0.03328891843557358,
+  "blocks.7.attention.w_q.weight": 0.028999419882893562,
+  "blocks.7.attention.w_v.weight": 0.025352291762828827,
+  "blocks.7.attention_norm.weight": 0.40566807985305786,
+  "blocks.7.feed_forward.w1.weight": 0.022943098098039627,
+  "blocks.7.feed_forward.w2.weight": 0.019236372783780098,
+  "blocks.7.feed_forward.w3.weight": 0.024431584402918816,
+  "blocks.7.feed_forward_norm.weight": 0.41024506092071533,
+  "blocks.8.attention.w_k.weight": 0.028775880113244057,
+  "blocks.8.attention.w_out.weight": 0.039014291018247604,
+  "blocks.8.attention.w_q.weight": 0.026884691789746284,
+  "blocks.8.attention.w_v.weight": 0.02573145553469658,
+  "blocks.8.attention_norm.weight": 0.409767210483551,
+  "blocks.8.feed_forward.w1.weight": 0.02186930924654007,
+  "blocks.8.feed_forward.w2.weight": 0.023152999579906464,
+  "blocks.8.feed_forward.w3.weight": 0.02385779283940792,
+  "blocks.8.feed_forward_norm.weight": 0.4094424843788147,
+  "embeddings.weight": 0.07002507150173187,
+  "lm_head.norm.weight": 1.3853260278701782,
+  "lm_head.w_out.weight": 0.07471665740013123
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 84,
+ "n_perturbed_params": 82191872,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 3.8165860759648975
+ },
+ "base_loss": 3.8165860759648975,
+ "loss": {
+  "0.005": [
+   3.8172678371833086,
+   3.8171116142565764,
+   3.8159586725472936
+  ],
+  "0.01": [
+   3.8176468903175227,
+   3.818043408056804,
+   3.8168794922845257
+  ],
+  "0.02": [
+   3.8206880907293117,
+   3.8205451970833733,
+   3.820196484005677
+  ],
+  "0.05": [
+   3.843052730337484,
+   3.843897525835096,
+   3.8416292241809287
+  ],
+  "0.1": [
+   3.9249599206092856,
+   3.929138217826567,
+   3.921051436673467
+  ],
+  "0.15": [
+   4.079257128156649,
+   4.096721040742732,
+   4.074365094313875
+  ],
+  "0.2": [
+   4.34255490302199,
+   4.388319989029223,
+   4.333827130957245
+  ],
+  "0.3": [
+   5.484848179687375,
+   5.656619338838878,
+   5.408889064578226
+  ],
+  "0.4": [
+   7.629972741131737,
+   7.77432493980658,
+   7.303515376176326
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.8172678371833086,
+    3.8171116142565764,
+    3.8159586725472936
+   ],
+   "0.01": [
+    3.8176468903175227,
+    3.818043408056804,
+    3.8168794922845257
+   ],
+   "0.02": [
+    3.8206880907293117,
+    3.8205451970833733,
+    3.820196484005677
+   ],
+   "0.05": [
+    3.843052730337484,
+    3.843897525835096,
+    3.8416292241809287
+   ],
+   "0.1": [
+    3.9249599206092856,
+    3.929138217826567,
+    3.921051436673467
+   ],
+   "0.15": [
+    4.079257128156649,
+    4.096721040742732,
+    4.074365094313875
+   ],
+   "0.2": [
+    4.34255490302199,
+    4.388319989029223,
+    4.333827130957245
+   ],
+   "0.3": [
+    5.484848179687375,
+    5.656619338838878,
+    5.408889064578226
+   ],
+   "0.4": [
+    7.629972741131737,
+    7.77432493980658,
+    7.303515376176326
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0001932986974959583,
+   "0.01": 0.0009371875880535008,
+   "0.02": 0.003890514641223586,
+   "0.05": 0.026273750819605368,
+   "0.1": 0.10846378240487553,
+   "0.15": 0.266861678439521,
+   "0.2": 0.5383145983712554,
+   "0.3": 1.7001994517365953,
+   "0.4": 3.7526849430733167
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9965,
+   1.0028
+  ],
+  "0.01": [
+   0.9965,
+   1.0028
+  ],
+  "0.02": [
+   0.9965,
+   1.0028
+  ],
+  "0.05": [
+   0.9965,
+   1.0028
+  ],
+  "0.1": [
+   0.9965,
+   1.0028
+  ],
+  "0.15": [
+   0.9965,
+   1.0028
+  ],
+  "0.2": [
+   0.9965,
+   1.0028
+  ],
+  "0.3": [
+   0.9965,
+   1.0028
+  ],
+  "0.4": [
+   0.9965,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0001932986974959583,
+  "0.01": 0.0009371875880535008,
+  "0.02": 0.003890514641223586,
+  "0.05": 0.026273750819605368,
+  "0.1": 0.10846378240487553,
+  "0.15": 0.266861678439521,
+  "0.2": 0.5383145983712554,
+  "0.3": 1.7001994517365953,
+  "0.4": 3.7526849430733167
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 79.2
+}

pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/perturb_blocks.json

@@ -0,0 +1,263 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101",
+ "scope": "blocks",
+ "noise_ref": null,
+ "noise_scale": "rms",
+ "noise_std_per_tensor_over_gamma": {
+  "blocks.0.attention.w_k.weight": 0.033884573727846146,
+  "blocks.0.attention.w_out.weight": 0.021545495837926865,
+  "blocks.0.attention.w_q.weight": 0.032501790672540665,
+  "blocks.0.attention.w_v.weight": 0.026910170912742615,
+  "blocks.0.feed_forward.w1.weight": 0.02422112412750721,
+  "blocks.0.feed_forward.w2.weight": 0.008776023052632809,
+  "blocks.0.feed_forward.w3.weight": 0.024775676429271698,
+  "blocks.1.attention.w_k.weight": 0.027325544506311417,
+  "blocks.1.attention.w_out.weight": 0.017743699252605438,
+  "blocks.1.attention.w_q.weight": 0.026878027245402336,
+  "blocks.1.attention.w_v.weight": 0.02216923050582409,
+  "blocks.1.feed_forward.w1.weight": 0.0247330479323864,
+  "blocks.1.feed_forward.w2.weight": 0.009560699574649334,
+  "blocks.1.feed_forward.w3.weight": 0.024706684052944183,
+  "blocks.2.attention.w_k.weight": 0.02852986566722393,
+  "blocks.2.attention.w_out.weight": 0.018322842195630074,
+  "blocks.2.attention.w_q.weight": 0.02715517021715641,
+  "blocks.2.attention.w_v.weight": 0.023582633584737778,
+  "blocks.2.feed_forward.w1.weight": 0.02440229430794716,
+  "blocks.2.feed_forward.w2.weight": 0.010609530843794346,
+  "blocks.2.feed_forward.w3.weight": 0.024790987372398376,
+  "blocks.3.attention.w_k.weight": 0.029783764854073524,
+  "blocks.3.attention.w_out.weight": 0.02107260376214981,
+  "blocks.3.attention.w_q.weight": 0.027262361720204353,
+  "blocks.3.attention.w_v.weight": 0.024727392941713333,
+  "blocks.3.feed_forward.w1.weight": 0.023838182911276817,
+  "blocks.3.feed_forward.w2.weight": 0.011221504770219326,
+  "blocks.3.feed_forward.w3.weight": 0.024804741144180298,
+  "blocks.4.attention.w_k.weight": 0.029391638934612274,
+  "blocks.4.attention.w_out.weight": 0.02368864417076111,
+  "blocks.4.attention.w_q.weight": 0.030474061146378517,
+  "blocks.4.attention.w_v.weight": 0.0250234492123127,
+  "blocks.4.feed_forward.w1.weight": 0.023672061040997505,
+  "blocks.4.feed_forward.w2.weight": 0.012774117290973663,
+  "blocks.4.feed_forward.w3.weight": 0.024888446554541588,
+  "blocks.5.attention.w_k.weight": 0.0299989040941,
+  "blocks.5.attention.w_out.weight": 0.02527996338903904,
+  "blocks.5.attention.w_q.weight": 0.03006657026708126,
+  "blocks.5.attention.w_v.weight": 0.02713628113269806,
+  "blocks.5.feed_forward.w1.weight": 0.023739339783787727,
+  "blocks.5.feed_forward.w2.weight": 0.01463079173117876,
+  "blocks.5.feed_forward.w3.weight": 0.025103239342570305,
+  "blocks.6.attention.w_k.weight": 0.029504111036658287,
+  "blocks.6.attention.w_out.weight": 0.02975505404174328,
+  "blocks.6.attention.w_q.weight": 0.02926630713045597,
+  "blocks.6.attention.w_v.weight": 0.027235006913542747,
+  "blocks.6.feed_forward.w1.weight": 0.02357933484017849,
+  "blocks.6.feed_forward.w2.weight": 0.016710400581359863,
+  "blocks.6.feed_forward.w3.weight": 0.024904726073145866,
+  "blocks.7.attention.w_k.weight": 0.030117973685264587,
+  "blocks.7.attention.w_out.weight": 0.03328891843557358,
+  "blocks.7.attention.w_q.weight": 0.028999419882893562,
+  "blocks.7.attention.w_v.weight": 0.025352291762828827,
+  "blocks.7.feed_forward.w1.weight": 0.022943098098039627,
+  "blocks.7.feed_forward.w2.weight": 0.019236372783780098,
+  "blocks.7.feed_forward.w3.weight": 0.024431584402918816,
+  "blocks.8.attention.w_k.weight": 0.028775880113244057,
+  "blocks.8.attention.w_out.weight": 0.039014291018247604,
+  "blocks.8.attention.w_q.weight": 0.026884691789746284,
+  "blocks.8.attention.w_v.weight": 0.02573145553469658,
+  "blocks.8.feed_forward.w1.weight": 0.02186930924654007,
+  "blocks.8.feed_forward.w2.weight": 0.023152999579906464,
+  "blocks.8.feed_forward.w3.weight": 0.02385779283940792
+ },
+ "gammas": [
+  0.005,
+  0.01,
+  0.02,
+  0.05,
+  0.1,
+  0.15,
+  0.2,
+  0.3,
+  0.4
+ ],
+ "seeds": [
+  0,
+  1,
+  2
+ ],
+ "n_tensors": 63,
+ "n_perturbed_params": 30670848,
+ "eval_tokens": {
+  "val_dolma_cc": 499712
+ },
+ "primary": "val_dolma_cc",
+ "w_bits": 6,
+ "base_losses": {
+  "val_dolma_cc": 3.8165860759648975
+ },
+ "base_loss": 3.8165860759648975,
+ "loss": {
+  "0.005": [
+   3.817336541468609,
+   3.817538655931506,
+   3.8168800165356633
+  ],
+  "0.01": [
+   3.8171605613161628,
+   3.816650633188959,
+   3.8168201697471975
+  ],
+  "0.02": [
+   3.8179068758463135,
+   3.8178052141611283,
+   3.818238981117444
+  ],
+  "0.05": [
+   3.8249593447154986,
+   3.8250348603532682,
+   3.8262709545527156
+  ],
+  "0.1": [
+   3.8528304287629402,
+   3.852761274560618,
+   3.854550848854923
+  ],
+  "0.15": [
+   3.9037916220974953,
+   3.902684196341472,
+   3.9077033297865786
+  ],
+  "0.2": [
+   3.9867743644386087,
+   3.98497371827326,
+   3.987170623964643
+  ],
+  "0.3": [
+   4.314607317431654,
+   4.309899569600771,
+   4.3006122689914745
+  ],
+  "0.4": [
+   5.02026613272549,
+   5.055846697512039,
+   4.985979411922304
+  ]
+ },
+ "loss_by_set": {
+  "val_dolma_cc": {
+   "0.005": [
+    3.817336541468609,
+    3.817538655931506,
+    3.8168800165356633
+   ],
+   "0.01": [
+    3.8171605613161628,
+    3.816650633188959,
+    3.8168201697471975
+   ],
+   "0.02": [
+    3.8179068758463135,
+    3.8178052141611283,
+    3.818238981117444
+   ],
+   "0.05": [
+    3.8249593447154986,
+    3.8250348603532682,
+    3.8262709545527156
+   ],
+   "0.1": [
+    3.8528304287629402,
+    3.852761274560618,
+    3.854550848854923
+   ],
+   "0.15": [
+    3.9037916220974953,
+    3.902684196341472,
+    3.9077033297865786
+   ],
+   "0.2": [
+    3.9867743644386087,
+    3.98497371827326,
+    3.987170623964643
+   ],
+   "0.3": [
+    4.314607317431654,
+    4.309899569600771,
+    4.3006122689914745
+   ],
+   "0.4": [
+    5.02026613272549,
+    5.055846697512039,
+    4.985979411922304
+   ]
+  }
+ },
+ "mean_delta_by_set": {
+  "val_dolma_cc": {
+   "0.005": 0.0006656620136951119,
+   "0.01": 0.00029104545254243774,
+   "0.02": 0.0013976144100644738,
+   "0.05": 0.008835643908929658,
+   "0.1": 0.03679477476126225,
+   "0.15": 0.08814030677695106,
+   "0.2": 0.1697201595939397,
+   "0.3": 0.49178697604306887,
+   "0.4": 1.2041113380883797
+  }
+ },
+ "rel_noise_over_gamma": {
+  "0.005": [
+   0.9966,
+   1.0028
+  ],
+  "0.01": [
+   0.9966,
+   1.0028
+  ],
+  "0.02": [
+   0.9966,
+   1.0028
+  ],
+  "0.05": [
+   0.9966,
+   1.0028
+  ],
+  "0.1": [
+   0.9966,
+   1.0028
+  ],
+  "0.15": [
+   0.9966,
+   1.0028
+  ],
+  "0.2": [
+   0.9966,
+   1.0028
+  ],
+  "0.3": [
+   0.9966,
+   1.0028
+  ],
+  "0.4": [
+   0.9966,
+   1.0028
+  ]
+ },
+ "mean_delta": {
+  "0.005": 0.0006656620136951119,
+  "0.01": 0.00029104545254243774,
+  "0.02": 0.0013976144100644738,
+  "0.05": 0.008835643908929658,
+  "0.1": 0.03679477476126225,
+  "0.15": 0.08814030677695106,
+  "0.2": 0.1697201595939397,
+  "0.3": 0.49178697604306887,
+  "0.4": 1.2041113380883797
+ },
+ "checks": {
+  "base_loss_finite": true,
+  "noise_scale_matches_gamma": true,
+  "base_restored": true
+ },
+ "ok": true,
+ "t_total_s": 71.7
+}

pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101/ptq.json

@@ -0,0 +1,52 @@
+{
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6-muon-lr5.0e-03-alr2.4e-03-s101",
+ "w_bits": 6,
+ "method": "RTN, symmetric absmax per output channel",
+ "scope": "transformer-block Linear weights",
+ "primary": "val_dolma_cc",
+ "base_losses": {
+  "val_dolma_cc": 3.8165860759648975
+ },
+ "losses_by_bits": {
+  "3": {
+   "val_dolma_cc": 4.549142319661365
+  },
+  "4": {
+   "val_dolma_cc": 3.9001482527445095
+  },
+  "5": {
+   "val_dolma_cc": 3.8357712028443145
+  },
+  "6": {
+   "val_dolma_cc": 3.8165860759648975
+  },
+  "8": {
+   "val_dolma_cc": 3.8169317139575623
+  }
+ },
+ "delta_by_bits": {
+  "3": {
+   "val_dolma_cc": 0.7325562436964672
+  },
+  "4": {
+   "val_dolma_cc": 0.08356217677961197
+  },
+  "5": {
+   "val_dolma_cc": 0.01918512687941698
+  },
+  "6": {
+   "val_dolma_cc": 0.0
+  },
+  "8": {
+   "val_dolma_cc": 0.00034563799266473794
+  }
+ },
+ "own_grid_max_rel_weight_change": 2.2763893525734602e-07,
+ "own_grid_delta_loss": 0.0,
+ "checks": {
+  "weights_restored": true,
+  "own_grid_weights_unchanged": true,
+  "own_grid_is_fixed_point": true
+ },
+ "ok": true
+}

pom-30M-D1.84B-f6B-muon-lr5.0e-03-alr2.4e-03-s101/bf16_start.json

@@ -0,0 +1,40 @@
+{
+ "run": "pom-30M-D1.84B-f6B-muon-lr5.0e-03-alr2.4e-03-s101",
+ "run_dir": "/data/user_data/jgai/precision-optimizer/runs/pom-30M-D1.84B-f6B-muon-lr5.0e-03-alr2.4e-03-s101",
+ "weights": "final-unsharded (fp32 master)",
+ "quantization": "none (weights, block-Linear inputs, q/k/v); pre-norm kept; BF16 autocast; uncompiled",
+ "run_bits": {
+  "w": 6,
+  "a": 6,
+  "qkv": 6
+ },
+ "pre_norm": true,
+ "val_files": [
+  "/data/user_data/jgai/precision-perturb/dolma/val_dolma_cc.npy",
+  "/data/user_data/jgai/precision-perturb/dolma/heldout_cc_en_tail_4M.npy"
+ ],
+ "n_inst": null,
+ "batch": 8,
+ "primary": "val_dolma_cc",
+ "n_seqs": {
+  "val_dolma_cc": 488,
+  "heldout_cc_en_tail_4M": 4096
+ },
+ "losses_bf16_master": {
+  "val_dolma_cc": 4.706428408139432,
+  "heldout_cc_en_tail_4M": 5.2704657109951345
+ },
+ "loss_bf16_master": 4.706428408139432,
+ "losses_eval_model": {
+  "val_dolma_cc": 4.668224568905942
+ },
+ "bf16_minus_quant": 0.03820383923349002,
+ "checks": {
+  "master_weights_fp32": true,
+  "loss_finite": true
+ },
+ "ok": true,
+ "host": "babel-o9-28",
+ "gpu": "NVIDIA L40S",
+ "t_s": 18.3
+}