Upload folder using huggingface_hub

150M_228M_15e4_10B_full/blumodel_muon_step00015248.pt

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

150M_228M_15e4_10B_full/blumodel_muon_step00016201.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84555bc64bcef2702981e736ac313fedba50cba6bd469c7c64b4d5ca8d3c3add
+size 2515184199

150M_228M_15e4_10B_full/blumodel_muon_step00017154.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80b9c23d170474babe961dc098c1baf1baff746c14cbf76b956c4b0a0d225ef7
+size 2515184199

150M_228M_15e4_10B_full/blumodel_muon_step00018107.pt

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

150M_228M_15e4_10B_full/blumodel_muon_step00019060.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cc676ec5c64b3020dabf764d682abbb4820bedd70460f216389cd04d57fa464
+size 2515184199

150M_228M_15e4_10B_full/blumodel_muon_step00019074.pt

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

150M_228M_3e3_6B_incmp.csv

@@ -0,0 +1,398 @@
+step,train_loss,val_loss,tok_per_sec,grad_norm,dt_ms,lr,elapsed_min,muon_lr
+10309,2.646035,3.057884,241083.302612,0.080317,2174.717180,0.002999,0.036246,0.002999
+10310,2.581861,3.057884,307079.635579,0.083498,1707.335620,0.002999,0.067322,0.002999
+10311,2.633330,3.057884,301997.020155,0.080284,1736.070110,0.002999,0.098820,0.002999
+10312,2.711791,3.057884,317031.979656,0.094343,1653.738530,0.002999,0.126388,0.002999
+10313,2.672399,3.057884,319283.029922,0.100907,1642.079130,0.002999,0.153763,0.002999
+10314,2.651201,3.057884,319127.757069,0.089898,1642.878090,0.002999,0.181149,0.002999
+10315,2.876376,3.057884,319772.240348,0.094825,1639.566960,0.002999,0.208485,0.002999
+10316,2.645422,3.057884,319992.920078,0.098303,1638.436250,0.002999,0.235806,0.002999
+10317,2.569267,3.057884,320027.812964,0.099491,1638.257610,0.002998,0.263117,0.002998
+10318,2.526511,3.057884,319971.904810,0.095968,1638.543860,0.002998,0.290432,0.002998
+10319,2.734799,3.057884,319961.467532,0.085967,1638.597310,0.002998,0.317747,0.002998
+10320,2.614316,3.057884,319905.998325,0.075867,1638.881430,0.002998,0.345069,0.002998
+10321,2.614264,3.057884,319885.991804,0.073269,1638.983930,0.002998,0.372391,0.002998
+10322,2.737565,3.057884,319925.286979,0.072798,1638.782620,0.002997,0.399712,0.002997
+10323,2.688076,3.057884,319928.158712,0.072226,1638.767910,0.002997,0.427032,0.002997
+10324,2.684522,3.057884,319891.013694,0.076770,1638.958200,0.002997,0.454354,0.002997
+10325,2.580991,3.057884,319816.716366,0.078332,1639.338950,0.002997,0.481682,0.002997
+10326,2.612901,3.057884,319771.171562,0.082629,1639.572440,0.002996,0.509014,0.002996
+10327,2.681869,3.057884,319822.838371,0.085332,1639.307570,0.002996,0.536342,0.002996
+10328,2.643904,3.057884,319831.128226,0.080611,1639.265080,0.002996,0.563668,0.002996
+10329,2.624868,3.057884,319825.544380,0.077997,1639.293700,0.002995,0.590995,0.002995
+10330,2.695429,3.057884,319867.236745,0.077656,1639.080030,0.002995,0.618318,0.002995
+10331,2.638742,3.057884,319782.458536,0.077614,1639.514570,0.002995,0.645650,0.002995
+10332,2.680027,3.057884,319788.111099,0.080066,1639.485590,0.002994,0.672980,0.002994
+10333,2.669517,3.057884,319659.210973,0.087468,1640.146700,0.002994,0.700322,0.002994
+10334,2.674735,3.057884,319674.298638,0.086068,1640.069290,0.002994,0.727663,0.002994
+10335,2.683662,3.057884,319687.947178,0.085377,1639.999270,0.002993,0.755002,0.002993
+10336,2.588056,3.057884,319702.214870,0.082420,1639.926080,0.002993,0.782341,0.002993
+10337,2.592314,3.057884,319685.405289,0.078556,1640.012310,0.002993,0.809681,0.002993
+10338,2.659404,3.057884,319722.970505,0.092736,1639.819620,0.002992,0.837018,0.002992
+10339,2.554655,3.057884,319656.926799,0.085486,1640.158420,0.002992,0.864361,0.002992
+10340,2.648113,3.057884,319627.440118,0.088174,1640.309730,0.002991,0.891706,0.002991
+10341,2.685987,3.057884,319721.757769,0.089004,1639.825840,0.002991,0.919042,0.002991
+10342,2.685290,3.057884,319669.213385,0.078556,1640.095380,0.002991,0.946384,0.002991
+10343,2.607982,3.057884,319640.492238,0.085381,1640.242750,0.002990,0.973728,0.002990
+10344,2.616746,3.057884,319631.695873,0.108711,1640.287890,0.002990,1.001073,0.002990
+10345,2.622490,3.057884,319681.003860,0.108881,1640.034890,0.002989,1.028413,0.002989
+10346,2.655280,3.057884,319721.030522,0.097260,1639.829570,0.002989,1.055750,0.002989
+10347,2.706025,3.057884,319694.019417,0.090086,1639.968120,0.002988,1.083090,0.002988
+10348,2.582737,3.057884,319620.970964,0.082688,1640.342930,0.002988,1.110435,0.002988
+10349,2.706316,3.057884,319644.491100,0.087002,1640.222230,0.002987,1.137779,0.002987
+10350,2.585762,3.057884,319714.249535,0.088104,1639.864350,0.002987,1.165116,0.002987
+10351,2.609111,3.057884,319646.500313,0.120095,1640.211920,0.002986,1.192459,0.002986
+10352,2.547157,3.057884,319653.208266,0.107101,1640.177500,0.002986,1.219802,0.002986
+10353,2.709230,3.057884,319682.259170,0.098309,1640.028450,0.002985,1.247142,0.002985
+10354,2.544000,3.057884,319639.428230,0.091700,1640.248210,0.002985,1.274486,0.002985
+10355,2.705685,3.057884,319610.174684,0.088208,1640.398340,0.002984,1.301833,0.002984
+10356,2.560872,3.057884,319621.327540,0.093999,1640.341100,0.002983,1.329178,0.002983
+10357,2.638376,3.057884,319630.614387,0.102006,1640.293440,0.002983,1.356523,0.002983
+10358,2.589150,3.057884,319641.821283,0.092981,1640.235930,0.002982,1.383867,0.002982
+10359,2.681022,3.057884,319590.472041,0.094087,1640.499470,0.002982,1.411215,0.002982
+10360,2.679786,3.057884,319579.416790,0.114467,1640.556220,0.002981,1.438564,0.002981
+10361,2.638465,3.057884,319607.172276,0.106257,1640.413750,0.002980,1.465911,0.002980
+10362,2.634351,3.057884,319598.297868,0.096087,1640.459300,0.002980,1.493259,0.002980
+10363,2.571059,3.057884,319635.369084,0.101446,1640.269040,0.002979,1.520603,0.002979
+10364,2.589033,3.057884,319598.290074,0.089744,1640.459340,0.002979,1.547951,0.002979
+10365,2.617486,3.057884,319631.425014,0.091137,1640.289280,0.002978,1.575295,0.002978
+10366,2.661411,3.057884,319602.293726,0.108345,1640.438790,0.002977,1.602643,0.002977
+10367,2.677341,3.057884,319611.776253,0.109394,1640.390120,0.002977,1.629989,0.002977
+10368,2.596525,3.057884,319608.423110,0.095481,1640.407330,0.002976,1.657335,0.002976
+10369,2.566483,3.057884,319630.871604,0.097378,1640.292120,0.002975,1.684680,0.002975
+10370,2.653575,3.057884,319718.375023,0.107972,1639.843190,0.002974,1.712017,0.002974
+10371,2.579601,3.057884,319624.476355,0.103648,1640.324940,0.002974,1.739362,0.002974
+10372,2.605006,3.057884,319593.943645,0.101006,1640.481650,0.002973,1.766710,0.002973
+10373,2.677480,3.057884,319613.576575,0.109509,1640.380880,0.002972,1.794056,0.002972
+10374,2.627289,3.057884,319651.382163,0.112040,1640.186870,0.002971,1.821399,0.002971
+10375,2.654639,3.057884,319594.356658,0.126767,1640.479530,0.002971,1.848747,0.002971
+10376,2.643901,3.057884,319604.908330,0.103787,1640.425370,0.002970,1.876094,0.002970
+10377,2.653200,3.057884,319623.595617,0.103036,1640.329460,0.002969,1.903439,0.002969
+10378,2.674783,3.057884,319568.973935,0.110426,1640.609830,0.002968,1.930789,0.002968
+10379,2.662562,3.057884,319691.087560,0.086865,1639.983160,0.002968,1.958128,0.002968
+10380,2.763343,3.057884,319639.800438,0.076092,1640.246300,0.002967,1.985472,0.002967
+10381,2.548690,3.057884,319660.791591,0.078707,1640.138590,0.002966,2.012814,0.002966
+10382,2.616013,3.057884,319655.710665,0.079765,1640.164660,0.002965,2.040157,0.002965
+10383,2.728987,3.057884,319608.460130,0.071894,1640.407140,0.002964,2.067504,0.002964
+10384,2.623497,3.057884,319543.809476,0.078327,1640.739030,0.002963,2.094856,0.002963
+10385,2.784320,3.057884,319582.654389,0.093729,1640.539600,0.002963,2.122205,0.002963
+10386,2.754317,3.057884,319643.581019,0.100288,1640.226900,0.002962,2.149548,0.002962
+10387,2.556110,3.057884,319645.560988,0.077960,1640.216740,0.002961,2.176892,0.002961
+10388,2.644885,3.057884,319624.544554,0.075712,1640.324590,0.002960,2.204237,0.002960
+10389,2.512394,3.057884,319626.980254,0.080357,1640.312090,0.002959,2.231583,0.002959
+10390,2.831544,3.057884,319677.107390,0.192073,1640.054880,0.002958,2.258923,0.002958
+10391,2.616250,3.057884,319653.508396,0.093133,1640.175960,0.002957,2.286267,0.002957
+10392,2.559199,3.057884,319623.663813,0.078946,1640.329110,0.002956,2.313612,0.002956
+10393,2.716237,3.057884,319537.655303,0.077445,1640.770630,0.002955,2.340965,0.002955
+10394,2.568694,3.057884,319631.333427,0.066505,1640.289750,0.002954,2.368311,0.002954
+10395,2.644906,3.057884,319559.509529,0.073140,1640.658420,0.002953,2.395661,0.002953
+10396,2.734189,3.057884,319635.448982,0.083121,1640.268630,0.002952,2.423006,0.002952
+10397,2.757720,3.057884,319688.598251,0.095128,1639.995930,0.002951,2.450345,0.002951
+10398,2.773420,3.057884,319636.783831,0.091602,1640.261780,0.002950,2.477690,0.002950
+10399,2.608002,3.057884,319619.895395,0.099007,1640.348450,0.002949,2.505035,0.002949
+10400,2.631424,3.057884,319687.974469,0.095844,1639.999130,0.002948,2.532375,0.002948
+10401,2.836423,3.057884,319619.513491,0.093251,1640.350410,0.002947,2.559721,0.002947
+10402,2.649410,3.057884,319634.525311,0.086392,1640.273370,0.002946,2.587065,0.002946
+10403,2.616033,3.057884,319566.994915,0.094652,1640.619990,0.002945,2.614415,0.002945
+10404,2.648589,3.057884,319668.293416,0.112149,1640.100100,0.002944,2.641757,0.002944
+10405,2.707739,3.057884,319613.888320,0.106043,1640.379280,0.002943,2.669103,0.002943
+10406,2.766337,3.057884,319599.421997,0.096687,1640.453530,0.002942,2.696451,0.002942
+10407,2.616031,3.057884,319633.936815,0.105159,1640.276390,0.002941,2.723796,0.002941
+10408,2.641221,3.057884,319665.843445,0.100178,1640.112670,0.002940,2.751137,0.002940
+10409,2.579833,3.057884,319622.420654,0.101381,1640.335490,0.002939,2.778483,0.002939
+10410,2.570470,3.057884,319693.637337,0.100985,1639.970080,0.002938,2.805822,0.002938
+10411,2.580414,3.057884,319600.620168,0.100138,1640.447380,0.002937,2.833169,0.002937
+10412,2.641308,3.057884,319647.969725,0.116407,1640.204380,0.002936,2.860513,0.002936
+10413,2.606977,3.057884,319660.251723,0.126621,1640.141360,0.002934,2.887855,0.002934
+10414,2.456744,3.057884,319618.100848,0.115935,1640.357660,0.002933,2.915200,0.002933
+10415,2.603387,3.057884,319651.119065,0.078222,1640.188220,0.002932,2.942543,0.002932
+10416,2.625212,3.057884,319620.815084,0.077804,1640.343730,0.002931,2.969889,0.002931
+10417,2.694703,3.057884,319580.899222,0.084717,1640.548610,0.002930,2.997238,0.002930
+10418,2.742951,3.057884,319627.492730,0.077662,1640.309460,0.002929,3.024583,0.002929
+10419,2.844461,3.057884,319640.293468,0.211969,1640.243770,0.002927,3.051927,0.002927
+10420,2.689941,3.057884,319667.112281,0.073109,1640.106160,0.002926,3.079268,0.002926
+10421,2.495415,3.057884,319628.151351,0.067772,1640.306080,0.002925,3.106613,0.002925
+10422,2.551864,3.057884,319608.697828,0.074357,1640.405920,0.002924,3.133960,0.002924
+10423,2.738378,3.057884,319573.931333,0.072473,1640.584380,0.002922,3.161310,0.002922
+10424,2.625021,3.057884,319587.374545,0.074038,1640.515370,0.002921,3.188659,0.002921
+10425,2.863070,3.057884,319577.536987,0.115944,1640.565870,0.002920,3.216008,0.002920
+10426,2.620465,3.057884,319627.687587,0.093652,1640.308460,0.002919,3.243353,0.002919
+10427,2.561589,3.057884,319630.690382,0.080830,1640.293050,0.002917,3.270698,0.002917
+10428,2.683123,3.057884,319650.974848,0.079047,1640.188960,0.002916,3.298040,0.002916
+10429,2.621562,3.057884,319594.691745,0.077472,1640.477810,0.002915,3.325388,0.002915
+10430,2.797495,3.057884,319611.090422,0.074964,1640.393640,0.002914,3.352734,0.002914
+10431,2.756585,3.057884,319625.168089,0.070199,1640.321390,0.002912,3.380080,0.002912
+10432,2.767882,3.057884,319570.760138,0.077297,1640.600660,0.002911,3.407431,0.002911
+10433,2.597656,3.057884,319678.487420,0.076636,1640.047800,0.002910,3.434771,0.002910
+10434,2.596557,3.057884,319609.054376,0.074409,1640.404090,0.002908,3.462118,0.002908
+10435,2.518670,3.057884,319629.762841,0.076837,1640.297810,0.002907,3.489463,0.002907
+10436,2.619940,3.057884,319620.937839,0.087342,1640.343100,0.002906,3.516808,0.002906
+10437,2.614896,3.057884,319592.459145,0.104175,1640.489270,0.002904,3.544156,0.002904
+10438,2.579712,3.057884,319627.859062,0.107138,1640.307580,0.002903,3.571501,0.002903
+10439,2.621733,3.057884,319636.663013,0.101048,1640.262400,0.002901,3.598845,0.002901
+10440,2.503057,3.057884,319646.225531,0.085263,1640.213330,0.002900,3.626188,0.002900
+10441,2.580845,3.057884,319627.319305,0.078154,1640.310350,0.002899,3.653533,0.002899
+10442,2.698635,3.057884,319575.758498,0.069221,1640.575000,0.002897,3.680882,0.002897
+10443,2.571500,3.057884,319636.867624,0.074886,1640.261350,0.002896,3.708227,0.002896
+10444,2.716341,3.057884,319638.179102,0.078948,1640.254620,0.002894,3.735571,0.002894
+10445,2.641122,3.057884,319578.472017,0.085113,1640.561070,0.002893,3.762920,0.002893
+10446,2.715120,3.057884,319607.320350,0.093111,1640.412990,0.002891,3.790266,0.002891
+10447,2.604709,3.057884,319638.744229,0.091830,1640.251720,0.002890,3.817610,0.002890
+10448,2.611566,3.057884,319654.297702,0.088400,1640.171910,0.002889,3.844953,0.002889
+10449,2.691082,3.057884,319607.043686,0.103328,1640.414410,0.002887,3.872300,0.002887
+10450,2.696886,3.057884,319611.068989,0.094205,1640.393750,0.002886,3.899647,0.002886
+10451,2.672131,3.057884,319666.001318,0.080755,1640.111860,0.002884,3.926989,0.002884
+10452,2.647804,3.057884,319612.109427,0.080559,1640.388410,0.002883,3.954335,0.002883
+10453,2.719227,3.057884,319645.071840,0.091364,1640.219250,0.002881,3.981679,0.002881
+10454,2.508566,3.057884,319658.228695,0.095662,1640.151740,0.002880,4.009021,0.002880
+10455,2.690474,3.057884,319650.873508,0.103981,1640.189480,0.002878,4.036363,0.002878
+10456,2.594931,3.057884,319644.282579,0.107731,1640.223300,0.002876,4.063707,0.002876
+10457,2.713377,3.057884,319689.202543,0.093809,1639.992830,0.002875,4.091047,0.002875
+10458,2.631460,3.057884,319642.766432,0.086537,1640.231080,0.002873,4.118391,0.002873
+10459,2.604504,3.057884,319600.727322,0.093811,1640.446830,0.002872,4.145738,0.002872
+10460,2.740415,3.057884,319638.385664,0.089191,1640.253560,0.002870,4.173082,0.002870
+10461,2.793469,3.057884,319638.978076,0.107256,1640.250520,0.002869,4.200426,0.002869
+10462,2.697177,3.057884,319537.458607,0.085313,1640.771640,0.002867,4.227779,0.002867
+10463,2.592596,3.057884,319642.569608,0.079506,1640.232090,0.002865,4.255122,0.002865
+10464,2.562007,3.057884,319654.461409,0.088091,1640.171070,0.002864,4.282465,0.002864
+10465,2.698494,3.057884,319624.503635,0.105216,1640.324800,0.002862,4.309811,0.002862
+10466,2.692487,3.057884,319595.698957,0.108205,1640.472640,0.002861,4.337158,0.002861
+10467,2.622339,3.057884,319611.799633,0.103044,1640.390000,0.002859,4.364505,0.002859
+10468,2.606605,3.057884,319644.964657,0.105802,1640.219800,0.002857,4.391848,0.002857
+10469,2.572092,3.057884,319602.496346,0.099942,1640.437750,0.002856,4.419195,0.002856
+10470,2.632645,3.057884,319623.531314,0.090137,1640.329790,0.002854,4.446540,0.002854
+10471,2.471226,3.057884,319610.624758,0.082590,1640.396030,0.002852,4.473887,0.002852
+10472,2.619698,3.057884,319604.132905,0.094196,1640.429350,0.002850,4.501234,0.002850
+10473,2.572552,3.057884,319644.278683,0.091742,1640.223320,0.002849,4.528577,0.002849
+10474,2.707855,3.057884,319655.570343,0.088665,1640.165380,0.002847,4.555920,0.002847
+10475,2.690413,3.057884,319619.661577,0.091093,1640.349650,0.002845,4.583265,0.002845
+10476,2.569057,3.057884,319588.333009,0.091402,1640.510450,0.002844,4.610614,0.002844
+10477,2.542245,3.057884,319650.351213,0.097850,1640.192160,0.002842,4.637957,0.002842
+10478,2.765682,3.057884,319579.742106,0.119198,1640.554550,0.002840,4.665306,0.002840
+10479,2.575944,3.057884,319606.743644,0.105098,1640.415950,0.002838,4.692653,0.002838
+10480,2.578832,3.057884,319619.934364,0.091078,1640.348250,0.002837,4.719998,0.002837
+10481,2.686681,3.057884,319649.289089,0.101555,1640.197610,0.002835,4.747342,0.002835
+10482,2.625229,3.057884,319674.920419,0.101058,1640.066100,0.002833,4.774682,0.002833
+10483,2.705569,3.057884,319615.123620,0.100343,1640.372940,0.002831,4.802028,0.002831
+10484,2.658226,3.057884,319577.663606,0.108719,1640.565220,0.002829,4.829377,0.002829
+10485,2.493976,3.057884,319637.415209,0.103367,1640.258540,0.002828,4.856722,0.002828
+10486,2.596253,3.057884,319582.358289,0.106996,1640.541120,0.002826,4.884070,0.002826
+10487,2.647026,3.057884,319631.935554,0.101760,1640.286660,0.002824,4.911415,0.002824
+10488,2.614134,3.057884,319580.318715,0.092354,1640.551590,0.002822,4.938764,0.002822
+10489,2.583889,3.057884,319599.538891,0.089481,1640.452930,0.002820,4.966112,0.002820
+10490,2.614840,3.057884,319622.496646,0.097606,1640.335100,0.002819,4.993458,0.002819
+10491,2.675368,3.057884,319612.810850,0.110222,1640.384810,0.002817,5.020804,0.002817
+10492,2.545220,3.057884,319650.583127,0.108203,1640.190970,0.002815,5.048147,0.002815
+10493,2.488529,3.057884,319627.558982,0.104375,1640.309120,0.002813,5.075492,0.002813
+10494,2.621747,3.057884,319635.828974,0.091405,1640.266680,0.002811,5.102836,0.002811
+10495,2.579574,3.057884,319643.072388,0.081315,1640.229510,0.002809,5.130180,0.002809
+10496,2.654875,3.057884,319584.288798,0.099719,1640.531210,0.002807,5.157528,0.002807
+10497,2.692438,3.057884,319601.226075,0.106274,1640.444270,0.002805,5.184876,0.002805
+10498,2.574060,3.057884,319667.199988,0.093198,1640.105710,0.002803,5.212217,0.002803
+10499,2.617799,3.057884,319605.342803,0.087670,1640.423140,0.002801,5.239564,0.002801
+10500,2.553835,3.051036,319630.963189,0.085366,1640.291650,0.002800,5.284262,0.002800
+10501,2.591541,3.051036,319871.241282,0.086118,1639.059510,0.002798,5.311586,0.002798
+10502,2.603822,3.051036,319906.146674,0.086761,1638.880670,0.002796,5.338908,0.002796
+10503,2.656048,3.051036,319874.078866,0.091617,1639.044970,0.002794,5.366231,0.002794
+10504,2.644125,3.051036,319808.150244,0.091598,1639.382860,0.002792,5.393561,0.002792
+10505,2.534880,3.051036,319901.231666,0.090393,1638.905850,0.002790,5.420883,0.002790
+10506,2.605783,3.051036,319836.852753,0.091925,1639.235740,0.002788,5.448210,0.002788
+10507,2.547550,3.051036,319867.751943,0.098600,1639.077390,0.002786,5.475535,0.002786
+10508,2.720389,3.051036,319909.076623,0.093485,1638.865660,0.002784,5.502856,0.002784
+10509,2.639339,3.051036,319870.973920,0.091895,1639.060880,0.002782,5.530180,0.002782
+10510,2.661986,3.051036,319937.992486,0.089784,1638.717540,0.002780,5.557499,0.002780
+10511,2.803216,3.051036,319920.578310,0.090028,1638.806740,0.002778,5.584819,0.002778
+10512,2.688303,3.051036,319907.854667,0.082248,1638.871920,0.002776,5.612140,0.002776
+10513,2.658484,3.051036,319958.120716,0.074884,1638.614450,0.002774,5.639457,0.002774
+10514,2.588421,3.051036,319944.648248,0.070931,1638.683450,0.002771,5.666775,0.002771
+10515,2.576880,3.051036,319787.204101,0.068283,1639.490240,0.002769,5.694106,0.002769
+10516,2.532889,3.051036,319845.071139,0.070746,1639.193620,0.002767,5.721433,0.002767
+10517,2.636373,3.051036,319900.938879,0.065863,1638.907350,0.002765,5.748754,0.002765
+10518,2.630598,3.051036,319871.411069,0.069664,1639.058640,0.002763,5.776078,0.002763
+10519,2.666814,3.051036,319926.960035,0.070672,1638.774050,0.002761,5.803398,0.002761
+10520,2.672949,3.051036,319839.679969,0.103110,1639.221250,0.002759,5.830725,0.002759
+10521,2.664687,3.051036,319914.698526,0.097227,1638.836860,0.002757,5.858045,0.002757
+10522,2.616051,3.051036,319914.079717,0.081522,1638.840030,0.002755,5.885365,0.002755
+10523,2.680714,3.051036,319937.447777,0.083714,1638.720330,0.002753,5.912684,0.002753
+10524,2.623666,3.051036,319858.482534,0.074971,1639.124890,0.002750,5.940009,0.002750
+10525,2.672380,3.051036,319842.239923,0.074018,1639.208130,0.002748,5.967335,0.002748
+10526,2.655396,3.051036,319952.165355,0.076316,1638.644950,0.002746,5.994653,0.002746
+10527,2.548869,3.051036,319883.185236,0.065984,1638.998310,0.002744,6.021976,0.002744
+10528,2.622389,3.051036,319904.192756,0.069845,1638.890680,0.002742,6.049297,0.002742
+10529,2.543951,3.051036,319888.616907,0.075224,1638.970480,0.002740,6.076619,0.002740
+10530,2.603323,3.051036,319905.147264,0.069526,1638.885790,0.002737,6.103940,0.002737
+10531,2.602228,3.051036,319849.322935,0.065717,1639.171830,0.002735,6.131267,0.002735
+10532,2.511177,3.051036,319887.820590,0.070995,1638.974560,0.002733,6.158589,0.002733
+10533,2.572714,3.051036,319889.963632,0.071372,1638.963580,0.002731,6.185912,0.002731
+10534,2.578027,3.051036,319866.934261,0.069808,1639.081580,0.002729,6.213236,0.002729
+10535,2.542672,3.051036,319887.865479,0.069784,1638.974330,0.002726,6.240559,0.002726
+10536,2.642068,3.051036,319859.229922,0.072882,1639.121060,0.002724,6.267884,0.002724
+10537,2.623832,3.051036,319930.247630,0.073846,1638.757210,0.002722,6.295203,0.002722
+10538,2.641641,3.051036,319902.734652,0.077861,1638.898150,0.002720,6.322525,0.002720
+10539,2.536454,3.051036,319882.968598,0.077846,1638.999420,0.002717,6.349848,0.002717
+10540,2.689143,3.051036,319918.895560,0.080883,1638.815360,0.002715,6.377169,0.002715
+10541,2.682338,3.051036,319865.597490,0.089776,1639.088430,0.002713,6.404494,0.002713
+10542,2.559465,3.051036,319876.645226,0.086228,1639.031820,0.002710,6.431817,0.002710
+10543,2.590423,3.051036,319894.064379,0.074885,1638.942570,0.002708,6.459139,0.002708
+10544,2.609137,3.051036,319926.563732,0.075140,1638.776080,0.002706,6.486459,0.002706
+10545,2.634950,3.051036,319873.518762,0.070040,1639.047840,0.002703,6.513783,0.002703
+10546,2.530984,3.051036,319940.895684,0.072124,1638.702670,0.002701,6.541101,0.002701
+10547,2.675320,3.051036,319913.029504,0.074220,1638.845410,0.002699,6.568421,0.002699
+10548,2.523057,3.051036,319861.961915,0.070509,1639.107060,0.002696,6.595746,0.002696
+10549,2.561991,3.051036,319878.637841,0.066621,1639.021610,0.002694,6.623070,0.002694
+10550,2.596606,3.051036,319919.777929,0.070872,1638.810840,0.002692,6.650390,0.002692
+10551,2.530825,3.051036,319887.475128,0.071982,1638.976330,0.002689,6.677712,0.002689
+10552,2.565084,3.051036,319914.995242,0.080476,1638.835340,0.002687,6.705033,0.002687
+10553,2.577397,3.051036,319917.812131,0.080556,1638.820910,0.002685,6.732353,0.002685
+10554,2.597488,3.051036,319952.313749,0.079997,1638.644190,0.002682,6.759670,0.002682
+10555,2.760314,3.051036,319901.339021,0.073902,1638.905300,0.002680,6.786992,0.002680
+10556,2.546053,3.051036,319911.793851,0.070460,1638.851740,0.002678,6.814312,0.002678
+10557,2.691957,3.051036,319878.762748,0.071473,1639.020970,0.002675,6.841637,0.002675
+10558,2.615556,3.051036,319908.085002,0.072449,1638.870740,0.002673,6.868958,0.002673
+10559,2.512166,3.051036,319942.496663,0.067485,1638.694470,0.002670,6.896276,0.002670
+10560,2.664476,3.051036,319919.258658,0.075668,1638.813500,0.002668,6.923596,0.002668
+10561,2.529170,3.051036,319909.919896,0.077503,1638.861340,0.002665,6.950917,0.002665
+10562,2.605489,3.051036,319883.298435,0.072904,1638.997730,0.002663,6.978241,0.002663
+10563,2.479856,3.051036,319893.919943,0.075315,1638.943310,0.002660,7.005563,0.002660
+10564,2.647278,3.051036,319965.052644,0.082513,1638.578950,0.002658,7.032879,0.002658
+10565,2.815350,3.051036,319930.884072,0.339185,1638.753950,0.002656,7.060198,0.002656
+10566,2.538144,3.051036,319846.721891,0.076109,1639.185160,0.002653,7.087524,0.002653
+10567,2.646501,3.051036,319916.016188,0.068290,1638.830110,0.002651,7.114845,0.002651
+10568,2.724876,3.051036,319890.250544,0.071005,1638.962110,0.002648,7.142167,0.002648
+10569,2.591948,3.051036,319878.319724,0.067990,1639.023240,0.002646,7.169491,0.002646
+10570,2.525646,3.051036,319902.555073,0.064275,1638.899070,0.002643,7.196812,0.002643
+10571,2.493958,3.051036,319926.817522,0.071735,1638.774780,0.002641,7.224132,0.002641
+10572,2.664876,3.051036,319940.112772,0.067992,1638.706680,0.002638,7.251450,0.002638
+10573,2.688772,3.051036,319959.958135,0.069894,1638.605040,0.002636,7.278767,0.002636
+10574,2.647939,3.051036,319906.127156,0.080451,1638.880770,0.002633,7.306088,0.002633
+10575,2.795584,3.051036,319946.788147,0.089154,1638.672490,0.002630,7.333406,0.002630
+10576,2.657212,3.051036,319947.565234,0.093406,1638.668510,0.002628,7.360723,0.002628
+10577,2.699704,3.051036,319928.637013,0.084709,1638.765460,0.002625,7.388042,0.002625
+10578,2.545095,3.051036,319905.959285,0.067302,1638.881630,0.002623,7.415364,0.002623
+10579,2.654643,3.051036,319868.755023,0.070425,1639.072250,0.002620,7.442688,0.002620
+10580,2.620667,3.051036,319853.561185,0.078753,1639.150110,0.002618,7.470014,0.002618
+10581,2.590570,3.051036,319912.211590,0.076353,1638.849600,0.002615,7.497335,0.002615
+10582,2.625221,3.051036,319877.213148,0.069685,1639.028910,0.002612,7.524670,0.002612
+10583,2.624538,3.051036,319885.384816,0.070212,1638.987040,0.002610,7.551993,0.002610
+10584,2.696926,3.051036,319922.573427,0.077568,1638.796520,0.002607,7.579313,0.002607
+10585,2.637998,3.051036,319951.495634,0.075602,1638.648380,0.002605,7.606630,0.002605
+10586,2.625880,3.051036,319892.559523,0.072780,1638.950280,0.002602,7.633953,0.002602
+10587,2.543417,3.051036,319930.997304,0.077396,1638.753370,0.002599,7.661272,0.002599
+10588,2.634770,3.051036,319865.931195,0.083200,1639.086720,0.002597,7.688596,0.002597
+10589,2.537612,3.051036,319931.973448,0.086491,1638.748370,0.002594,7.715916,0.002594
+10590,2.513633,3.051036,319879.916174,0.088868,1639.015060,0.002591,7.743240,0.002591
+10591,2.589366,3.051036,319882.705120,0.090307,1639.000770,0.002589,7.770563,0.002589
+10592,2.695950,3.051036,319905.047714,0.090116,1638.886300,0.002586,7.797884,0.002586
+10593,2.564269,3.051036,319994.318460,0.087173,1638.429090,0.002583,7.825198,0.002583
+10594,2.665426,3.051036,319902.598016,0.072874,1638.898850,0.002581,7.852520,0.002581
+10595,2.831981,3.051036,319921.505588,0.074998,1638.801990,0.002578,7.879840,0.002578
+10596,2.514539,3.051036,319895.612190,0.077179,1638.934640,0.002575,7.907162,0.002575
+10597,2.710741,3.051036,319842.196996,0.080509,1639.208350,0.002573,7.934488,0.002573
+10598,2.496753,3.051036,319859.019169,0.077677,1639.122140,0.002570,7.961814,0.002570
+10599,2.580015,3.051036,319896.268013,0.076440,1638.931280,0.002567,7.989136,0.002567
+10600,2.671192,3.051036,319979.979768,0.075587,1638.502510,0.002564,8.016451,0.002564
+10601,2.679725,3.051036,319858.843544,0.081872,1639.123040,0.002562,8.043776,0.002562
+10602,2.659037,3.051036,319902.166638,0.080890,1638.901060,0.002559,8.071098,0.002559
+10603,2.641108,3.051036,319917.566164,0.072960,1638.822170,0.002556,8.098418,0.002556
+10604,2.565979,3.051036,319955.225016,0.071166,1638.629280,0.002554,8.125735,0.002554
+10605,2.522801,3.051036,319917.259680,0.070900,1638.823740,0.002551,8.153056,0.002551
+10606,2.649387,3.051036,319897.413757,0.073748,1638.925410,0.002548,8.180378,0.002548
+10607,2.512681,3.051036,319832.792494,0.077130,1639.256550,0.002545,8.207705,0.002545
+10608,2.555712,3.051036,319773.065346,0.087002,1639.562730,0.002542,8.235040,0.002542
+10609,2.616781,3.051036,319949.033509,0.102688,1638.660990,0.002540,8.262357,0.002540
+10610,2.617699,3.051036,319860.712997,0.105337,1639.113460,0.002537,8.289682,0.002537
+10611,2.610990,3.051036,319924.480717,0.086641,1638.786750,0.002534,8.317002,0.002534
+10612,2.633626,3.051036,319880.372863,0.082724,1639.012720,0.002531,8.344325,0.002531
+10613,2.585816,3.051036,319935.231863,0.086433,1638.731680,0.002528,8.371644,0.002528
+10614,2.740581,3.051036,319885.736127,0.074876,1638.985240,0.002526,8.398967,0.002526
+10615,2.657877,3.051036,319863.640162,0.069786,1639.098460,0.002523,8.426292,0.002523
+10616,2.518871,3.051036,319810.331236,0.074594,1639.371680,0.002520,8.453621,0.002520
+10617,2.637612,3.051036,319851.637177,0.075239,1639.159970,0.002517,8.480947,0.002517
+10618,2.510497,3.051036,319875.222502,0.070262,1639.039110,0.002514,8.508272,0.002514
+10619,2.660985,3.051036,319864.656877,0.067809,1639.093250,0.002511,8.535597,0.002511
+10620,2.566689,3.051036,319933.332251,0.074619,1638.741410,0.002509,8.562915,0.002509
+10621,2.616558,3.051036,319946.296124,0.075489,1638.675010,0.002506,8.590234,0.002506
+10622,2.577765,3.051036,319908.672557,0.083309,1638.867730,0.002503,8.617555,0.002503
+10623,2.569064,3.051036,319906.552688,0.088770,1638.878590,0.002500,8.644876,0.002500
+10624,2.644313,3.051036,319954.742730,0.079174,1638.631750,0.002497,8.672193,0.002497
+10625,2.600619,3.051036,319903.652065,0.071480,1638.893450,0.002494,8.699514,0.002494
+10626,2.528054,3.051036,319873.975433,0.067265,1639.045500,0.002491,8.726838,0.002491
+10627,2.437984,3.051036,319861.850683,0.063434,1639.107630,0.002488,8.754163,0.002488
+10628,2.533826,3.051036,319901.504935,0.066275,1638.904450,0.002485,8.781484,0.002485
+10629,2.619308,3.051036,319866.645440,0.063366,1639.083060,0.002483,8.808809,0.002483
+10630,2.612548,3.051036,319902.262282,0.061399,1638.900570,0.002480,8.836131,0.002480
+10631,2.688943,3.051036,319917.408040,0.064647,1638.822980,0.002477,8.863451,0.002477
+10632,2.664616,3.051036,319912.521967,0.066990,1638.848010,0.002474,8.890772,0.002474
+10633,2.594924,3.051036,319893.080658,0.072295,1638.947610,0.002471,8.918095,0.002471
+10634,2.611122,3.051036,319874.718990,0.071003,1639.041690,0.002468,8.945419,0.002468
+10635,2.570190,3.051036,319931.516612,0.067295,1638.750710,0.002465,8.972738,0.002465
+10636,2.671571,3.051036,319630.390294,0.063783,1640.294590,0.002462,9.000083,0.002462
+10637,2.517806,3.051036,319808.271194,0.067984,1639.382240,0.002459,9.027413,0.002459
+10638,2.718278,3.051036,319768.562037,0.074768,1639.585820,0.002456,9.054746,0.002456
+10639,2.578080,3.051036,319795.406283,0.070484,1639.448190,0.002453,9.082079,0.002453
+10640,2.564280,3.051036,319810.826743,0.063558,1639.369140,0.002450,9.109409,0.002450
+10641,2.535107,3.051036,319922.567569,0.078593,1638.796550,0.002447,9.136729,0.002447
+10642,2.620808,3.051036,319849.067317,0.078063,1639.173140,0.002444,9.164055,0.002444
+10643,2.642947,3.051036,319856.894105,0.068785,1639.133030,0.002441,9.191381,0.002441
+10644,2.641507,3.051036,319888.954563,0.070789,1638.968750,0.002438,9.218704,0.002438
+10645,2.421164,3.051036,319978.624475,0.075252,1638.509450,0.002435,9.246019,0.002435
+10646,2.469529,3.051036,319863.685046,0.068891,1639.098230,0.002432,9.273344,0.002432
+10647,2.558671,3.051036,319909.299152,0.076300,1638.864520,0.002429,9.300665,0.002429
+10648,2.653425,3.051036,319838.476103,0.073404,1639.227420,0.002426,9.327992,0.002426
+10649,2.666234,3.051036,319836.304484,0.065529,1639.238550,0.002423,9.355319,0.002423
+10650,2.587936,3.051036,319775.314124,0.073040,1639.551200,0.002420,9.382651,0.002420
+10651,2.774417,3.051036,319917.593493,0.095976,1638.822030,0.002417,9.409972,0.002417
+10652,2.539855,3.051036,319776.774966,0.061902,1639.543710,0.002414,9.437305,0.002414
+10653,2.685279,3.051036,319834.452876,0.063734,1639.248040,0.002411,9.464632,0.002411
+10654,2.726453,3.051036,319805.436728,0.067314,1639.396770,0.002408,9.491962,0.002408
+10655,2.661300,3.051036,319869.014576,0.073526,1639.070920,0.002405,9.519287,0.002405
+10656,2.544700,3.051036,319911.561559,0.070344,1638.852930,0.002402,9.546607,0.002402
+10657,2.502671,3.051036,319895.118371,0.072678,1638.937170,0.002398,9.573930,0.002398
+10658,2.537918,3.051036,319849.970764,0.071267,1639.168510,0.002395,9.601255,0.002395
+10659,2.595579,3.051036,319904.817382,0.068044,1638.887480,0.002392,9.628576,0.002392
+10660,2.550340,3.051036,319914.538455,0.066611,1638.837680,0.002389,9.655897,0.002389
+10661,2.609774,3.051036,319908.233355,0.066748,1638.869980,0.002386,9.683218,0.002386
+10662,2.601235,3.051036,319858.189825,0.063435,1639.126390,0.002383,9.710543,0.002383
+10663,2.647851,3.051036,319884.924210,0.066107,1638.989400,0.002380,9.737866,0.002380
+10664,2.568462,3.051036,319859.659231,0.062815,1639.118860,0.002377,9.765192,0.002377
+10665,2.614563,3.051036,319962.272027,0.063989,1638.593190,0.002374,9.792508,0.002374
+10666,2.678935,3.051036,319922.524622,0.071495,1638.796770,0.002370,9.819828,0.002370
+10667,2.500076,3.051036,319974.062649,0.066418,1638.532810,0.002367,9.847143,0.002367
+10668,2.561331,3.051036,319864.732984,0.065154,1639.092860,0.002364,9.874468,0.002364
+10669,2.660706,3.051036,319922.389921,0.069331,1638.797460,0.002361,9.901788,0.002361
+10670,2.620308,3.051036,319898.737132,0.066663,1638.918630,0.002358,9.929110,0.002358
+10671,2.481086,3.051036,319933.217065,0.065410,1638.742000,0.002355,9.956428,0.002355
+10672,2.541093,3.051036,319862.168767,0.070190,1639.106000,0.002352,9.983753,0.002352
+10673,2.683227,3.051036,319842.083828,0.072099,1639.208930,0.002348,10.011081,0.002348
+10674,2.566789,3.051036,319925.277218,0.068605,1638.782670,0.002345,10.038400,0.002345
+10675,2.580904,3.051036,319937.801155,0.072783,1638.718520,0.002342,10.065719,0.002342
+10676,2.597879,3.051036,319881.016913,0.074251,1639.009420,0.002339,10.093042,0.002339
+10677,2.633514,3.051036,319858.494243,0.072311,1639.124830,0.002336,10.120367,0.002336
+10678,2.640936,3.051036,319843.508209,0.079450,1639.201630,0.002332,10.147694,0.002332
+10679,2.654224,3.051036,319876.522273,0.079399,1639.032450,0.002329,10.175018,0.002329
+10680,2.556746,3.051036,319905.164833,0.081221,1638.885700,0.002326,10.202339,0.002326
+10681,2.599396,3.051036,319826.303318,0.086565,1639.289810,0.002323,10.229667,0.002323
+10682,2.592875,3.051036,319887.449756,0.094078,1638.976460,0.002320,10.256990,0.002320
+10683,2.569088,3.051036,319857.268771,0.094845,1639.131110,0.002316,10.284316,0.002316
+10684,2.662574,3.051036,319904.946212,0.086158,1638.886820,0.002313,10.311638,0.002313
+10685,2.543761,3.051036,319893.170442,0.080795,1638.947150,0.002310,10.338960,0.002310
+10686,2.545193,3.051036,319846.181396,0.078171,1639.187930,0.002307,10.366286,0.002307
+10687,2.580250,3.051036,319824.252825,0.076685,1639.300320,0.002303,10.393615,0.002303
+10688,2.549900,3.051036,319939.035054,0.074892,1638.712200,0.002300,10.420934,0.002300
+10689,2.707404,3.051036,319883.446765,0.071871,1638.996970,0.002297,10.448257,0.002297
+10690,2.600739,3.051036,319844.883821,0.070457,1639.194580,0.002294,10.475584,0.002294
+10691,2.674275,3.051036,319963.730674,0.066224,1638.585720,0.002290,10.502900,0.002290
+10692,2.572013,3.051036,319887.791312,0.064391,1638.974710,0.002287,10.530223,0.002287
+10693,2.521792,3.051036,319880.415800,0.069601,1639.012500,0.002284,10.557546,0.002284
+10694,2.500074,3.051036,319838.401960,0.073791,1639.227800,0.002281,10.584874,0.002281
+10695,2.579055,3.051036,319931.516612,0.068063,1638.750710,0.002277,10.612194,0.002277
+10696,2.675575,3.051036,319952.065775,0.069365,1638.645460,0.002274,10.639511,0.002274
+10697,2.736763,3.051036,319867.935386,0.064741,1639.076450,0.002271,10.666835,0.002271
+10698,2.598721,3.051036,319905.207777,0.061174,1638.885480,0.002268,10.694157,0.002268
+10699,2.507962,3.051036,319882.685602,0.057756,1639.000870,0.002264,10.721480,0.002264
+10700,2.495724,3.051036,319934.460692,0.080465,1638.735630,0.002261,10.748799,0.002261
+10701,2.673309,3.051036,319869.771770,0.066389,1639.067040,0.002258,10.776123,0.002258
+10702,2.641712,3.051036,319810.916480,0.062941,1639.368680,0.002254,10.803452,0.002254
+10703,2.574183,3.051036,319910.970089,0.063474,1638.855960,0.002251,10.830774,0.002251
+10704,2.662001,3.051036,319849.806855,0.069015,1639.169350,0.002248,10.858100,0.002248
+10705,2.520388,3.051036,319911.481525,0.062682,1638.853340,0.002244,10.885421,0.002244

150M_228M_3e3_6B_incomp/blumodel_muon_step00010706.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:086c1617701ba929341a9bd40e6ec9b9d56ac0f3c536be7e03d54d0753a2f855
+size 2515184199

150M_228M_75e5_10B_half/blumodel_muon_step00000953.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19c6294243727cd2bc282c9a530ca3c98521a85b9f2c0df4dbe82e8e048a4dcf
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00001906.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00002859.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00003812.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00004765.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35c731f2894426e91642b9705ae6e0deb7e0caa7c8c70ff6d065890ec6eb0f7b
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00005718.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00006671.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cd83878afac59c02cb76d00b723ec319b95412ad7c2bf5b0cd59aeb58d10c96
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00007624.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:482c93f862bb26681afd5af4d853a36096cc91a22c88c871125069a50c9cc4da
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00008577.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c050edb018b294e3a6e84149d3bed6dc8e68fe4444a851c86bd5b41c9703445
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00009530.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19bd8cdf882c01d35d8d028aa8213bef4b52b85d5421042c76a68fae5ee5919d
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00010483.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19116842c7aeb68d9ffe03e812646c1095c505d80af2aeb6cf706e6a4b6a44a6
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00011436.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2427673fceb45df25926a5ddde0d3e763fa235c93160ed0f555edc564be9b89c
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00012389.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00013342.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00014295.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cb1589be099c8c7ad095906fa3813394c56e992a281c8e65d1d469217153dec
+size 2515183879

150M_228M_75e5_10B_half/blumodel_muon_step00015248.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00016201.pt

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

150M_228M_75e5_10B_half/blumodel_muon_step00016332.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ab005c8850f5e064ba08efd426345f3b64d5678bb01f086d6ae9000f4975615
+size 2515183879

train_muon.py

@@ -0,0 +1,1829 @@
+
+#!/usr/bin/env python3
+# =============================================================================
+# train_muon.py — the MUON arm of train.py. Same model, same data order, same
+# schedule shape, same seed: the ONLY thing that changes is how a gradient
+# becomes a weight update.
+#
+# Fork discipline — everything outside section 5 is byte-identical to train.py,
+# so `diff train.py train_muon.py` is exactly the experiment being run. If you
+# edit a kernel or the model in one file, port it to the other or the comparison
+# stops meaning anything.
+#
+# The optimizer itself is torch.optim.Muon — stock, not reimplemented. This file
+# supplies only the two things it does not do for you: the ROUTING (which
+# parameters it may have) and the fp32 MASTER WEIGHTS (it updates the master,
+# not the bf16 param). Requires a torch new enough to ship it; main() checks and
+# says so.
+#
+# What Muon does: AdamW rescales each weight ENTRY by its own gradient history.
+# Muon rescales each weight MATRIX by its singular values — it takes the
+# momentum buffer M and replaces it with the nearest semi-orthogonal matrix
+# (all singular values driven to ~1) via a 5-step Newton-Schulz iteration, then
+# steps along that. Directions the gradient has barely explored get the same
+# step size as the dominant ones, which is where the reported speedup comes
+# from. Only 2D hidden weights can be orthogonalized, so this is a HYBRID:
+#
+#     Muon   : every 2D Linear weight inside the blocks   (75.5M params, 66%)
+#     AdamW  : the token embedding / lm_head              (38.6M params)
+#              + RMSNorm weights + QK-norm gammas (1D)    (22.3k params)
+#
+# The embedding is deliberately NOT given to Muon: its rows are per-token, not a
+# linear map between two feature spaces, and orthogonalizing across the vocab
+# axis mixes unrelated tokens. That split is the standard Muon recipe, and it
+# keeps the AdamW arm's exact update on 34% of the parameters — so a difference
+# in the loss curve is attributable to the hidden matrices alone.
+#
+# Everything train_ddp.py / train_zero1.py import from blu_common / blu_data /
+# blu_kernels / blu_model is inlined here, in that order, so this file runs on
+# its own:
+#     section 1  fused QK-norm + RoPE (Triton)  <- blu_kernels.py
+#     section 2  config + WSD schedule          <- blu_common.py
+#     section 3  data loader                    <- blu_data.py
+#     section 4  model                          <- blu_model.py
+#     section 5  optimizer / clipping / checkpointing
+#     section 6  training
+#
+# Architecture:
+#   token embedding
+#   -> N x { RMSNorm -> [fused QK-norm + RoPE] -> causal GQA (flash) -> Wo -> +x
+#            RMSNorm -> SwiGLU FFN (gate_up -> silu(gate)*up -> down)  -> +x }
+#   -> final RMSNorm
+#   -> lm_head (weight-tied with the token embedding)
+#   -> cross-entropy
+#   -> backward -> AdamW (decoupled weight decay) on a WSD (warmup / stable /
+#      cosine-decay) schedule
+#
+# Kept from the AdamW arm, unchanged, because they are confounds if they move:
+#   * fp32 master weights for every non-fp32 param — Muon updates the MASTER,
+#     reading the same fp32 landing buffer AdamW reads, so both optimizers see
+#     bit-identical gradients;
+#   * decoupled weight decay 0.1 on dim>=2 params, applied to the master;
+#   * ONE global-norm gradient clip over every parameter, BEFORE the optimizer.
+#     (Muon's orthogonalization makes the update scale-invariant, so clipping is
+#     nearly a no-op for its group — but it still gates the AdamW group, and
+#     removing it would change two things at once.)
+#
+# Memory, at the 114M config this repo actually runs (printed at startup):
+#
+#     parameters (bf16 Linear + fp32 embed/norms)  0.28 GiB
+#     gradients (same dtypes as the params)        0.28 GiB
+#     fp32 master weights (the bf16 Linears)       0.28 GiB
+#     fp32 gradient landing buffers                0.28 GiB
+#     Muon momentum   (75.5M params, fp32 x1)      0.28 GiB
+#     AdamW m + v     (38.7M params, fp32 x2)      0.29 GiB
+#     -------------------------------------------------------
+#     static total                                 1.70 GiB   (AdamW arm: 1.98)
+#
+# Muon carries ONE state tensor per param where Adam carries two, so the saving
+# scales with how much of the model is hidden matrices: 0.28 GiB here (14% of
+# static), because this config spends a third of its parameters on a 50304-row
+# embedding that stays on AdamW. At the 1.73B config in the original header the
+# same split saves ~5.5 GiB.
+#
+# Activations dominate either way (63.96 GiB measured at B=16 x T=4096), so the
+# OOM escape hatches are unchanged: --B 8, --loss-chunks 4, --activation-checkpoint.
+#
+# Fidelity notes (why a run here is statistically but not bit-wise identical):
+#   * init draws come from torch's Philox, not the C++ single advancing
+#     std::mt19937, so the two runs start from different samples of the same
+#     distributions;
+#   * QK-norm reads the projection output in bf16 here (autocast) where the C++
+#     casts it up to fp32 first; the reduction and the rotation are fp32 in both.
+#
+# Determinism: two runs of the same launch line on the same GPU produce
+# bit-identical losses. Seeding is only half of that — the default CUDA kernels
+# for the embedding backward and the flash/mem-efficient attention backward
+# accumulate partial sums with atomicAdd, so the order they land in (and with it
+# the fp32 rounding) follows the scheduler, not the program. --determinism
+# strict (the default) forces the ordered variants of those kernels; it costs a
+# few percent of step time, which --determinism off buys back.
+#
+# Env overrides (the bluscript_single.cpp names, so an existing launch line keeps
+# working; the equivalent CLI flag wins if both are given):
+#   BLU_DEVICE=<idx>          CUDA device to train on            (default 0)
+#   BLU_MAX_STEPS=<n>         short benchmark run (rescales WSD) (default off)
+#   BLU_B=<n>                 micro-batch size                   (default 4)
+#   BLU_GLOBAL_BATCH=<tokens> tokens per optimizer step          (default 524288)
+#   BLU_DIAG=1                per-micro-batch loss + depth probes
+#
+# Multi-GPU: DDP, one process per GPU, via torchrun. RANK / LOCAL_RANK /
+# WORLD_SIZE (set by torchrun, not this file) switch the bootstrap in main();
+# their absence is the single-process path above, unchanged. --device /
+# BLU_DEVICE are ignored under torchrun — LOCAL_RANK picks each rank's GPU.
+# --global-batch is still the TOTAL tokens per optimizer step across every
+# rank (grad_accum = global_batch / (B*T*world_size)), so the same
+# --global-batch reproduces the same schedule at any world size. Data is
+# rank-sharded (DataLoaderLite reads a disjoint, rank-offset slice of the
+# shards); gradients and the reported train/val loss are averaged across
+# ranks; only rank 0 writes the CSV and checkpoints. Same launch line as
+# single-GPU otherwise:
+#   torchrun --standalone --nproc_per_node=8 train_muon.py
+#
+# Learning rate — the one knob that decides whether the comparison is fair.
+#
+# Muon's update is spectrally normalized, so its natural scale has nothing to do
+# with AdamW's. torch.optim.Muon ships both published conventions behind its
+# `adjust_lr_fn`; --muon-adjust-lr passes the string straight through:
+#
+#   match_rms_adamw      Moonshot's: gamma <- 0.2 * gamma * sqrt(max(A,B)). The
+#     (DEFAULT HERE)     Newton-Schulz output has singular values ~1, so its
+#                        entries have RMS 1/sqrt(max(A,B)); this lifts them to
+#                        RMS 0.2, which is what a converged AdamW update looks
+#                        like. The consequence: --muon-lr is on the SAME SCALE
+#                        as --max-lr, and the default (2e-4, equal to the AdamW
+#                        arm's max_lr) gives both arms the same step size on the
+#                        first run — only the update DIRECTION differs. That is
+#                        the controlled experiment; run it before tuning.
+#   original             Keller's modded-nanogpt scaling, sqrt(max(1, A/B)).
+#                        Entries stay at RMS 1/sqrt(max(A,B)), so an equivalent
+#                        LR is ~5.5x larger at d_model=768 (0.2*sqrt(768)).
+#                        Use it to reproduce published Muon numbers.
+#
+# NOTE: torch's own default for adjust_lr_fn is `original`. This file defaults
+# to `match_rms_adamw` instead, because an LR-matched first run is the point of
+# having two arms. If you cite a number from here, cite the flag with it.
+#
+# Muon usually tolerates 2-4x more LR than AdamW once it is working, so the
+# honest protocol is: baseline run, matched run, then a short --max-steps 200
+# sweep over --muon-lr before claiming a win.
+#
+# Launch (the two arms; separate CSV + checkpoint dirs so they never collide):
+#   python3 train.py                                  # AdamW baseline
+#   python3 train_muon.py                             # Muon, LR-matched
+#   python3 train_muon.py --muon-lr 6e-4              # after a sweep
+#   python3 train_muon.py --max-steps 200 --muon-lr 4e-4   # short sweep point
+# =============================================================================
+
+import argparse
+import glob
+import math
+import os
+import random
+import re
+import signal
+import time
+from contextlib import nullcontext
+from dataclasses import dataclass
+
+# cuBLAS reads CUBLAS_WORKSPACE_CONFIG once, when it creates its handle on the
+# first GEMM, and torch refuses to run fp32 cuBLAS GEMMs under
+# use_deterministic_algorithms() unless it is one of the two allowed values. It
+# has to be in the environment before torch touches CUDA, hence: before the
+# import, not inside main(). An explicit setting from the launch line wins.
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.checkpoint import checkpoint
+
+try:
+    import triton
+    import triton.language as tl
+    _HAS_TRITON = True
+except ImportError:                                   # eager fallback below
+    _HAS_TRITON = False
+
+
+# =============================================================================
+# 1. Fused QK-norm (RMSNorm over head_dim) + RoPE, in Triton.
+#
+# PyTorch counterpart of the C++ `autograd::fused_rms_rope`. One kernel does,
+# per (batch, head, position) row:
+#
+#     s = rsqrt(mean(x^2) + eps)          # RMS over head_dim
+#     y = x * s * gamma                   # learnable per-head_dim scale
+#     o = rope_neox(y, cos[t], sin[t])    # NeoX / Llama half-split rotation
+#
+# Attention is deliberately NOT fused in here — it stays a separate op
+# (F.scaled_dot_product_attention on the flash backend), mirroring the C++ split
+# between `fused_rms_rope` and `gqa_fused_flash_attn_forward/backward`.
+#
+# Why fuse: the eager formulation (`rms_rope_eager`) materialises an fp32 copy of
+# x, the squared tensor, the normalised tensor, the gamma-scaled tensor and a
+# `cat` of the two rotated halves — roughly 6 extra [B,H,T,D] buffers held for
+# backward. The fused version saves only x (which the projection's backward
+# already holds) plus one fp32 scalar per row.
+#
+# Layout notes
+#   - x is [B,H,T,D] and may be a non-contiguous transpose of a [B,T,H,D] view;
+#     arbitrary (b,h,t) strides are honoured, D must be stride-1.
+#   - D must be a power of two (the C++ kernel requires 64 or 128).
+#   - cos/sin are [T_ctx, D//2] fp32, as produced by `build_rope_cache`.
+#   - the row index is r = (b*H + h)*T + t, so the output is contiguous
+#     [B,H,T,D] in that same order.
+# =============================================================================
+if _HAS_TRITON:
+
+    @triton.jit
+    def _fused_rms_rope_fwd(
+        X, GAMMA, COS, SIN, OUT, RSTD,
+        sx_b, sx_h, sx_t,                   # input strides, in elements
+        T, H, N,                            # seq len, n_heads, total rows
+        eps,
+        D: tl.constexpr, HALF: tl.constexpr, BLOCK_R: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)          # [R]
+        mask_r = rows < N
+
+        # r = (b*H + h)*T + t  ->  recover t (for RoPE) and the (b, h) base.
+        t = rows % T
+        bh = rows // T
+        b = bh // H
+        h = bh % H
+        base = b * sx_b + h * sx_h + t * sx_t                  # [R]
+
+        dh = tl.arange(0, HALF)                               # [HALF]
+        p1 = X + base[:, None] + dh[None, :]
+        x1 = tl.load(p1, mask=mask_r[:, None], other=0.0).to(tl.float32)
+        x2 = tl.load(p1 + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)
+
+        # RMS over the FULL head_dim (both halves), fp32 accumulation.
+        ss = tl.sum(x1 * x1, 1) + tl.sum(x2 * x2, 1)          # [R]
+        rstd = 1.0 / tl.sqrt(ss / D + eps)                    # [R]
+
+        g1 = tl.load(GAMMA + dh).to(tl.float32)
+        g2 = tl.load(GAMMA + HALF + dh).to(tl.float32)
+        y1 = x1 * rstd[:, None] * g1[None, :]
+        y2 = x2 * rstd[:, None] * g2[None, :]
+
+        cp = t[:, None] * HALF + dh[None, :]
+        c = tl.load(COS + cp, mask=mask_r[:, None], other=0.0)
+        s = tl.load(SIN + cp, mask=mask_r[:, None], other=0.0)
+
+        # NeoX half-split rotation.
+        o1 = y1 * c - y2 * s
+        o2 = y2 * c + y1 * s
+
+        op = OUT + rows[:, None] * D + dh[None, :]
+        tl.store(op, o1.to(OUT.dtype.element_ty), mask=mask_r[:, None])
+        tl.store(op + HALF, o2.to(OUT.dtype.element_ty), mask=mask_r[:, None])
+        tl.store(RSTD + rows, rstd, mask=mask_r)
+
+    # -------------------------------------------------------------------------
+    # backward
+    #
+    #   RoPE is an orthogonal rotation, so its vjp is the inverse rotation:
+    #       dy1 = do1*cos + do2*sin ;  dy2 = do2*cos - do1*sin
+    #   then, with n = x * rstd (the normalised row),
+    #       dgamma += dy * n
+    #       dn      = dy * gamma
+    #       dx      = rstd * (dn - n * mean(dn * n))
+    # -------------------------------------------------------------------------
+    @triton.jit
+    def _fused_rms_rope_bwd(
+        X, GAMMA, COS, SIN, DOUT, RSTD, DX, DGP,
+        sx_b, sx_h, sx_t,                   # input strides
+        sd_b, sd_h, sd_t,                   # dx strides
+        T, H, N,
+        D: tl.constexpr, HALF: tl.constexpr, BLOCK_R: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
+        mask_r = rows < N
+
+        t = rows % T
+        bh = rows // T
+        b = bh // H
+        h = bh % H
+        base = b * sx_b + h * sx_h + t * sx_t
+
+        dh = tl.arange(0, HALF)
+        p1 = X + base[:, None] + dh[None, :]
+        x1 = tl.load(p1, mask=mask_r[:, None], other=0.0).to(tl.float32)
+        x2 = tl.load(p1 + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)
+
+        rstd = tl.load(RSTD + rows, mask=mask_r, other=1.0)
+        n1 = x1 * rstd[:, None]
+        n2 = x2 * rstd[:, None]
+
+        dp = DOUT + rows[:, None] * D + dh[None, :]
+        do1 = tl.load(dp, mask=mask_r[:, None], other=0.0).to(tl.float32)
+        do2 = tl.load(dp + HALF, mask=mask_r[:, None], other=0.0).to(tl.float32)
+
+        cp = t[:, None] * HALF + dh[None, :]
+        c = tl.load(COS + cp, mask=mask_r[:, None], other=0.0)
+        s = tl.load(SIN + cp, mask=mask_r[:, None], other=0.0)
+
+        # inverse rotation
+        dy1 = do1 * c + do2 * s
+        dy2 = do2 * c - do1 * s
+
+        # per-program partial dgamma; summed by the caller (deterministic, no atomics)
+        tl.store(DGP + pid * D + dh, tl.sum(dy1 * n1, 0))
+        tl.store(DGP + pid * D + HALF + dh, tl.sum(dy2 * n2, 0))
+
+        g1 = tl.load(GAMMA + dh).to(tl.float32)
+        g2 = tl.load(GAMMA + HALF + dh).to(tl.float32)
+        dn1 = dy1 * g1[None, :]
+        dn2 = dy2 * g2[None, :]
+
+        dot = (tl.sum(dn1 * n1, 1) + tl.sum(dn2 * n2, 1)) / D      # [R]
+        dx1 = rstd[:, None] * (dn1 - n1 * dot[:, None])
+        dx2 = rstd[:, None] * (dn2 - n2 * dot[:, None])
+
+        dbase = b * sd_b + h * sd_h + t * sd_t
+        q1 = DX + dbase[:, None] + dh[None, :]
+        tl.store(q1, dx1.to(DX.dtype.element_ty), mask=mask_r[:, None])
+        tl.store(q1 + HALF, dx2.to(DX.dtype.element_ty), mask=mask_r[:, None])
+
+
+_BLOCK_R = 16
+_NUM_WARPS = 4
+
+
+def _launch_fwd(x, gamma, cos, sin, eps):
+    B, H, T, D = x.shape
+    N = B * H * T
+    out = torch.empty((B, H, T, D), dtype=x.dtype, device=x.device)
+    rstd = torch.empty(N, dtype=torch.float32, device=x.device)
+    grid = (triton.cdiv(N, _BLOCK_R),)
+    _fused_rms_rope_fwd[grid](
+        x, gamma, cos, sin, out, rstd,
+        x.stride(0), x.stride(1), x.stride(2),
+        T, H, N, eps,
+        D=D, HALF=D // 2, BLOCK_R=_BLOCK_R, num_warps=_NUM_WARPS,
+    )
+    return out, rstd
+
+
+def _launch_bwd(x, gamma, cos, sin, dout, rstd):
+    B, H, T, D = x.shape
+    N = B * H * T
+    dout = dout.contiguous()
+    dx = torch.empty_like(x)
+    nprog = triton.cdiv(N, _BLOCK_R)
+    dgp = torch.empty((nprog, D), dtype=torch.float32, device=x.device)
+    _fused_rms_rope_bwd[(nprog,)](
+        x, gamma, cos, sin, dout, rstd, dx, dgp,
+        x.stride(0), x.stride(1), x.stride(2),
+        dx.stride(0), dx.stride(1), dx.stride(2),
+        T, H, N,
+        D=D, HALF=D // 2, BLOCK_R=_BLOCK_R, num_warps=_NUM_WARPS,
+    )
+    return dx, dgp.sum(0).to(gamma.dtype)
+
+
+class _FusedRMSRoPE(torch.autograd.Function):
+    """QK-norm + RoPE for q and k in one autograd node (two kernel launches).
+    Mirrors the C++ `fused_rms_rope(q, k, q_gamma, k_gamma, cache)`: q/k keep
+    their incoming dtype (bf16 under autocast), the reduction and the rotation
+    run in fp32 internally, gamma and the cache stay fp32.
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, q_gamma, k_gamma, cos, sin, eps):
+        qo, q_rstd = _launch_fwd(q, q_gamma, cos, sin, eps)
+        ko, k_rstd = _launch_fwd(k, k_gamma, cos, sin, eps)
+        ctx.save_for_backward(q, k, q_gamma, k_gamma, cos, sin, q_rstd, k_rstd)
+        return qo, ko
+
+    @staticmethod
+    def backward(ctx, dqo, dko):
+        q, k, q_gamma, k_gamma, cos, sin, q_rstd, k_rstd = ctx.saved_tensors
+        dq, dqg = _launch_bwd(q, q_gamma, cos, sin, dqo, q_rstd)
+        dk, dkg = _launch_bwd(k, k_gamma, cos, sin, dko, k_rstd)
+        return dq, dk, dqg, dkg, None, None, None
+
+
+def fused_rms_rope(q, k, q_gamma, k_gamma, cos, sin, eps):
+    """q: [B,Hq,T,D], k: [B,Hkv,T,D] -> same shapes, contiguous.
+    `cos`/`sin` are [>=T, D//2] fp32; only the first T rows are read.
+    """
+    assert q.shape[-1] == k.shape[-1], "q and k must share head_dim"
+    D = q.shape[-1]
+    assert D % 2 == 0 and (D & (D - 1)) == 0, f"head_dim must be a power of two, got {D}"
+    assert q.stride(-1) == 1 and k.stride(-1) == 1, "head_dim must be stride-1"
+    T = q.shape[2]
+    return _FusedRMSRoPE.apply(q, k, q_gamma, k_gamma,
+                               cos[:T].contiguous(), sin[:T].contiguous(), eps)
+
+
+def rms_rope_eager(x, gamma, cos, sin, eps):
+    """Unfused QK-norm then NeoX RoPE, fp32 throughout — the reference the fused
+    kernel is measured against (`--rms-rope eager`), and the fallback when
+    Triton is unavailable."""
+    T, D = x.shape[2], x.shape[-1]
+    xf = x.float()
+    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * gamma
+    x1, x2 = xf[..., : D // 2], xf[..., D // 2:]
+    c = cos[:T][None, None, :, :]
+    s = sin[:T][None, None, :, :]
+    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)
+
+
+def build_rope_cache(seq_len, head_dim, theta, device):
+    """NeoX / Llama RoPE cache: cos, sin of shape [seq_len, head_dim//2], fp32."""
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
+                                             device=device) / head_dim))
+    pos = torch.arange(seq_len, dtype=torch.float32, device=device)
+    freqs = torch.outer(pos, inv_freq)                       # [T, D/2]
+    return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous()
+
+
+# =============================================================================
+# 2. Configuration — 1:1 with ModelConfig in bluscript_single.cpp — and the WSD
+#    learning-rate schedule.
+# =============================================================================
+@dataclass
+class ModelConfig:
+    # ---- model ----
+    d_model: int = 768
+    n_layers: int = 24
+    q_heads: int = 12
+    kv_heads: int = 4                     # grouped-query: G = q/kv = 4
+    head_dim: int = 64                   # the fused kernels need 64 or 128
+    ffn_hidden: int = 2048                # ~8/3 * d_model, rounded to /256
+    context_length: int = 4096            # RoPE cache length; must be >= T
+    rope_theta: float = 500000.0
+    vocab_size: int = 50304               # 50257 (tiktoken) padded for the GPU
+    weight_tying: bool = False
+
+    rms_eps: float = 1e-5                 # block RMSNorm epsilon
+    qk_eps: float = 1e-6                  # QK-norm (fused) epsilon
+
+    # ---- optimization / data ----
+    B: int = 4                            # micro-batch
+    T: int = 4096                         # sequence length (<= context_length)
+    global_batch: int = 524288            # tokens per optimizer step
+    grad_clip: float = 1.0
+    weight_decay: float = 0.1
+
+    # WSD: linear warmup -> constant at max_lr -> cosine decay to min_lr.
+    # Requires warmup_steps + stable_steps < max_steps - 1 (the remainder is
+    # the decay tail).
+    max_lr: float = 0.003
+    min_lr: float = 0.0003               # 10% of max_lr
+    
+    muon_lr: float = 0.003
+    muon_min_lr: float = 0.0003
+    
+    max_steps: int = 11444
+    warmup_steps: int = 500
+    stable_steps: int = int(max_steps * 0.9) - warmup_steps           # steps held at max_lr after warmup
+     
+
+    # ---- Muon (hidden 2D weights); the AdamW fields above still drive the
+    #      embedding / lm_head / 1D group, unchanged from the AdamW arm ----
+    # Default == max_lr: under adjust_lr_fn=match_rms_adamw the two are the same unit,
+    # so out of the box both arms take the same size step and differ only in
+    # direction. Same 10% floor, same WSD shape, so the schedules stay aligned.
+   
+    # The rest are torch.optim.Muon's own arguments, passed through unchanged.
+    # momentum/nesterov/ns_steps are torch's defaults; adjust_lr_fn is not (see
+    # the header — torch defaults to "original").
+    muon_momentum: float = 0.95           # beta on the momentum buffer
+    muon_ns_steps: int = 5                # Newton-Schulz iterations
+    muon_adjust_lr: str = "match_rms_adamw"   # "match_rms_adamw" | "original"
+
+    val_freq: int = 250
+    val_loss_steps: int = 20
+
+    # ---- checkpointing ----
+    checkpointing: bool = True
+    ckpt_freq: int = int(max_steps / 20)                 # save every 1800 optimizer steps
+    ckpt_keep: int = 40                    # retain the most-recent N checkpoints
+    # Separate from the AdamW arm's dir: the two checkpoint formats differ (Muon
+    # momentum vs Adam m/v) and auto-resume picks the newest file it finds.
+    ckpt_dir: str = "ckpt_blu_muon"
+    ckpt_prefix: str = "blumodel_muon"
+
+    # ---- runtime knobs (not in the C++ struct) ----
+    fused_rms_rope: bool = True
+    activation_checkpoint: bool = False
+    loss_chunks: int = 1
+    # Depth probe: which layers --diag reports gradient / weight statistics for.
+    probe_layers: tuple = (0, 8, 17, 26, 35)
+
+
+def get_lr(step, cfg, max_lr=None, min_lr=None):
+    """WSD: linear Warmup -> Stable (constant max_lr) -> cosine Decay to min_lr.
+    Byte-for-byte the C++ get_lr().
+    `max_lr`/`min_lr` override the config pair so the Muon group can ride the
+    SAME schedule shape — same warmup length, same stable window, same cosine
+    tail — at its own amplitude. Called with no override it is identical to the
+    AdamW arm's get_lr(), which is what the embedding/norm group uses.
+    """
+    max_lr = cfg.max_lr if max_lr is None else max_lr
+    min_lr = cfg.min_lr if min_lr is None else min_lr
+    if step < cfg.warmup_steps:
+        return max_lr * (step + 1) / cfg.warmup_steps
+    if step < cfg.stable_steps + cfg.warmup_steps:
+        return max_lr
+    tail = cfg.max_steps - (cfg.stable_steps + cfg.warmup_steps) - 1
+    decay_ratio = (step - (cfg.stable_steps + cfg.warmup_steps)) / max(tail, 1)
+    decay_ratio = min(max(decay_ratio, 0.0), 1.0)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))       # 1 -> 0
+    return min_lr + coeff * (max_lr - min_lr)
+
+
+# =============================================================================
+# 3. Data — port of the C++ DataLoaderLite (Tensor-Implementations/include/
+#    DataLoader.h), including its pinned-host / side-stream prefetch ring.
+#
+# Shards are HEADERLESS flat uint16 token files whose basename contains the
+# split ("train" / "val"). For a batch at position p:
+#     x = tokens[p : p+B*T]      y = tokens[p+1 : p+B*T+1]
+# p strides by B*T (world_size == 1 here). A shard tail that cannot fit a full
+# stride is dropped and the loader wraps to the next shard.
+#
+# The prefetch mirrors the C++ ring: NBUF pinned staging buffers and NBUF device
+# tensors, H2D issued on a side stream, CUDA events ordering producer against
+# consumer. Each slot also carries the (shard, token offset) it was filled from,
+# which is what --diag prints to name the bytes behind a loss spike.
+# =============================================================================
+class DataLoaderLite:
+    NBUF = 4
+
+    def __init__(self, B, T, split, data_root, device, master=True,
+                 max_tokens_per_shard=100_000_000, rank=0, world_size=1):
+        assert split in ("train", "val")
+        self.B, self.T = B, T
+        self.split = split
+        self.max_tokens = max_tokens_per_shard
+        self.device = device
+        # DDP: every rank reads the SAME shard list but a DISJOINT, interleaved
+        # slice of it (offset by rank, strided by world_size), so a training
+        # step's world_size*grad_accum micro-batches never repeat a window.
+        self.rank = rank
+        self.world_size = world_size
+
+        self.shards = sorted(
+            p for p in glob.glob(os.path.join(data_root, "*.bin"))
+            if split in os.path.basename(p)
+        )
+        if not self.shards:
+            raise RuntimeError(f"no .bin shards found for split {split} in {data_root}")
+        if master:
+            print(f"found {len(self.shards)} shards for split {split}")
+
+        BT = B * T
+        self.pin_x = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
+        self.pin_y = [torch.empty(BT, dtype=torch.int64).pin_memory() for _ in range(self.NBUF)]
+        self.dev_x = [torch.empty(B, T, dtype=torch.int64, device=device) for _ in range(self.NBUF)]
+        self.dev_y = [torch.empty(B, T, dtype=torch.int64, device=device) for _ in range(self.NBUF)]
+        self.copy_stream = torch.cuda.Stream(device=device)
+        self.h2d_done = [torch.cuda.Event() for _ in range(self.NBUF)]
+        self.h2d_recorded = [False] * self.NBUF
+        self.where = [(0, 0)] * self.NBUF        # (shard_index, token_offset)
+        self.reset()
+
+    # -- shard plumbing ------------------------------------------------------
+    def _load_shard(self, i):
+        self.cur = i
+        arr = np.memmap(self.shards[i], dtype=np.uint16, mode="r")
+        self.tok = arr
+        self.ntok = min(len(arr), self.max_tokens)
+
+    def _advance_shard(self):
+        self._load_shard((self.cur + 1) % len(self.shards))
+        self.pos = self.B * self.T * self.rank
+
+    def reset(self):
+        self.copy_stream.synchronize()
+        self._load_shard(0)
+        self.pos = self.B * self.T * self.rank
+        self.primed = False
+        self.buf = 0
+
+    # -- prefetch ------------------------------------------------------------
+    def _stage(self, idx):
+        BT = self.B * self.T
+        stride = BT * self.world_size
+        if self.ntok < BT + 1:
+            raise RuntimeError(f"shard too small for one batch: {self.shards[self.cur]}")
+        if self.pos + BT + 1 > self.ntok:
+            self._advance_shard()
+
+        # Host-side WAR guard: the pinned buffer may still be the source of the
+        # H2D issued the last time this slot was filled (that copy is gated on
+        # the device, not the host, so a host running ahead of the GPU would
+        # otherwise overwrite bytes still in flight and hand the model a batch
+        # spliced from two different windows).
+        if self.h2d_recorded[idx]:
+            self.h2d_done[idx].synchronize()
+
+        self.where[idx] = (self.cur, self.pos)
+        buf = np.asarray(self.tok[self.pos: self.pos + BT + 1]).astype(np.int64)
+        self.pin_x[idx].copy_(torch.from_numpy(buf[:-1]))
+        self.pin_y[idx].copy_(torch.from_numpy(buf[1:]))
+
+        with torch.cuda.stream(self.copy_stream):
+            self.dev_x[idx].view(-1).copy_(self.pin_x[idx], non_blocking=True)
+            self.dev_y[idx].view(-1).copy_(self.pin_y[idx], non_blocking=True)
+            self.h2d_done[idx].record(self.copy_stream)
+        self.h2d_recorded[idx] = True
+
+        self.pos += stride
+        if self.pos + BT + 1 > self.ntok:
+            self._advance_shard()
+
+    def next_batch(self):
+        """-> (x, y, shard_index, token_offset), x/y int64 [B,T] on device."""
+        if not self.primed:
+            for i in range(self.NBUF):
+                self._stage(i)
+            self.primed = True
+            self.buf = 0
+        else:
+            # Refill the slot returned LAST call. Its consumer's kernels are
+            # already queued on the compute stream, so the producer must wait
+            # for them before overwriting the device tensor.
+            self.copy_stream.wait_stream(torch.cuda.current_stream())
+            self._stage((self.buf - 1) % self.NBUF)
+
+        torch.cuda.current_stream().wait_event(self.h2d_done[self.buf])
+        i = self.buf
+        self.buf = (self.buf + 1) % self.NBUF
+        return self.dev_x[i], self.dev_y[i], self.where[i][0], self.where[i][1]
+
+    def skip_batches(self, n):
+        """Fast-forward `n` batches without touching the GPU (resume path)."""
+        self.copy_stream.synchronize()
+        self.primed = False
+        BT = self.B * self.T
+        stride = BT * self.world_size
+        for _ in range(n):
+            if self.pos + BT + 1 > self.ntok:
+                self._advance_shard()
+            self.pos += stride
+            if self.pos + BT + 1 > self.ntok:
+                self._advance_shard()
+
+
+# =============================================================================
+# 4. Model.
+#
+# Precision (bf16 mixed precision; both --precision modes run under autocast):
+#   bf16-params : Linear weights stored bf16, embedding / RMSNorm / QK-norm
+#                 scales stay fp32, fp32 master weights live in the optimizer
+#                 — the exact scheme bluscript_single.cpp uses (default).
+#   bf16        : all params fp32, autocast casts per-op (standard PyTorch AMP).
+# In both cases: matmuls + attention bf16, RMSNorm / QK-norm reductions, RoPE,
+# softmax and the loss in fp32. No GradScaler (bf16 has fp32's exponent range).
+# =============================================================================
+_FLASH = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm over d_model, fp32 weight. F.rms_norm is a fused ATen op and is
+    not autocast-eligible, so it runs in the input dtype — fp32 on the residual
+    stream, matching the C++ nn::RMSNorm."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        self.dim = (dim,)
+
+    def forward(self, x):
+        return F.rms_norm(x, self.dim, self.weight, self.eps)
+
+
+class CausalGQA(nn.Module):
+    """RMSNorm -> fused(QK-norm + RoPE) -> causal grouped-query attention -> Wo -> +x"""
+
+    def __init__(self, cfg, resid_scale):
+        super().__init__()
+        self.cfg = cfg
+        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
+        self.wQ = nn.Linear(cfg.d_model, cfg.q_heads * cfg.head_dim, bias=False)
+        self.wK = nn.Linear(cfg.d_model, cfg.kv_heads * cfg.head_dim, bias=False)
+        self.wV = nn.Linear(cfg.d_model, cfg.kv_heads * cfg.head_dim, bias=False)
+        self.wO = nn.Linear(cfg.q_heads * cfg.head_dim, cfg.d_model, bias=False)
+        # QK-norm scales, initialised to 1.0 (identity RMS scale), fp32.
+        self.q_gamma = nn.Parameter(torch.ones(cfg.head_dim))
+        self.k_gamma = nn.Parameter(torch.ones(cfg.head_dim))
+        # GPT-2/Llama init: projections std=0.02, residual proj scaled by
+        # (2*n_layers)^-0.5 to keep residual-stream variance ~constant.
+        for lin in (self.wQ, self.wK, self.wV):
+            nn.init.normal_(lin.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.wO.weight, mean=0.0, std=0.02 * resid_scale)
+
+    def forward(self, x, cos, sin):
+        cfg = self.cfg
+        B, T, _ = x.shape
+        h = self.norm(x)
+
+        # [B,T,H*D] -> [B,H,T,D]. The transpose stays a view; the fused kernel
+        # honours arbitrary (b,h,t) strides, so no .contiguous() is paid here.
+        q = self.wQ(h).view(B, T, cfg.q_heads, cfg.head_dim).transpose(1, 2)
+        k = self.wK(h).view(B, T, cfg.kv_heads, cfg.head_dim).transpose(1, 2)
+        v = self.wV(h).view(B, T, cfg.kv_heads, cfg.head_dim).transpose(1, 2)
+
+        if cfg.fused_rms_rope:
+            q, k = fused_rms_rope(q, k, self.q_gamma, self.k_gamma, cos, sin, cfg.qk_eps)
+        else:
+            q = rms_rope_eager(q, self.q_gamma, cos, sin, cfg.qk_eps)
+            k = rms_rope_eager(k, self.k_gamma, cos, sin, cfg.qk_eps)
+
+        # Separate fused flash kernel. enable_gqa maps query head i to kv head
+        # i // (q_heads/kv_heads), the same grouping as the C++ kernel.
+        with sdpa_kernel(_FLASH):
+            o = F.scaled_dot_product_attention(
+                q.to(v.dtype), k.to(v.dtype), v, is_causal=True, enable_gqa=True)
+
+        o = o.transpose(1, 2).reshape(B, T, cfg.q_heads * cfg.head_dim)
+        return x + self.wO(o)
+
+
+class SwiGLUMLP(nn.Module):
+    """RMSNorm -> gate_up -> silu(gate)*up -> down -> +x   ([gate | up] layout)"""
+
+    def __init__(self, cfg, resid_scale):
+        super().__init__()
+        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
+        self.gate_up = nn.Linear(cfg.d_model, 2 * cfg.ffn_hidden, bias=False)
+        self.down = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)
+        nn.init.normal_(self.gate_up.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.down.weight, mean=0.0, std=0.02 * resid_scale)
+
+    def forward(self, x):
+        gate, up = self.gate_up(self.norm(x)).chunk(2, dim=-1)
+        return x + self.down(F.silu(gate) * up)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg, resid_scale):
+        super().__init__()
+        self.attn = CausalGQA(cfg, resid_scale)
+        self.mlp = SwiGLUMLP(cfg, resid_scale)
+
+    def forward(self, x, cos, sin):
+        return self.mlp(self.attn(x, cos, sin))
+
+
+class BLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        resid_scale = 1.0 / math.sqrt(2.0 * cfg.n_layers)
+
+        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        nn.init.normal_(self.wte.weight, mean=0.0, std=0.02)
+
+        self.blocks = nn.ModuleList([Block(cfg, resid_scale) for _ in range(cfg.n_layers)])
+        self.norm_f = RMSNorm(cfg.d_model, cfg.rms_eps)
+        if not cfg.weight_tying:
+            self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+            nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02)
+
+        cos, sin = build_rope_cache(cfg.context_length, cfg.head_dim, cfg.rope_theta, "cpu")
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def apply_param_dtype(self, mode):
+        """`bf16-params`: store Linear weights in bf16 (embedding, RMSNorm and
+        QK-norm scales stay fp32) — the two-dtype-group layout of the C++ model.
+        Under autocast this is transparent: the cast to bf16 becomes a no-op."""
+        if mode != "bf16-params":
+            return
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.weight.data = m.weight.data.to(torch.bfloat16)
+
+    def forward(self, idx, targets=None):
+        cfg = self.cfg
+        T = idx.shape[1]
+        cos, sin = self.rope_cos[:T], self.rope_sin[:T]
+
+        x = self.wte(idx)                                    # fp32 [B,T,d_model]
+        for blk in self.blocks:
+            if cfg.activation_checkpoint and self.training:
+                x = checkpoint(blk, x, cos, sin, use_reentrant=False)
+            else:
+                x = blk(x, cos, sin)
+        x = self.norm_f(x)
+
+        head_w = self.wte.weight if cfg.weight_tying else self.lm_head.weight
+        if targets is None:
+            return F.linear(x, head_w), None
+        if cfg.loss_chunks > 1:
+            return None, _chunked_ce(x, head_w, targets, cfg.loss_chunks)
+        # Matches the C++: bf16 matmul, logits promoted to fp32 for the loss
+        # (cross_entropy is an autocast-fp32 op, so the promotion is implicit).
+        logits = F.linear(x, head_w)
+        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1))
+        return logits, loss
+
+
+def _chunked_ce(x, weight, targets, chunks):
+    """lm_head + cross-entropy in row chunks, logits recomputed in backward.
+    Peak logit memory drops from [B*T, V] to [B*T/chunks, V] at the cost of one
+    extra lm_head matmul. Off by default (`--loss-chunks 1`) so the memory
+    profile matches the C++ script, which materialises the full fp32 logits.
+    """
+    xf = x.reshape(-1, x.shape[-1])
+    tf = targets.reshape(-1)
+    n = xf.shape[0]
+    size = (n + chunks - 1) // chunks
+    total = xf.new_zeros((), dtype=torch.float32)
+    for lo in range(0, n, size):
+        hi = min(lo + size, n)
+        part = checkpoint(_ce_chunk, xf[lo:hi], weight, tf[lo:hi], use_reentrant=False)
+        total = total + part * (hi - lo)
+    return total / n
+
+
+def _ce_chunk(x, weight, targets):
+    return F.cross_entropy(F.linear(x, weight).float(), targets)
+
+
+# =============================================================================
+# 5. Optimizer (Muon on the hidden matrices, AdamW on the rest), gradient
+#    clipping, checkpointing.
+# =============================================================================
+def _is_muon_param(name, p):
+    """Muon orthogonalizes a matrix, so it needs the parameter to BE a linear
+    map between two feature spaces:
+      * dim >= 2 — RMSNorm weights and QK-norm gammas are 1D and have no
+        singular values to equalise;
+      * not the embedding / lm_head — its rows index tokens, not features, so
+        the left singular vectors mix unrelated vocabulary entries. Weight
+        tying makes this one tensor serving both ends; either way it is AdamW's.
+    Everything else — wQ/wK/wV/wO, gate_up, down — is Muon's: 75.5M of the
+    114.2M parameters at the default config.
+    """
+    return p.dim() >= 2 and not name.startswith(("wte.", "lm_head."))
+
+
+def _state_bytes(opt):
+    """Bytes held in a torch optimizer's per-parameter state — Adam's m/v, or
+    Muon's momentum buffer. Reads 0 before the first step(): both allocate
+    lazily."""
+    return sum(v.numel() * v.element_size()
+               for st in opt.state.values()
+               for v in st.values() if torch.is_tensor(v))
+
+
+class MasterMuon:
+    """torch.optim.Muon on the hidden matrices + torch.optim.AdamW on the rest,
+    both driving fp32 master weights.
+    Neither optimizer is reimplemented here. What this class owns is the part
+    torch cannot know about:
+      * ROUTING — torch.optim.Muon accepts whatever you hand it and the docs are
+        explicit that only 2D hidden weights belong ("bias, and embedding,
+        should be optimized by a standard method such as AdamW"). The split is
+        `_is_muon_param`.
+      * MASTER WEIGHTS — the parameters are bf16 under --precision bf16-params,
+        and a bf16 tensor cannot absorb an lr*update of 1e-4 relative size. So
+        both optimizers are constructed over the fp32 MASTERS, fed from the fp32
+        landing buffers, and the masters are cast back over the params after the
+        step. Unchanged from the AdamW arm — that machinery is what makes the
+        two arms comparable, since both optimizers then read a bit-identical
+        fp32 gradient and write a bit-identical bf16 parameter.
+    Weight decay is decoupled in both (torch.optim.Muon applies it before the
+    update, as AdamW does) and is applied to the master for the same reason: a
+    decay written to the bf16 param is overwritten by the master write-back
+    microseconds later and silently no-ops.
+    State cost: Muon holds one momentum buffer per param, against Adam's two.
+    """
+
+    def __init__(self, named_params, cfg):
+        self.cfg = cfg
+        self.triples = []                 # (param, master, landing_grad), ALL params
+        self.muon_masters = []            # masters routed to Muon
+        decay, nodecay = [], []
+        self.n_muon = self.n_decay = self.n_nodecay = 0
+        for name, p in named_params:
+            if p.dtype == torch.float32:
+                m, g = p, None
+            else:
+                m = p.detach().float().clone()
+                g = torch.zeros_like(m)
+                m.grad = g
+            self.triples.append((p, m, g))
+            if _is_muon_param(name, p):
+                self.muon_masters.append(m)
+                self.n_muon += m.numel()
+            elif p.dim() >= 2:
+                decay.append(m)
+                self.n_decay += m.numel()
+            else:
+                nodecay.append(m)
+                self.n_nodecay += m.numel()
+        self.opt = torch.optim.AdamW(
+            [{"params": decay, "weight_decay": cfg.weight_decay},
+             {"params": nodecay, "weight_decay": 0.0}],
+            lr=cfg.max_lr, betas=(0.9, 0.95), eps=1e-8, fused=True)
+        # Everything below is a torch.optim.Muon default except adjust_lr_fn
+        # (torch: "original") and weight_decay, which is pinned to the AdamW
+        # arm's 0.1 so the decay is identical across both groups and both arms.
+        self.muon = torch.optim.Muon(
+            self.muon_masters,
+            lr=cfg.muon_lr,
+            weight_decay=cfg.weight_decay,
+            momentum=cfg.muon_momentum,
+            nesterov=True,
+            ns_steps=cfg.muon_ns_steps,
+            adjust_lr_fn=cfg.muon_adjust_lr)
+
+    @torch.no_grad()
+    def step(self, lr, muon_lr):
+        for p, m, g in self.triples:
+            if g is not None:
+                g.copy_(p.grad)               # bf16 grad -> fp32 landing buffer
+        # AdamW group (embedding / lm_head / norms / gammas) — identical to the
+        # baseline's update, on identical gradients.
+        for grp in self.opt.param_groups:
+            grp["lr"] = lr
+        self.opt.step()
+        # Muon group. Its nesterov path may consume the gradient in place; that
+        # is harmless here because grad_norm and the --diag probes are read
+        # before step() and the landing buffers are overwritten wholesale on the
+        # next one.
+        for grp in self.muon.param_groups:
+            grp["lr"] = muon_lr
+        self.muon.step()
+        for p, m, g in self.triples:
+            if g is not None:
+                p.data.copy_(m)               # fp32 master -> bf16 param
+
+    def state_breakdown(self):
+        """-> (master_bytes, landing_grad_bytes, adam_mv_bytes, muon_mom_bytes).
+        The first two are zero under --precision bf16, where the master IS the
+        parameter and the gradient is already fp32. Both optimizers allocate
+        their state lazily on the first step(), so the last two read 0 until
+        then — the startup report uses the expected_* predictors instead.
+        """
+        master = landing = 0
+        for p, m, g in self.triples:
+            if m is not p:
+                master += m.numel() * m.element_size()
+                landing += g.numel() * g.element_size()
+        mv = _state_bytes(self.opt)
+        mom = _state_bytes(self.muon)
+        return master, landing, mv, mom
+
+    def state_bytes(self):
+        return sum(self.state_breakdown())
+
+    def expected_mv_bytes(self):
+        """Adam m/v are lazily allocated, so the startup report has to predict
+        them rather than measure them. Muon's params are excluded — they carry a
+        momentum buffer instead, reported separately."""
+        muon_ids = {id(m) for m in self.muon_masters}
+        return sum(m.numel() * m.element_size() * 2
+                   for _, m, _ in self.triples if id(m) not in muon_ids)
+
+    def expected_muon_mom_bytes(self):
+        """One fp32 momentum buffer per Muon parameter — the prediction the
+        startup report prints before step 0 has allocated it."""
+        return sum(m.numel() * m.element_size() for m in self.muon_masters)
+
+    # -- checkpointing -------------------------------------------------------
+    def state_dict(self):
+        # `optimizer` tags the arm: resuming a Muon run from an AdamW checkpoint
+        # would restore the weights and then silently drop the momentum, so
+        # CheckpointManager refuses on a mismatch rather than producing a run
+        # that looks fine and is not comparable to either arm.
+        return {"optimizer": "muon",
+                "adam": self.opt.state_dict(),
+                "muon": self.muon.state_dict(),
+                "masters": [None if m is p else m for p, m, g in self.triples]}
+
+    @torch.no_grad()
+    def load_state_dict(self, blob):
+        """-> True if the fp32 masters were restored from the checkpoint."""
+        if blob.get("optimizer") != "muon":
+            raise SystemExit(
+                "ERROR: this checkpoint was written by the AdamW arm (train.py). "
+                "Point --ckpt-dir somewhere else, or resume it with train.py.")
+        self.opt.load_state_dict(blob["adam"])
+        if blob.get("muon") is not None:
+            self.muon.load_state_dict(blob["muon"])
+        masters = blob.get("masters")
+        if masters is None or len(masters) != len(self.triples):
+            return False
+        for (p, m, g), saved in zip(self.triples, masters):
+            if saved is not None and m is not p:
+                m.copy_(saved.to(m.device))
+        return True
+
+    @torch.no_grad()
+    def resync_masters(self):
+        """Re-seed every master from its parameter.
+        Needed after loading a checkpoint that carries no masters: the masters
+        were built in __init__ from the RANDOM INIT, so without this the first
+        step() would write them straight back over the restored weights and
+        silently throw the run away.
+        """
+        for p, m, g in self.triples:
+            if m is not p:
+                m.copy_(p.detach().float())
+
+
+def clip_grads_(params, max_norm):
+    """ONE global L2 norm over every parameter, then one rescale — the textbook
+    definition. (It differs from the ZeRO script, which clips each dtype group
+    against its own norm and can therefore let the true global norm reach
+    sqrt(n_groups) * grad_clip.) The per-tensor norms are promoted to fp32
+    before they are combined, or a reduction over ~250 bf16 norms would lose
+    ~3 digits."""
+    grads = [p.grad for p in params if p.grad is not None]
+    norms = torch._foreach_norm(grads)
+    total = torch.linalg.vector_norm(torch.stack([n.float() for n in norms]))
+    if max_norm and max_norm > 0:
+        clip = (max_norm / (total + 1e-6)).clamp(max=1.0)
+        torch._foreach_mul_(grads, clip)
+    return total
+
+
+class CheckpointManager:
+    """Synchronous save / keep-N / auto-resume, mirroring the C++
+    CheckpointManager. A checkpoint holds the model parameters, Adam's m/v, the
+    fp32 masters, and the step it should resume at."""
+
+    def __init__(self, ckpt_dir, prefix, keep):
+        self.dir, self.prefix, self.keep = ckpt_dir, prefix, keep
+        os.makedirs(ckpt_dir, exist_ok=True)
+
+    def _path(self, step):
+        return os.path.join(self.dir, f"{self.prefix}_step{step:08d}.pt")
+
+    def _existing(self):
+        out = []
+        for p in glob.glob(os.path.join(self.dir, f"{self.prefix}_step*.pt")):
+            m = re.search(r"_step(\d+)\.pt$", os.path.basename(p))
+            if m:
+                out.append((int(m.group(1)), p))
+        return sorted(out)
+
+    def save(self, step, model, opt, loss, last_val, cfg):
+        tmp = self._path(step) + ".tmp"
+        torch.save({"step": step, "loss": loss, "last_val": last_val,
+                    # B/T/global_batch pin the data stride this checkpoint's
+                    # step count assumes. world_size is deliberately NOT
+                    # recorded: grad_accum = global_batch // (B*T*world_size)
+                    # already makes tokens-per-step (and skip_batches'
+                    # target offset) independent of how many ranks are
+                    # running, so a GPU-count change alone is not a mismatch.
+                    "run_cfg": {"B": cfg.B, "T": cfg.T,
+                                "global_batch": cfg.global_batch},
+                    "model": model.state_dict(),
+                    "opt": opt.state_dict()}, tmp)
+        # Rename last: a crash mid-write leaves a .tmp, never a truncated
+        # checkpoint that the next load_latest would pick as the newest.
+        os.replace(tmp, self._path(step))
+        if self.keep > 0:
+            for _, stale in self._existing()[:-self.keep]:
+                os.remove(stale)
+
+    def load_latest(self, model, opt, cfg, master=True):
+        """-> (start_step, loss, last_val). (0, 0.0, nan) when there is
+        nothing to resume."""
+        existing = self._existing()
+        if not existing:
+            return 0, 0.0, float("nan")
+        _, path = existing[-1]
+        return self.load(path, model, opt, cfg, master)
+
+    def load(self, path, model, opt, cfg, master=True):
+        """Load a specific checkpoint file (e.g. from --resume-from), rather
+        than auto-discovering the latest one in self.dir.
+        -> (start_step, loss, last_val)."""
+        blob = torch.load(path, map_location="cpu", weights_only=True)
+        model.load_state_dict(blob["model"])
+        if not opt.load_state_dict(blob["opt"]):
+            opt.resync_masters()
+        run_cfg = blob.get("run_cfg")
+        if master and run_cfg is not None:
+            mismatches = [f"{k} checkpoint={v!r} current={getattr(cfg, k)!r}"
+                         for k, v in run_cfg.items() if getattr(cfg, k) != v]
+            if mismatches:
+                print("[WARN] resume config differs from the checkpoint's — "
+                      "the data offset this step count assumes will be wrong:")
+                for m in mismatches:
+                    print(f"         {m}")
+        print(f"[resume] loaded {path}")
+        return (int(blob["step"]), float(blob["loss"]),
+                float(blob.get("last_val", float("nan"))))
+
+
+# =============================================================================
+# 6. Training.
+# =============================================================================
+def _env_int(name, default=0):
+    v = os.environ.get(name)
+    if v in (None, ""):
+        return default
+    try:
+        return int(v)
+    except ValueError:
+        print(f"[WARN] ignoring non-integer {name}={v!r}")
+        return default
+
+
+def seed_everything(seed):
+    """Every RNG the run can draw from, off one seed.
+    torch's CPU generator is the one that matters here (the model is built on
+    CPU and moved), but `random` / numpy are seeded too so that anything added
+    later — a shuffled shard order, a sampled probe — is reproducible by
+    default rather than by accident.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def apply_determinism(mode):
+    """Pin the run-to-run drift that seeding does not cover.
+    Three sources, all on the backward pass:
+      * nn.Embedding's dense backward scatters one atomicAdd per token into
+        [vocab, d_model]; with weight tying that is also the lm_head gradient,
+        so it is the largest reduction in the step;
+      * the flash / mem-efficient attention backward splits dq over the key
+        blocks and atomicAdds them together;
+      * cuDNN picks kernels by wall-clock timing when `benchmark` is on, so the
+        chosen algorithm can differ between runs of the same shape.
+    Each is a float reduction in an order the scheduler chooses, so two runs
+    diverge in the last mantissa bits of the first step and, through the
+    optimizer, visibly by a few hundred steps.
+    `strict` raises if some op has no deterministic kernel — better than
+    reporting a reproducible run that is not one. `warn` downgrades that to a
+    warning and keeps training; `off` restores the default fast kernels.
+    """
+    if mode == "off":
+        return
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    # Routes embedding-backward and SDPA-backward to their ordered variants, and
+    # enforces the cuBLAS workspace setting made at the top of this file.
+    torch.use_deterministic_algorithms(True, warn_only=(mode == "warn"))
+
+
+def build_argparser():
+    ap = argparse.ArgumentParser(
+        description="BLU single-GPU training — PyTorch port of bluscript_single.cpp")
+    ap.add_argument("--data-root",
+                    default=("/root/.cache/huggingface/hub/datasets--navingv--flux/"
+                             "snapshots/5232b3d133d2eb71905725f288ff6b56dc0df773"))
+    ap.add_argument("--device", type=int, default=None,
+                    help="CUDA device index (env: BLU_DEVICE)")
+    ap.add_argument("--precision", choices=["bf16-params", "bf16"], default="bf16-params",
+                    help="bf16-params (default): Linear weights stored bf16 with "
+                         "fp32 master weights in the optimizer — the C++ scheme. "
+                         "bf16: fp32 params + autocast (standard PyTorch AMP).")
+    ap.add_argument("--rms-rope", choices=["fused", "eager"], default="fused",
+                    help="fused = Triton QK-norm+RoPE kernel (default), "
+                         "eager = unfused reference, for A/B on the kernel")
+    ap.add_argument("--compile", action="store_true", help="torch.compile each block")
+    ap.add_argument("--activation-checkpoint", action="store_true",
+                    help="recompute each block in backward (trades ~30%% time for memory)")
+    ap.add_argument("--loss-chunks", type=int, default=1,
+                    help="chunk lm_head+CE into N pieces with recompute; 1 (default) "
+                         "materialises the full logits, as the C++ does")
+    # bluscript_single.cpp defaults unless overridden.
+    ap.add_argument("--B", type=int, default=None, help="micro-batch (env: BLU_B)")
+    ap.add_argument("--T", type=int, default=None)
+    ap.add_argument("--d-model", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--n-layers", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--q-heads", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--kv-heads", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--ffn-hidden", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--vocab-size", type=int, default=None, help="smoke-test override")
+    ap.add_argument("--global-batch", type=int, default=None,
+                    help="tokens per optimizer step (env: BLU_GLOBAL_BATCH)")
+    ap.add_argument("--max-steps", type=int, default=None,
+                    help="caps the run for benchmarking; warmup/stable shrink "
+                         "proportionally so the WSD decay tail stays non-empty "
+                         "(env: BLU_MAX_STEPS)")
+    # ---- Muon ----
+    ap.add_argument("--muon-lr", type=float, default=None,
+                    help="peak LR for the hidden 2D weights. Under the default "
+                         "--muon-adjust-lr match_rms_adamw this is the same unit as "
+                         "--max-lr, and defaults to the same value (2e-4) so "
+                         "both arms take a matched first run; sweep it upward "
+                         "from there (Muon usually takes 2-4x AdamW's LR).")
+    ap.add_argument("--muon-min-lr", type=float, default=None,
+                    help="floor of the Muon cosine tail (default 10%% of --muon-lr)")
+    ap.add_argument("--muon-momentum", type=float, default=None,
+                    help="torch.optim.Muon momentum (default 0.95, torch's own)")
+    ap.add_argument("--muon-ns-steps", type=int, default=None,
+                    help="torch.optim.Muon ns_steps, Newton-Schulz iterations "
+                         "(default 5, torch's own); fewer = looser "
+                         "orthogonalization, cheaper")
+    ap.add_argument("--muon-adjust-lr", choices=["match_rms_adamw", "original"],
+                    default=None,
+                    help="torch.optim.Muon adjust_lr_fn. match_rms_adamw "
+                         "(default HERE): 0.2*sqrt(max(A,B)), putting --muon-lr "
+                         "on the same scale as --max-lr. original (default in "
+                         "torch): sqrt(max(1,A/B)), needs a ~5.5x larger "
+                         "--muon-lr at d_model=768.")
+    ap.add_argument("--val-freq", type=int, default=None)
+    ap.add_argument("--val-loss-steps", type=int, default=None)
+    ap.add_argument("--diag", action="store_true",
+                    help="per-micro-batch loss + depth probes (env: BLU_DIAG=1)")
+    ap.add_argument("--warmup-measure", type=int, default=3,
+                    help="steps excluded from the steady-state throughput summary")
+    ap.add_argument("--csv", default="blu_muon_log.csv",
+                    help="CSV log path; appended (not truncated) on resume. The "
+                         "first 8 columns match the AdamW arm's CSV exactly, so "
+                         "the two files plot against each other directly; "
+                         "muon_lr is appended as a 9th.")
+    ap.add_argument("--seed", type=int, default=1234)
+    ap.add_argument("--determinism", choices=["strict", "warn", "off"], default="strict",
+                    help="strict (default): deterministic CUDA kernels, so a rerun "
+                         "of the same launch line reproduces the loss curve bit for "
+                         "bit; raises if an op has no deterministic implementation. "
+                         "warn: same, but non-deterministic ops warn instead of "
+                         "raising. off: fastest kernels, results drift between runs.")
+    # checkpointing
+    ap.add_argument("--no-checkpointing", action="store_true")
+    ap.add_argument("--ckpt-dir", default=None)
+    ap.add_argument("--ckpt-prefix", default=None)
+    ap.add_argument("--ckpt-freq", type=int, default=None)
+    ap.add_argument("--ckpt-keep", type=int, default=None)
+    ap.add_argument("--resume-from", default=None,
+                    help="path to a specific *_step*.pt checkpoint to resume "
+                         "from, overriding auto-discovery of the latest "
+                         "checkpoint in --ckpt-dir")
+    return ap
+
+
+def config_from_args(args):
+    cfg = ModelConfig()
+    for f in ("B", "T", "d_model", "n_layers", "q_heads", "kv_heads", "ffn_hidden",
+              "vocab_size", "global_batch", "val_freq", "val_loss_steps",
+              "ckpt_dir", "ckpt_prefix", "ckpt_freq", "ckpt_keep",
+              "muon_lr", "muon_min_lr", "muon_momentum", "muon_ns_steps",
+              "muon_adjust_lr"):
+        v = getattr(args, f, None)
+        if v is not None:
+            setattr(cfg, f, v)
+    # A --muon-lr given without a floor keeps the 10% ratio the WSD tail assumes,
+    # instead of silently decaying to the default 2e-5 from wherever it starts.
+    if args.muon_lr is not None and args.muon_min_lr is None:
+        cfg.muon_min_lr = 0.1 * cfg.muon_lr
+    cfg.fused_rms_rope = (args.rms_rope == "fused")
+    cfg.activation_checkpoint = args.activation_checkpoint
+    cfg.loss_chunks = args.loss_chunks
+    cfg.checkpointing = not args.no_checkpointing
+    cfg.context_length = max(cfg.context_length, cfg.T)
+
+    # Mirrors the BLU_MAX_STEPS branch in bluscript_single.cpp: shrink warmup and
+    # stable with the run, or a short benchmark leaves no decay tail (the C++
+    # asserts on it; get_lr here would clamp, which is worse — it would silently
+    # report a schedule the config does not describe).
+    if args.max_steps is not None:
+        ms = args.max_steps
+        cfg.max_steps = ms
+        cfg.warmup_steps = max(1, ms // 10)
+        cfg.stable_steps = max(1, ms // 2)
+        if cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1:
+            cfg.stable_steps = max(1, cfg.max_steps - cfg.warmup_steps - 2)
+        print(f"[bench] max_steps={ms} (warmup={cfg.warmup_steps} "
+              f"stable={cfg.stable_steps})")
+    return cfg
+
+
+def validate_config(cfg, world_size=1, master=True):
+    """Cheap checks; each one is a silent-corruption or out-of-bounds failure
+    mode if it does not hold."""
+    bad = []
+    if cfg.T > cfg.context_length:
+        bad.append(f"T ({cfg.T}) > context_length ({cfg.context_length}): "
+                   "the RoPE cache is too short")
+    if cfg.q_heads % cfg.kv_heads != 0:
+        bad.append(f"q_heads ({cfg.q_heads}) not divisible by kv_heads ({cfg.kv_heads})")
+    if cfg.head_dim not in (64, 128):
+        bad.append(f"the fused kernels support head_dim 64 or 128, got {cfg.head_dim}")
+    # Under DDP each rank runs B*T tokens per micro-step, so the world's tokens
+    # per micro-step are B*T*world_size; global_batch must land on a whole
+    # number of those or grad_accum below truncates and quietly under-shoots
+    # the configured optimizer-step token count.
+    if cfg.global_batch % (cfg.B * cfg.T * world_size) != 0:
+        bad.append(f"global_batch ({cfg.global_batch}) not divisible by "
+                   f"B*T*world_size ({cfg.B * cfg.T * world_size})")
+    if cfg.warmup_steps < 1:
+        bad.append("warmup_steps must be >= 1 (get_lr divides by it)")
+    if cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1:
+        bad.append(f"warmup_steps + stable_steps "
+                   f"({cfg.warmup_steps + cfg.stable_steps}) leaves no decay tail "
+                   f"before max_steps ({cfg.max_steps})")
+    if not hasattr(torch.optim, "Muon"):
+        bad.append(f"torch {torch.__version__} has no torch.optim.Muon — it first "
+                   "appears in the 2.9 docs. Upgrade torch, or run the AdamW "
+                   "arm (train.py).")
+    if cfg.muon_ns_steps < 1:
+        bad.append("muon_ns_steps must be >= 1 (0 leaves the update unorthogonalized, "
+                   "which is not Muon)")
+    if cfg.muon_momentum <= 0.0 or cfg.muon_momentum >= 1.0:
+        bad.append(f"muon_momentum must be in (0,1), got {cfg.muon_momentum}")
+    if cfg.q_heads * cfg.head_dim != cfg.d_model and master:
+        print(f"WARNING: q_heads*head_dim ({cfg.q_heads * cfg.head_dim}) "
+              f"!= d_model ({cfg.d_model})")
+    if bad:
+        raise SystemExit("ERROR: " + "\nERROR: ".join(bad))
+
+
+def main():
+    args = build_argparser().parse_args()
+
+    # Env overrides, for parity with the C++ launch line. An explicit CLI flag
+    # wins over the environment.
+    if args.B is None:
+        args.B = _env_int("BLU_B", 0) or None
+    if args.global_batch is None:
+        args.global_batch = _env_int("BLU_GLOBAL_BATCH", 0) or None
+    if args.max_steps is None:
+        args.max_steps = _env_int("BLU_MAX_STEPS", 0) or None
+    diag = args.diag or _env_int("BLU_DIAG", 0) != 0
+
+    if not torch.cuda.is_available():
+        raise SystemExit("ERROR: no CUDA device available")
+
+    # ---- DDP bootstrap ----
+    # torchrun sets RANK/LOCAL_RANK/WORLD_SIZE; their absence means a plain
+    # single-process launch, where --device / BLU_DEVICE pick the GPU instead.
+    # torch.cuda.set_device MUST happen before init_process_group: NCCL creates
+    # its context on whatever device is current at that call, so doing it after
+    # would leave every rank on this node contending for device 0.
+    ddp = int(os.environ.get("RANK", -1)) != -1
+    if ddp:
+        ddp_rank = int(os.environ["RANK"])
+        ddp_local_rank = int(os.environ["LOCAL_RANK"])
+        ddp_world_size = int(os.environ["WORLD_SIZE"])
+        master_process = ddp_rank == 0
+        if ddp_local_rank >= torch.cuda.device_count():
+            raise SystemExit(f"ERROR: LOCAL_RANK={ddp_local_rank} but only "
+                             f"{torch.cuda.device_count()} CUDA device(s) visible")
+        torch.cuda.set_device(ddp_local_rank)
+        device = torch.device("cuda", ddp_local_rank)
+        dist.init_process_group(backend="nccl")
+        if master_process and (args.device is not None
+                               or os.environ.get("BLU_DEVICE") not in (None, "")):
+            print("[WARN] --device / BLU_DEVICE is ignored under torchrun — "
+                  "LOCAL_RANK selects each rank's GPU.")
+    else:
+        ddp_rank, ddp_local_rank, ddp_world_size = 0, 0, 1
+        master_process = True
+        if args.device is None:
+            args.device = _env_int("BLU_DEVICE", 0)
+        if args.device >= torch.cuda.device_count():
+            raise SystemExit(f"ERROR: BLU_DEVICE={args.device} but only "
+                             f"{torch.cuda.device_count()} CUDA device(s) visible")
+        torch.cuda.set_device(args.device)
+        device = torch.device("cuda", args.device)
+
+    seed_everything(args.seed)
+    apply_determinism(args.determinism)
+    # TF32 loses mantissa bits but is deterministic, so it is orthogonal to the
+    # setting above and stays on.
+    torch.set_float32_matmul_precision("high")
+
+    cfg = config_from_args(args)
+    if cfg.fused_rms_rope and not _HAS_TRITON:
+        if master_process:
+            print("[WARN] Triton not available — falling back to the eager QK-norm+RoPE")
+        cfg.fused_rms_rope = False
+    validate_config(cfg, ddp_world_size, master_process)
+
+    # Tokens consumed per optimizer step = B * T * grad_accum * world_size.
+    grad_accum = cfg.global_batch // (cfg.B * cfg.T * ddp_world_size)
+    tokens_per_step = cfg.B * cfg.T * grad_accum * ddp_world_size
+
+    if master_process:
+        print("================ Configuration ================")
+        mode_str = (f"DDP x{ddp_world_size} (torchrun, NCCL)" if ddp
+                   else "SINGLE GPU (no torchrun / DDP / ZeRO)")
+        print(f"  mode            : {mode_str}")
+        print(f"  cuda device     : {ddp_local_rank if ddp else args.device} "
+              f"({torch.cuda.get_device_name(device)})")
+        print(f"  d_model         : {cfg.d_model}")
+        print(f"  n_layers        : {cfg.n_layers}")
+        print(f"  q_heads/kv_heads: {cfg.q_heads} / {cfg.kv_heads} "
+              f"(G={cfg.q_heads // cfg.kv_heads})")
+        print(f"  head_dim        : {cfg.head_dim}")
+        print(f"  ffn_hidden      : {cfg.ffn_hidden}  (SwiGLU)")
+        print(f"  context_length  : {cfg.context_length}")
+        print(f"  rope_theta      : {cfg.rope_theta:g}  (NeoX half-split)")
+        print(f"  vocab_size      : {cfg.vocab_size}")
+        print(f"  weight_tying    : {'ON' if cfg.weight_tying else 'OFF'}")
+        print(f"  B x T           : {cfg.B} x {cfg.T}")
+        print(f"  global_batch    : {cfg.global_batch}  (grad_accum={grad_accum}, "
+              f"world_size={ddp_world_size})")
+        print("  optimizer       : MUON (hidden 2D) + AdamW (embed/head/1D)")
+        print(f"  max_lr / min_lr : {cfg.max_lr:g} / {cfg.min_lr:g}   (AdamW group)")
+        print(f"  muon lr / min   : {cfg.muon_lr:g} / {cfg.muon_min_lr:g}   "
+              f"(adjust_lr_fn={cfg.muon_adjust_lr})")
+        print(f"  muon momentum   : {cfg.muon_momentum:g}  | NS steps {cfg.muon_ns_steps}")
+        print(f"  warmup / steps  : {cfg.warmup_steps} / {cfg.max_steps}")
+        print(f"  WSD stable_steps: {cfg.stable_steps}  "
+              f"(decay tail={cfg.max_steps - cfg.warmup_steps - cfg.stable_steps})")
+        print(f"  precision       : {args.precision}")
+        print(f"  seed / determ.  : {args.seed} / {args.determinism}")
+        print(f"  rms+rope        : {'fused' if cfg.fused_rms_rope else 'eager'}"
+              f" | compile {args.compile} | act-ckpt {cfg.activation_checkpoint}"
+              f" | loss_chunks {cfg.loss_chunks}")
+        print("===============================================", flush=True)
+
+    # ---- model ----
+    # Built and seeded identically on every rank (same seed, deterministic CPU
+    # init); DDP's constructor broadcasts rank 0's parameters/buffers anyway,
+    # so this is belt-and-suspenders against any future source of divergence.
+    if master_process:
+        print(f"Initializing model on CUDA device {device.index}...", flush=True)
+    raw_model = BLU(cfg).to(device)
+    raw_model.apply_param_dtype(args.precision)
+    params = list(raw_model.parameters())
+    n_params = sum(p.numel() for p in params)
+    if master_process:
+        print(f"Number of parameters: {n_params}")
+    if args.compile:
+        for blk in raw_model.blocks:
+            blk.compile()
+
+    # DDP wraps the (already dtype-cast, already compiled) raw module. Buffer
+    # broadcasting is off: rope_cos/rope_sin are deterministic functions of
+    # cfg, identical on every rank already, so syncing them every forward
+    # would just be wasted communication.
+    # if ddp:
+    #     model = DDP(raw_model, device_ids=[ddp_local_rank], broadcast_buffers=False)
+    # else:
+    #     model = raw_model
+
+    if ddp:
+        model = DDP(raw_model, device_ids=[ddp_local_rank], broadcast_buffers=False)
+        torch.cuda.synchronize()   # let DDP's constructor-time NCCL broadcast fully land
+                                # before DataLoaderLite's independent copy_stream starts
+                                # issuing its own async H2D work — see [crash investigation]
+    else:
+        model = raw_model
+
+
+    # ---- optimizer ----
+    # named_parameters(), not parameters(): the routing is by name (the
+    # embedding and lm_head stay on AdamW). Duplicates are dropped by default,
+    # so a weight-tied head is visited once, as `wte.weight`. Built from
+    # raw_model so names never carry DDP's "module." prefix.
+    opt = MasterMuon(raw_model.named_parameters(), cfg)
+    if master_process:
+        print(f"muon params: {opt.n_muon} | adamw decayed: {opt.n_decay}, "
+              f"non-decayed: {opt.n_nodecay}")
+    # Static footprint. Nothing here is sharded — plain DDP replicates the
+    # full optimizer state on every rank (the whole cost of not running
+    # ZeRO-1), so this is the per-GPU charge before a single activation is
+    # allocated.
+    gb = 2 ** 30
+    param_bytes = sum(p.numel() * p.element_size() for p in params)
+    grad_bytes = param_bytes                       # grads share the param dtype
+    master_bytes, landing_bytes, _, _ = opt.state_breakdown()
+    # Both optimizers allocate their state on the first step(), so at this point
+    # the measured values are 0 and the report has to predict them.
+    mv_bytes = opt.expected_mv_bytes()
+    mom_bytes = opt.expected_muon_mom_bytes()
+    static_bytes = (param_bytes + grad_bytes + master_bytes + landing_bytes
+                    + mv_bytes + mom_bytes)
+    if master_process:
+        print(f"[{'ddp' if ddp else 'single'}] static footprint  (per GPU: "
+              f"params + grads + optimizer state)")
+        print(f"    parameters              : {param_bytes / gb:7.2f} GiB")
+        print(f"    gradients               : {grad_bytes / gb:7.2f} GiB")
+        print(f"    fp32 master weights     : {master_bytes / gb:7.2f} GiB")
+        print(f"    fp32 grad landing bufs  : {landing_bytes / gb:7.2f} GiB")
+        print(f"    Muon momentum           : {mom_bytes / gb:7.2f} GiB")
+        print(f"    Adam exp_avg + exp_avg_sq: {mv_bytes / gb:6.2f} GiB")
+        print(f"    static total            : {static_bytes / gb:7.2f} GiB", flush=True)
+
+    # ---- data ----
+    # Sharded by rank: each rank reads a disjoint, interleaved slice of the
+    # same shard list (see DataLoaderLite), so DDP's cross-rank gradient
+    # average is over world_size DIFFERENT windows, not world_size copies of
+    # the same one.
+    train_loader = DataLoaderLite(cfg.B, cfg.T, "train", args.data_root, device,
+                                  master=master_process, rank=ddp_rank,
+                                  world_size=ddp_world_size)
+    # rank=0, world_size=1 UNCONDITIONALLY, on every rank: val_loss_steps
+    # windows at fixed token offsets 0, BT, 2BT, ... Sharding this by
+    # ddp_rank/ddp_world_size like the train loader would make the reported
+    # val_loss a function of how many GPUs happen to be running (world_size=2
+    # averages 2x as many, differently-offset windows as world_size=1) —
+    # the two numbers would not be comparable across a GPU-count change, which
+    # is what produced the apparent val_loss jump on the 2-GPU -> 1-GPU
+    # resume. Every rank now redundantly evaluates the identical fixed
+    # windows, so the metric means the same thing regardless of world_size.
+    val_loader = DataLoaderLite(cfg.B, cfg.T, "val", args.data_root, device,
+                                master=master_process, rank=0, world_size=1)
+
+    # ---- checkpointing: manager + auto-resume ----
+    # Every rank independently loads the same checkpoint file from the same
+    # (assumed shared) filesystem — a read, so no race — which keeps model,
+    # optimizer state, and start_step identical across ranks without an
+    # explicit broadcast. Only rank 0 ever WRITES one (below).
+    start_step, ckpt = 0, None
+    last_val = float("nan")
+    if cfg.checkpointing or args.resume_from:
+        ckpt = CheckpointManager(cfg.ckpt_dir, cfg.ckpt_prefix, cfg.ckpt_keep)
+        if args.resume_from:
+            if master_process:
+                print(f"[resume] loading explicit checkpoint: {args.resume_from}")
+            start_step, resume_loss, last_val = ckpt.load(
+                args.resume_from, raw_model, opt, cfg, master_process)
+        else:
+            start_step, resume_loss, last_val = ckpt.load_latest(
+                raw_model, opt, cfg, master_process)
+        if start_step > 0:
+            if master_process:
+                print(f"[resume] at step {start_step} (loss {resume_loss:.6f}, "
+                      f"last val {last_val:.6f})")
+            # Skip the batches already consumed by steps [0, start_step).
+            # global_batch // (B*T*world_size) makes tokens-per-step (and so
+            # this target offset) independent of world_size, so a GPU-count
+            # change alone lands on the correct byte offset as long as B, T,
+            # and global_batch are unchanged (checked above).
+            train_loader.skip_batches(start_step * grad_accum)
+
+    # ---- CSV log ----
+    # Appended (not truncated) when resuming so a restart does not erase the
+    # history of the run it is continuing. Rank 0 only — every rank would
+    # otherwise open and write the same path.
+    log_file = None
+    if args.csv and master_process:
+        resuming = start_step > 0 and os.path.exists(args.csv)
+        log_file = open(args.csv, "a" if resuming else "w")
+        if not resuming:
+            log_file.write("step,train_loss,val_loss,tok_per_sec,grad_norm,dt_ms,lr,"
+                           "elapsed_min,muon_lr\n")
+        log_file.flush()
+
+    amp = torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+    probes = [L for L in cfg.probe_layers if L < cfg.n_layers]
+    steady_dts = []
+    t_train_start = time.perf_counter()
+
+    # ---- Ctrl-C: finish the in-flight step, save a checkpoint, exit clean ----
+    # Registered here (not earlier) so a Ctrl-C during setup/data-loading still
+    # gets Python's default immediate-abort behaviour — only the training loop
+    # below gets the graceful path. A second Ctrl-C re-arms the default handler
+    # and re-raises, for when the checkpoint save itself needs interrupting.
+    interrupted = False
+
+    def _handle_sigint(signum, frame):
+        nonlocal interrupted
+        if interrupted:
+            signal.signal(signal.SIGINT, signal.SIG_DFL)
+            raise KeyboardInterrupt
+        interrupted = True
+        if master_process:
+            print("\n[ctrl-c] finishing this step, then checkpointing and exiting "
+                  "(press again to abort immediately)...", flush=True)
+
+    # A terminal Ctrl-C is delivered to every process in the foreground
+    # process group, so under torchrun every rank sets `interrupted` at
+    # roughly the same time and the ranks stay in lockstep for the collective
+    # ops (DDP's grad all-reduce, the checkpoint barrier below) still ahead.
+    signal.signal(signal.SIGINT, _handle_sigint)
+
+    if master_process:
+        print("\nStarting training" +
+              (f" (resumed at step {start_step})" if start_step else "") + "...\n", flush=True)
+
+    for step in range(start_step, cfg.max_steps):
+        # ---------- periodic validation ----------
+        if cfg.val_loss_steps > 0 and (step % cfg.val_freq == 0 or step == cfg.max_steps - 1):
+            model.eval()
+            val_loader.reset()
+            acc = 0.0
+            with torch.no_grad(), amp:
+                for _ in range(cfg.val_loss_steps):
+                    x, y, _, _ = val_loader.next_batch()
+                    acc += model(x, y)[1].item() / cfg.val_loss_steps
+            model.train()
+            # val_loader is rank=0/world_size=1 on every rank (see its
+            # construction above), so every rank just computed the identical
+            # fixed windows redundantly — no cross-rank reduction needed, and
+            # val_loss no longer depends on how many ranks are running.
+            last_val = acc
+            if master_process:
+                print(f"  [val] step {step} | val_loss {last_val:.6f}", flush=True)
+
+        # ---------- training step ----------
+        t0 = time.perf_counter()
+        model.zero_grad(set_to_none=True)
+        # Fresh zeros rather than an in-place scale of last step's accumulator:
+        # multiplying by zero does NOT clear a NaN, so a single non-finite
+        # micro-batch would poison the accumulator for the rest of the run.
+        loss_accum = torch.zeros((), device=device)
+        micro_loss, micro_where = [], []
+
+        for micro in range(grad_accum):
+            x, y, shard, offset = train_loader.next_batch()
+            # DDP allreduces gradients on every .backward(); for grad_accum > 1
+            # only the LAST micro-batch of the step should pay for that — the
+            # others accumulate purely locally under no_sync().
+            sync_ctx = (nullcontext() if (not ddp or micro == grad_accum - 1)
+                       else model.no_sync())
+            with sync_ctx:
+                with amp:
+                    _, loss = model(x, y)
+                loss_accum += loss.detach()
+                # Per-micro-batch loss (--diag). The reported train loss is the mean
+                # over grad_accum micro-steps, so a single pathological micro-batch
+                # is diluted while its gradient dominates the accumulated norm
+                # almost undiluted — which is how a modest loss bump can sit next to
+                # a large norm spike. `micro_where` names the shard and token offset
+                # so the responsible bytes can be decoded offline.
+                #
+                # Costs a device sync per micro-step, so it is off by default.
+                if diag:
+                    micro_loss.append(loss.item())
+                    micro_where.append((shard, offset))
+                (loss / grad_accum).backward()
+
+        if ddp:
+            # Average each rank's accumulated (summed) loss BEFORE dividing by
+            # grad_accum, so loss_val is the mean over world_size*grad_accum
+            # micro-batches — the true loss behind this optimizer step, not
+            # just rank 0's slice. Also doubles as a sync point: a NaN on any
+            # rank poisons the average for every rank, so the check below
+            # raises everywhere instead of only where it originated (which
+            # would otherwise hang the survivors on the next collective).
+            dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG)
+        loss_val = loss_accum.item() / grad_accum
+        if math.isnan(loss_val) or math.isinf(loss_val):
+            if log_file:
+                log_file.close()
+            # Every rank reaches this identically (the all_reduce above is a
+            # sync point and NaN propagates through AVG), so every rank exits
+            # here together rather than the survivors hanging on the next
+            # collective.
+            if ddp:
+                dist.destroy_process_group()
+            raise SystemExit(f"ERROR: NaN/Inf loss at step {step}")
+
+        # Depth probe (--diag), read BEFORE the update, so these are the
+        # pre-update gradients. Compare the profile on a spiking step against a
+        # quiet one: a blow-up that begins mid-stack and grows toward the loss
+        # points at the forward, while a flat multiplicative lift across all
+        # depths points at the optimizer/update.
+        depth = []
+        if diag:
+            for L in probes:
+                a = raw_model.blocks[L].attn
+                gn = a.wQ.weight.grad
+                depth.append((L,
+                              gn.float().norm().item() if gn is not None else -1.0,
+                              a.q_gamma.abs().max().item(),
+                              a.wQ.weight.abs().max().float().item()))
+
+        grad_norm = clip_grads_(params, cfg.grad_clip).item()
+        lr = get_lr(step, cfg)                                   # AdamW group
+        muon_lr = get_lr(step, cfg, cfg.muon_lr, cfg.muon_min_lr)
+        opt.step(lr, muon_lr)
+
+        torch.cuda.synchronize()   # make the per-step timing reflect completed GPU work
+        dt = time.perf_counter() - t0
+        tok_per_sec = tokens_per_step / dt
+        elapsed_min = (time.perf_counter() - t_train_start) / 60.0
+        if step - start_step >= args.warmup_measure:
+            steady_dts.append(dt)
+
+        # ---------- per-step report (rank 0 only — every rank's numbers agree
+        # on loss/grad_norm post-allreduce, so printing from every rank would
+        # only duplicate lines, not add information) ----------
+        if master_process:
+            vs = "     -  " if last_val != last_val else f"{last_val:.6f}"
+            print(f"step {step:6d} | train {loss_val:.6f} | val {vs} "
+                  f"| tok/s {tok_per_sec:10.2f} | norm {grad_norm:.4f} "
+                  f"| dt {dt * 1e3:8.2f} ms | lr {lr:.4e} | mlr {muon_lr:.4e} "
+                  f"| mem {torch.cuda.max_memory_allocated() / gb:5.2f}/"
+                  f"{torch.cuda.max_memory_reserved() / gb:5.2f} GiB", flush=True)
+
+            if diag and micro_loss:
+                # Flag micro-batches that stand out against this step's own median.
+                med = sorted(micro_loss)[len(micro_loss) // 2]
+                marks = " ".join(("*" if m > med + 0.5 else "") + f"{m:.3f}"
+                                 for m in micro_loss)
+                print(f"  [micro] med {med:.4f} | {marks}")
+                for m, (ml, (shard, offset)) in enumerate(zip(micro_loss, micro_where)):
+                    if ml > med + 0.5:
+                        print(f"  [micro] OUTLIER m={m} loss {ml:.4f} (median {med:.4f}) "
+                              f"from shard {shard} token offset {offset}")
+            if diag and depth:
+                print("  [depth] ||g(wQ)|| by layer:"
+                      + "".join(f" L{L} {g:.3e}" for L, g, _, _ in depth)
+                      + " | max|q_gamma|:"
+                      + "".join(f" L{L} {gm:.4f}" for L, _, gm, _ in depth)
+                      + " | max|wQ|:"
+                      + "".join(f" L{L} {w:.4f}" for L, _, _, w in depth), flush=True)
+
+            if log_file:
+                log_file.write(f"{step},{loss_val:.6f},{last_val:.6f},{tok_per_sec:.6f},"
+                               f"{grad_norm:.6f},{dt * 1e3:.6f},{lr:.6f},{elapsed_min:.6f},"
+                               f"{muon_lr:.6f}\n")
+                log_file.flush()
+
+        if ddp:
+            # `interrupted` is set by this process's own SIGINT handler, so in
+            # principle one rank could see it before another (a signal isn't
+            # guaranteed simultaneous delivery, even though in practice a
+            # terminal Ctrl-C hits every rank in the same process group at
+            # once). Left asymmetric, it would desync the checkpoint barrier
+            # and the `break` below against ranks still issuing DDP's regular
+            # per-step collectives — a hang, not just a wrong answer. This
+            # runs unconditionally every step (never gated on a per-rank
+            # value), so it can never itself mismatch across ranks.
+            flag = torch.tensor(1.0 if interrupted else 0.0, device=device)
+            dist.all_reduce(flag, op=dist.ReduceOp.MAX)
+            interrupted = flag.item() > 0
+
+        # ---------- periodic checkpoint (also fires once on Ctrl-C) ----------
+        # Labeled step+1 (the next iteration to run) so load_latest resumes
+        # exactly where this save left off — no duplicated or skipped iteration.
+        # Captures weights + Adam m/v AFTER this step's update. Rank 0 only —
+        # every rank's raw_model/opt state is identical post-allreduce, so a
+        # second writer would only race the first on the same path. The
+        # barrier holds every rank at this point until the save (or its
+        # failure) is done, so a slow write can't let other ranks race ahead
+        # into the next step's collectives before the checkpoint is durable.
+        if cfg.checkpointing and ((step + 1) % cfg.ckpt_freq == 0
+                                  or step == cfg.max_steps - 1
+                                  or interrupted):
+            if master_process:
+                try:
+                    ckpt.save(step + 1, raw_model, opt, loss_val, last_val, cfg)
+                except Exception as e:                   # noqa: BLE001 — advisory
+                    print(f"[WARN] checkpoint save at step {step + 1} failed: {e}")
+            if ddp:
+                dist.barrier()
+
+        if interrupted:
+            if not cfg.checkpointing and master_process:
+                print(f"[ctrl-c] checkpointing is disabled (--no-checkpointing); "
+                      f"stopping at step {step} without saving.", flush=True)
+            break
+
+        # Reclaim the extra allocator high-water mark from the first couple of
+        # steps' one-time warmup allocations (autograd graph caches, etc.).
+        if step in (start_step, start_step + 1):
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+
+    torch.cuda.synchronize()
+    if log_file:
+        log_file.close()
+
+    # ---- steady-state summary ----
+    # dt/throughput are this rank's own timing; under DDP all ranks run in
+    # lockstep (grad sync + the checkpoint barrier), so rank 0's numbers are
+    # representative and `tokens_per_step` already counts every rank's tokens.
+    if steady_dts and master_process:
+        xs = sorted(steady_dts)
+        med = xs[len(xs) // 2]
+        mean = sum(xs) / len(xs)
+        alloc = torch.cuda.max_memory_allocated() / gb
+        reserved = torch.cuda.max_memory_reserved() / gb
+        # Measured now that Adam's m/v exist, rather than the startup estimate.
+        static = (param_bytes + grad_bytes + opt.state_bytes()) / gb
+        mode_str = f"DDP x{ddp_world_size}" if ddp else "single GPU"
+        print("\n" + "=" * 78)
+        print(f"  STEADY-STATE SUMMARY — {mode_str}  optimizer=muon+adamw  "
+              f"precision={args.precision}  "
+              f"rms+rope={'fused' if cfg.fused_rms_rope else 'eager'}")
+        print(f"  torch.optim.Muon  lr={cfg.muon_lr:g} "
+              f"adjust_lr_fn={cfg.muon_adjust_lr} "
+              f"momentum={cfg.muon_momentum:g} ns_steps={cfg.muon_ns_steps}")
+        print(f"  ({len(xs)} steps, first {args.warmup_measure} excluded)")
+        print("=" * 78)
+        print(f"  step time          : {med * 1e3:10.2f} ms  (mean {mean * 1e3:.2f})")
+        print(f"  throughput         : {tokens_per_step / med:10.1f} tok/s")
+        print(f"  tokens / step      : {tokens_per_step:10d}")
+        print(f"  static (p+g+opt)   : {static:10.2f} GiB  (per GPU)")
+        print(f"  activations (peak) : {max(alloc - static, 0.0):10.2f} GiB")
+        print(f"  peak mem allocated : {alloc:10.2f} GiB")
+        print(f"  peak mem reserved  : {reserved:10.2f} GiB")
+        print("=" * 78)
+
+    if ddp:
+        dist.destroy_process_group()
+    if master_process:
+        print("\n=== Training complete ===")
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+