---

--- Training AdamW on TinyStories ---
Device: cuda | Params: 298400 | Steps: 30000 | LR: 0.001
Step 100   | Loss: 2.3693
Step 200   | Loss: 2.2290
Step 300   | Loss: 2.2405
Step 400   | Loss: 2.2655
Step 500   | Loss: 2.1799
Step 600   | Loss: 2.1831
Step 700   | Loss: 2.1735
Step 800   | Loss: 2.1372
Step 900   | Loss: 2.1364
Step 1000  | Loss: 2.1368
Saved checkpoint: checkpoint_AdamW_step1000.pt
Step 1100  | Loss: 2.2150
Step 1200  | Loss: 2.1900
Step 1300  | Loss: 2.1310
Step 1400  | Loss: 2.1300
Step 1500  | Loss: 2.1200
Step 1600  | Loss: 2.1532
Step 1700  | Loss: 2.1184
Step 1800  | Loss: 2.0754
Step 1900  | Loss: 2.0871
Step 2000  | Loss: 2.1420
Saved checkpoint: checkpoint_AdamW_step2000.pt
Step 2100  | Loss: 2.1150
Step 2200  | Loss: 2.1718
Step 2300  | Loss: 2.1013
Step 2400  | Loss: 2.0593
Step 2500  | Loss: 2.1299
Step 2600  | Loss: 2.0749
Step 2700  | Loss: 2.1173
Step 2800  | Loss: 2.1176
Step 2900  | Loss: 2.1677
Step 3000  | Loss: 2.0864
Saved checkpoint: checkpoint_AdamW_step3000.pt
Step 3100  | Loss: 2.1030
Step 3200  | Loss: 2.0285
Step 3300  | Loss: 2.1043
Step 3400  | Loss: 2.0900
Step 3500  | Loss: 2.1592
Step 3600  | Loss: 2.1009
Step 3700  | Loss: 2.0683
Step 3800  | Loss: 2.0703
Step 3900  | Loss: 2.0656
Step 4000  | Loss: 2.0858
Saved checkpoint: checkpoint_AdamW_step4000.pt
Step 4100  | Loss: 2.0786
Step 4200  | Loss: 2.0698
Step 4300  | Loss: 2.0725
Step 4400  | Loss: 2.0349
Step 4500  | Loss: 2.0330
Step 4600  | Loss: 1.9249
Step 4700  | Loss: 1.6299
Step 4800  | Loss: 1.1039
Step 4900  | Loss: 0.6402
Step 5000  | Loss: 0.4044
Saved checkpoint: checkpoint_AdamW_step5000.pt
Step 5100  | Loss: 0.1376
Step 5200  | Loss: 0.0806
Step 5300  | Loss: 0.0761
Step 5400  | Loss: 0.0521
Step 5500  | Loss: 0.0516
Step 5600  | Loss: 0.0550
Step 5700  | Loss: 0.0413
Step 5800  | Loss: 0.0453
Step 5900  | Loss: 0.0485
Step 6000  | Loss: 0.0376
Saved checkpoint: checkpoint_AdamW_step6000.pt
Step 6100  | Loss: 0.0391
Step 6200  | Loss: 0.0376
Step 6300  | Loss: 0.0280
Step 6400  | Loss: 0.0278
Step 6500  | Loss: 0.0289
Step 6600  | Loss: 0.0200
Step 6700  | Loss: 0.0247
Step 6800  | Loss: 0.0272
Step 6900  | Loss: 0.0301
Step 7000  | Loss: 0.0258
Saved checkpoint: checkpoint_AdamW_step7000.pt
Step 7100  | Loss: 0.0219
Step 7200  | Loss: 0.0255
Step 7300  | Loss: 0.0251
Step 7400  | Loss: 0.0266
Step 7500  | Loss: 0.0267
Step 7600  | Loss: 0.0239
Step 7700  | Loss: 0.0249
Step 7800  | Loss: 0.0219
Step 7900  | Loss: 0.0229
Step 8000  | Loss: 0.0206
Saved checkpoint: checkpoint_AdamW_step8000.pt
Step 8100  | Loss: 0.0236
Step 8200  | Loss: 0.0387
Step 8300  | Loss: 0.0171
Step 8400  | Loss: 0.0253
Step 8500  | Loss: 0.0165
Step 8600  | Loss: 0.0183
Step 8700  | Loss: 0.0268
Step 8800  | Loss: 0.0236
Step 8900  | Loss: 0.0220
Step 9000  | Loss: 0.0213
Saved checkpoint: checkpoint_AdamW_step9000.pt
Step 9100  | Loss: 0.0178
Step 9200  | Loss: 0.0164
Step 9300  | Loss: 0.0229
Step 9400  | Loss: 0.0232
Step 9500  | Loss: 0.0185
Step 9600  | Loss: 0.0183
Step 9700  | Loss: 0.0190
Step 9800  | Loss: 0.0204
Step 9900  | Loss: 0.0235
Step 10000 | Loss: 0.0184
Saved checkpoint: checkpoint_AdamW_step10000.pt
Step 10100 | Loss: 0.0214
Step 10200 | Loss: 0.0158
Step 10300 | Loss: 0.0201
Step 10400 | Loss: 0.0199
Step 10500 | Loss: 0.0213
Step 10600 | Loss: 0.0213
Step 10700 | Loss: 0.0190
Step 10800 | Loss: 0.0202
Step 10900 | Loss: 0.0218
Step 11000 | Loss: 0.0186
Saved checkpoint: checkpoint_AdamW_step11000.pt
Step 11100 | Loss: 0.0166
Step 11200 | Loss: 0.0201
Step 11300 | Loss: 0.0242
Step 11400 | Loss: 0.0247
Step 11500 | Loss: 0.0177
Step 11600 | Loss: 0.0179
Step 11700 | Loss: 0.0160
Step 11800 | Loss: 0.0197
Step 11900 | Loss: 0.0143
Step 12000 | Loss: 0.0238
Saved checkpoint: checkpoint_AdamW_step12000.pt
Step 12100 | Loss: 0.0187
Step 12200 | Loss: 0.0151
Step 12300 | Loss: 0.0171
Step 12400 | Loss: 0.0246
Step 12500 | Loss: 0.0162
Step 12600 | Loss: 0.0201
Step 12700 | Loss: 0.0184
Step 12800 | Loss: 0.0180
Step 12900 | Loss: 0.0160
Step 13000 | Loss: 0.0179
Saved checkpoint: checkpoint_AdamW_step13000.pt
Step 13100 | Loss: 0.0178
Step 13200 | Loss: 0.0188
Step 13300 | Loss: 0.0142
Step 13400 | Loss: 0.0169
Step 13500 | Loss: 0.0217
Step 13600 | Loss: 0.0276
Step 13700 | Loss: 0.0144
Step 13800 | Loss: 0.0159
Step 13900 | Loss: 0.0197
Step 14000 | Loss: 0.0167
Saved checkpoint: checkpoint_AdamW_step14000.pt
Step 14100 | Loss: 0.0209
Step 14200 | Loss: 0.0159
Step 14300 | Loss: 0.0198
Step 14400 | Loss: 0.0136
Step 14500 | Loss: 0.0177
Step 14600 | Loss: 0.0185
Step 14700 | Loss: 0.0133
Step 14800 | Loss: 0.0164
Step 14900 | Loss: 0.0160
Step 15000 | Loss: 0.0139
Saved checkpoint: checkpoint_AdamW_step15000.pt
Step 15100 | Loss: 0.0228
Step 15200 | Loss: 0.0141
Step 15300 | Loss: 0.0182
Step 15400 | Loss: 0.0147
Step 15500 | Loss: 0.0200
Step 15600 | Loss: 0.0150
Step 15700 | Loss: 0.0140
Step 15800 | Loss: 0.0149
Step 15900 | Loss: 0.0154
Step 16000 | Loss: 0.0129
Saved checkpoint: checkpoint_AdamW_step16000.pt
Step 16100 | Loss: 0.0173
Step 16200 | Loss: 0.0139
Step 16300 | Loss: 0.0147
Step 16400 | Loss: 0.0126
Step 16500 | Loss: 0.0138
Step 16600 | Loss: 0.0174
Step 16700 | Loss: 0.0166
Step 16800 | Loss: 0.0126
Step 16900 | Loss: 0.0120
Step 17000 | Loss: 0.0141
Saved checkpoint: checkpoint_AdamW_step17000.pt
Step 17100 | Loss: 0.0153
Step 17200 | Loss: 0.0138
Step 17300 | Loss: 0.0176
Step 17400 | Loss: 0.0137
Step 17500 | Loss: 0.0170
Step 17600 | Loss: 0.0153
Step 17700 | Loss: 0.0152
Step 17800 | Loss: 0.0158
Step 17900 | Loss: 0.0143
Step 18000 | Loss: 0.0178
Saved checkpoint: checkpoint_AdamW_step18000.pt
Step 18100 | Loss: 0.0161
Step 18200 | Loss: 0.0124
Step 18300 | Loss: 0.0145
Step 18400 | Loss: 0.0139
Step 18500 | Loss: 0.0164
Step 18600 | Loss: 0.0143
Step 18700 | Loss: 0.0118
Step 18800 | Loss: 0.0137
Step 18900 | Loss: 0.0171
Step 19000 | Loss: 0.0141
Saved checkpoint: checkpoint_AdamW_step19000.pt
Step 19100 | Loss: 0.0169
Step 19200 | Loss: 0.0131
Step 19300 | Loss: 0.0163
Step 19400 | Loss: 0.0151
Step 19500 | Loss: 0.0157
Step 19600 | Loss: 0.0131
Step 19700 | Loss: 0.0117
Step 19800 | Loss: 0.0146
Step 19900 | Loss: 0.0100
Step 20000 | Loss: 0.0135
Saved checkpoint: checkpoint_AdamW_step20000.pt
Step 20100 | Loss: 0.0164
Step 20200 | Loss: 0.0129
Step 20300 | Loss: 0.0172
Step 20400 | Loss: 0.0119
Step 20500 | Loss: 0.0100
Step 20600 | Loss: 0.0150
Step 20700 | Loss: 0.0127
Step 20800 | Loss: 0.0130
Step 20900 | Loss: 0.0142
Step 21000 | Loss: 0.0156
Saved checkpoint: checkpoint_AdamW_step21000.pt
Step 21100 | Loss: 0.0123
Step 21200 | Loss: 0.0158
Step 21300 | Loss: 0.0143
Step 21400 | Loss: 0.0132
Step 21500 | Loss: 0.0104
Step 21600 | Loss: 0.0180
Step 21700 | Loss: 0.0109
Step 21800 | Loss: 0.0184
Step 21900 | Loss: 0.0139
Step 22000 | Loss: 0.0133
Saved checkpoint: checkpoint_AdamW_step22000.pt
Step 22100 | Loss: 0.0131
Step 22200 | Loss: 0.0112
Step 22300 | Loss: 0.0139
Step 22400 | Loss: 0.0158
Step 22500 | Loss: 0.0118
Step 22600 | Loss: 0.0117
Step 22700 | Loss: 0.0145
Step 22800 | Loss: 0.0161
Step 22900 | Loss: 0.0135
Step 23000 | Loss: 0.0146
Saved checkpoint: checkpoint_AdamW_step23000.pt
Step 23100 | Loss: 0.0147
Step 23200 | Loss: 0.0139
Step 23300 | Loss: 0.0178
Step 23400 | Loss: 0.0121
Step 23500 | Loss: 0.0156
Step 23600 | Loss: 0.0123
Step 23700 | Loss: 0.0139
Step 23800 | Loss: 0.0135
Step 23900 | Loss: 0.0125
Step 24000 | Loss: 0.0146
Saved checkpoint: checkpoint_AdamW_step24000.pt
Step 24100 | Loss: 0.0089
Step 24200 | Loss: 0.0178
Step 24300 | Loss: 0.0152
Step 24400 | Loss: 0.0140
Step 24500 | Loss: 0.0158
Step 24600 | Loss: 0.0134
Step 24700 | Loss: 0.0134
Step 24800 | Loss: 0.0134
Step 24900 | Loss: 0.0114
Step 25000 | Loss: 0.0109
Saved checkpoint: checkpoint_AdamW_step25000.pt
Step 25100 | Loss: 0.0143
Step 25200 | Loss: 0.0135
Step 25300 | Loss: 0.0165
Step 25400 | Loss: 0.0143
Step 25500 | Loss: 0.0143
Step 25600 | Loss: 0.0114
Step 25700 | Loss: 0.0109
Step 25800 | Loss: 0.0112
Step 25900 | Loss: 0.0135
Step 26000 | Loss: 0.0143
Saved checkpoint: checkpoint_AdamW_step26000.pt
Step 26100 | Loss: 0.0101
Step 26200 | Loss: 0.0137
Step 26300 | Loss: 0.0150
Step 26400 | Loss: 0.0121
Step 26500 | Loss: 0.0098
Step 26600 | Loss: 0.0131
Step 26700 | Loss: 0.0099
Step 26800 | Loss: 0.0166
Step 26900 | Loss: 0.0115
Step 27000 | Loss: 0.0155
Saved checkpoint: checkpoint_AdamW_step27000.pt
Step 27100 | Loss: 0.0154
Step 27200 | Loss: 0.0147
Step 27300 | Loss: 0.0124
Step 27400 | Loss: 0.0139
Step 27500 | Loss: 0.0136
Step 27600 | Loss: 0.0147
Step 27700 | Loss: 0.0129
Step 27800 | Loss: 0.0117
Step 27900 | Loss: 0.0126
Step 28000 | Loss: 0.0131
Saved checkpoint: checkpoint_AdamW_step28000.pt
Step 28100 | Loss: 0.0124
Step 28200 | Loss: 0.0143
Step 28300 | Loss: 0.0133
Step 28400 | Loss: 0.0132
Step 28500 | Loss: 0.0121
Step 28600 | Loss: 0.0113
Step 28700 | Loss: 0.0086
Step 28800 | Loss: 0.0128
Step 28900 | Loss: 0.0160
Step 29000 | Loss: 0.0124
Saved checkpoint: checkpoint_AdamW_step29000.pt
Step 29100 | Loss: 0.0109
Step 29200 | Loss: 0.0138
Step 29300 | Loss: 0.0129
Step 29400 | Loss: 0.0124
Step 29500 | Loss: 0.0113
Step 29600 | Loss: 0.0153
Step 29700 | Loss: 0.0103
Step 29800 | Loss: 0.0105
Step 29900 | Loss: 0.0122
Step 30000 | Loss: 0.0114
Saved checkpoint: checkpoint_AdamW_step30000.pt

--- Training Muon on TinyStories ---
Device: cuda | Params: 298400 | Steps: 30000 | LR: 0.001
torch.optim.Muon NOT found. Falling back to custom implementation.
Step 100   | Loss: 2.5669
Step 200   | Loss: 2.3135
Step 300   | Loss: 2.3385
Step 400   | Loss: 2.3726
Step 500   | Loss: 2.2891
Step 600   | Loss: 2.2959
Step 700   | Loss: 2.3183
Step 800   | Loss: 2.2587
Step 900   | Loss: 2.2709
Step 1000  | Loss: 2.3068
Saved checkpoint: checkpoint_Muon_step1000.pt
Step 1100  | Loss: 2.3524
Step 1200  | Loss: 2.3086
Step 1300  | Loss: 2.2748
Step 1400  | Loss: 2.2906
Step 1500  | Loss: 2.2590
Step 1600  | Loss: 2.2915
Step 1700  | Loss: 2.2636
Step 1800  | Loss: 2.2068
Step 1900  | Loss: 2.2202
Step 2000  | Loss: 2.2575
Saved checkpoint: checkpoint_Muon_step2000.pt
Step 2100  | Loss: 2.2340
Step 2200  | Loss: 2.2854
Step 2300  | Loss: 2.2188
Step 2400  | Loss: 2.2045
Step 2500  | Loss: 2.2734
Step 2600  | Loss: 2.2041
Step 2700  | Loss: 2.2606
Step 2800  | Loss: 2.2404
Step 2900  | Loss: 2.2879
Step 3000  | Loss: 2.2052
Saved checkpoint: checkpoint_Muon_step3000.pt
Step 3100  | Loss: 2.2112
Step 3200  | Loss: 2.1528
Step 3300  | Loss: 2.2177
Step 3400  | Loss: 2.1883
Step 3500  | Loss: 2.2606
Step 3600  | Loss: 2.2101
Step 3700  | Loss: 2.2019
Step 3800  | Loss: 2.1808
Step 3900  | Loss: 2.1735
Step 4000  | Loss: 2.1871
Saved checkpoint: checkpoint_Muon_step4000.pt
Step 4100  | Loss: 2.1788
Step 4200  | Loss: 2.1826
Step 4300  | Loss: 2.1780
Step 4400  | Loss: 2.1613
Step 4500  | Loss: 2.2029
Step 4600  | Loss: 2.2319
Step 4700  | Loss: 2.2051
Step 4800  | Loss: 2.1433
Step 4900  | Loss: 2.1839
Step 5000  | Loss: 2.1566
Saved checkpoint: checkpoint_Muon_step5000.pt
Step 5100  | Loss: 2.1327
Step 5200  | Loss: 2.1451
Step 5300  | Loss: 2.2499
Step 5400  | Loss: 2.2026
Step 5500  | Loss: 2.1571
Step 5600  | Loss: 2.1604
Step 5700  | Loss: 2.1719
Step 5800  | Loss: 2.1695
Step 5900  | Loss: 2.1955
Step 6000  | Loss: 2.1353
Saved checkpoint: checkpoint_Muon_step6000.pt
Step 6100  | Loss: 2.1519
Step 6200  | Loss: 2.2194
Step 6300  | Loss: 2.1198
Step 6400  | Loss: 2.1362
Step 6500  | Loss: 2.1516
Step 6600  | Loss: 2.1233
Step 6700  | Loss: 2.1004
Step 6800  | Loss: 2.1213
Step 6900  | Loss: 2.2155
Step 7000  | Loss: 2.1779
Saved checkpoint: checkpoint_Muon_step7000.pt
Step 7100  | Loss: 2.1487
Step 7200  | Loss: 2.1697
Step 7300  | Loss: 2.1147
Step 7400  | Loss: 2.1679
Step 7500  | Loss: 2.1445
Step 7600  | Loss: 2.1434
Step 7700  | Loss: 2.1220
Step 7800  | Loss: 2.2224
Step 7900  | Loss: 2.1088
Step 8000  | Loss: 2.1331
Saved checkpoint: checkpoint_Muon_step8000.pt
Step 8100  | Loss: 2.1400
Step 8200  | Loss: 2.2085
Step 8300  | Loss: 2.1500
Step 8400  | Loss: 2.2345
Step 8500  | Loss: 2.1616
Step 8600  | Loss: 2.1778
Step 8700  | Loss: 2.1072
Step 8800  | Loss: 2.1704
Step 8900  | Loss: 2.1573
Step 9000  | Loss: 2.1650
Saved checkpoint: checkpoint_Muon_step9000.pt
Step 9100  | Loss: 2.1213
Step 9200  | Loss: 2.0966
Step 9300  | Loss: 2.1523
Step 9400  | Loss: 2.1244
Step 9500  | Loss: 2.1367
Step 9600  | Loss: 2.1371
Step 9700  | Loss: 2.1932
Step 9800  | Loss: 2.0963
Step 9900  | Loss: 2.1017
Step 10000 | Loss: 2.1091
Saved checkpoint: checkpoint_Muon_step10000.pt
Step 10100 | Loss: 2.0915
Step 10200 | Loss: 2.0933
Step 10300 | Loss: 2.1365
Step 10400 | Loss: 2.1352
Step 10500 | Loss: 2.1360
Step 10600 | Loss: 2.1059
Step 10700 | Loss: 2.1100
Step 10800 | Loss: 2.0640
Step 10900 | Loss: 2.1101
Step 11000 | Loss: 2.0433
Saved checkpoint: checkpoint_Muon_step11000.pt
Step 11100 | Loss: 2.0581
Step 11200 | Loss: 2.0297
Step 11300 | Loss: 2.0179
Step 11400 | Loss: 1.8833
Step 11500 | Loss: 1.7789
Step 11600 | Loss: 1.5766
Step 11700 | Loss: 1.3656
Step 11800 | Loss: 0.7278
Step 11900 | Loss: 0.3071
Step 12000 | Loss: 0.1900
Saved checkpoint: checkpoint_Muon_step12000.pt
Step 12100 | Loss: 0.1328
Step 12200 | Loss: 0.1085
Step 12300 | Loss: 0.0942
Step 12400 | Loss: 0.1042
Step 12500 | Loss: 0.0672
Step 12600 | Loss: 0.0603
Step 12700 | Loss: 0.0621
Step 12800 | Loss: 0.0505
Step 12900 | Loss: 0.0426
Step 13000 | Loss: 0.0416
Saved checkpoint: checkpoint_Muon_step13000.pt
Step 13100 | Loss: 0.0372
Step 13200 | Loss: 0.0483
Step 13300 | Loss: 0.0376
Step 13400 | Loss: 0.0298
Step 13500 | Loss: 0.0416
Step 13600 | Loss: 0.0435
Step 13700 | Loss: 0.0324
Step 13800 | Loss: 0.0313
Step 13900 | Loss: 0.0271
Step 14000 | Loss: 0.0339
Saved checkpoint: checkpoint_Muon_step14000.pt
Step 14100 | Loss: 0.0312
Step 14200 | Loss: 0.0283
Step 14300 | Loss: 0.0304
Step 14400 | Loss: 0.0273
Step 14500 | Loss: 0.0319
Step 14600 | Loss: 0.0282
Step 14700 | Loss: 0.0261
Step 14800 | Loss: 0.0292
Step 14900 | Loss: 0.0292
Step 15000 | Loss: 0.0225
Saved checkpoint: checkpoint_Muon_step15000.pt
Step 15100 | Loss: 0.0322
Step 15200 | Loss: 0.0267
Step 15300 | Loss: 0.0287
Step 15400 | Loss: 0.0228
Step 15500 | Loss: 0.0262
Step 15600 | Loss: 0.0247
Step 15700 | Loss: 0.0250
Step 15800 | Loss: 0.0258
Step 15900 | Loss: 0.0231
Step 16000 | Loss: 0.0212
Saved checkpoint: checkpoint_Muon_step16000.pt
Step 16100 | Loss: 0.0253
Step 16200 | Loss: 0.0224
Step 16300 | Loss: 0.0213
Step 16400 | Loss: 0.0247
Step 16500 | Loss: 0.0253
Step 16600 | Loss: 0.0252
Step 16700 | Loss: 0.0231
Step 16800 | Loss: 0.0233
Step 16900 | Loss: 0.0238
Step 17000 | Loss: 0.0200
Saved checkpoint: checkpoint_Muon_step17000.pt
Step 17100 | Loss: 0.0244
Step 17200 | Loss: 0.0203
Step 17300 | Loss: 0.0223
Step 17400 | Loss: 0.0188
Step 17500 | Loss: 0.0207
Step 17600 | Loss: 0.0191
Step 17700 | Loss: 0.0182
Step 17800 | Loss: 0.0249
Step 17900 | Loss: 0.0164
Step 18000 | Loss: 0.0241
Saved checkpoint: checkpoint_Muon_step18000.pt
Step 18100 | Loss: 0.0236
Step 18200 | Loss: 0.0209
Step 18300 | Loss: 0.0232
Step 18400 | Loss: 0.0212
Step 18500 | Loss: 0.0213
Step 18600 | Loss: 0.0202
Step 18700 | Loss: 0.0213
Step 18800 | Loss: 0.0248
Step 18900 | Loss: 0.0197
Step 19000 | Loss: 0.0176
Saved checkpoint: checkpoint_Muon_step19000.pt
Step 19100 | Loss: 0.0239
Step 19200 | Loss: 0.0224
Step 19300 | Loss: 0.0212
Step 19400 | Loss: 0.0230
Step 19500 | Loss: 0.0196
Step 19600 | Loss: 0.0188
Step 19700 | Loss: 0.0187
Step 19800 | Loss: 0.0220
Step 19900 | Loss: 0.0178
Step 20000 | Loss: 0.0223
Saved checkpoint: checkpoint_Muon_step20000.pt
Step 20100 | Loss: 0.0192
Step 20200 | Loss: 0.0182
Step 20300 | Loss: 0.0220
Step 20400 | Loss: 0.0183
Step 20500 | Loss: 0.0200
Step 20600 | Loss: 0.0229
Step 20700 | Loss: 0.0190
Step 20800 | Loss: 0.0205
Step 20900 | Loss: 0.0213
Step 21000 | Loss: 0.0215
Saved checkpoint: checkpoint_Muon_step21000.pt
Step 21100 | Loss: 0.0187
Step 21200 | Loss: 0.0216
Step 21300 | Loss: 0.0196
Step 21400 | Loss: 0.0196
Step 21500 | Loss: 0.0182
Step 21600 | Loss: 0.0203
Step 21700 | Loss: 0.0173
Step 21800 | Loss: 0.0206
Step 21900 | Loss: 0.0193
Step 22000 | Loss: 0.0191
Saved checkpoint: checkpoint_Muon_step22000.pt
Step 22100 | Loss: 0.0184
Step 22200 | Loss: 0.0189
Step 22300 | Loss: 0.0202
Step 22400 | Loss: 0.0181
Step 22500 | Loss: 0.0203
Step 22600 | Loss: 0.0188
Step 22700 | Loss: 0.0202
Step 22800 | Loss: 0.0214
Step 22900 | Loss: 0.0199
Step 23000 | Loss: 0.0208
Saved checkpoint: checkpoint_Muon_step23000.pt
Step 23100 | Loss: 0.0215
Step 23200 | Loss: 0.0188
Step 23300 | Loss: 0.0221
Step 23400 | Loss: 0.0193
Step 23500 | Loss: 0.0216
Step 23600 | Loss: 0.0186
Step 23700 | Loss: 0.0183
Step 23800 | Loss: 0.0196
Step 23900 | Loss: 0.0191
Step 24000 | Loss: 0.0179
Saved checkpoint: checkpoint_Muon_step24000.pt
Step 24100 | Loss: 0.0133
Step 24200 | Loss: 0.0231
Step 24300 | Loss: 0.0185
Step 24400 | Loss: 0.0180
Step 24500 | Loss: 0.0222
Step 24600 | Loss: 0.0187
Step 24700 | Loss: 0.0169
Step 24800 | Loss: 0.0181
Step 24900 | Loss: 0.0185
Step 25000 | Loss: 0.0176
Saved checkpoint: checkpoint_Muon_step25000.pt
Step 25100 | Loss: 0.0173
Step 25200 | Loss: 0.0225
Step 25300 | Loss: 0.0213
Step 25400 | Loss: 0.0187
Step 25500 | Loss: 0.0191
Step 25600 | Loss: 0.0172
Step 25700 | Loss: 0.0190
Step 25800 | Loss: 0.0186
Step 25900 | Loss: 0.0200
Step 26000 | Loss: 0.0192
Saved checkpoint: checkpoint_Muon_step26000.pt
Step 26100 | Loss: 0.0173
Step 26200 | Loss: 0.0191
Step 26300 | Loss: 0.0199
Step 26400 | Loss: 0.0200
Step 26500 | Loss: 0.0161
Step 26600 | Loss: 0.0179
Step 26700 | Loss: 0.0186
Step 26800 | Loss: 0.0209
Step 26900 | Loss: 0.0198
Step 27000 | Loss: 0.0194
Saved checkpoint: checkpoint_Muon_step27000.pt
Step 27100 | Loss: 0.0174
Step 27200 | Loss: 0.0192
Step 27300 | Loss: 0.0167
Step 27400 | Loss: 0.0172
Step 27500 | Loss: 0.0187
Step 27600 | Loss: 0.0193
Step 27700 | Loss: 0.0215
Step 27800 | Loss: 0.0201
Step 27900 | Loss: 0.0212
Step 28000 | Loss: 0.0189
Saved checkpoint: checkpoint_Muon_step28000.pt
Step 28100 | Loss: 0.0178
Step 28200 | Loss: 0.0197
Step 28300 | Loss: 0.0197
Step 28400 | Loss: 0.0210
Step 28500 | Loss: 0.0186
Step 28600 | Loss: 0.0167
Step 28700 | Loss: 0.0140
Step 28800 | Loss: 0.0196
Step 28900 | Loss: 0.0191
Step 29000 | Loss: 0.0173
Saved checkpoint: checkpoint_Muon_step29000.pt
Step 29100 | Loss: 0.0171
Step 29200 | Loss: 0.0183
Step 29300 | Loss: 0.0186
Step 29400 | Loss: 0.0186
Step 29500 | Loss: 0.0163
Step 29600 | Loss: 0.0212
Step 29700 | Loss: 0.0169
Step 29800 | Loss: 0.0168
Step 29900 | Loss: 0.0156
Step 30000 | Loss: 0.0158
Saved checkpoint: checkpoint_Muon_step30000.pt

--- Training EML-V4 on TinyStories ---
Device: cuda | Params: 298400 | Steps: 30000 | LR: 0.001

--- Training EML-V4 on TinyStories ---
Device: cuda | Params: 298400 | Steps: 30000 | LR: 0.001
Step 10    | Loss: 39.7029
Step 20    | Loss: 19.2250
Step 30    | Loss: 7.3312
Step 40    | Loss: 4.0654
Step 50    | Loss: 3.3341
Step 60    | Loss: 2.8625
Step 70    | Loss: 2.8004
Step 80    | Loss: 2.6663
Step 90    | Loss: 2.6336
Step 100   | Loss: 2.6204
Step 110   | Loss: 2.5595
Step 120   | Loss: 2.5080
Step 130   | Loss: 2.4726
Step 140   | Loss: 2.5371
Step 150   | Loss: 2.4911
Step 160   | Loss: 2.4022
Step 170   | Loss: 2.4066
Step 180   | Loss: 2.4641
Step 190   | Loss: 2.4955
Step 200   | Loss: 2.3934
Step 210   | Loss: 2.4958
Step 220   | Loss: 2.4474
Step 230   | Loss: 2.4777
Step 240   | Loss: 2.4746
Step 250   | Loss: 2.3977
Step 260   | Loss: 2.4253
Step 270   | Loss: 2.5334
Step 280   | Loss: 2.4408
Step 290   | Loss: 2.4066
Step 300   | Loss: 2.3954
Step 310   | Loss: 2.3758
Step 320   | Loss: 2.4115
Step 330   | Loss: 2.3960
Step 340   | Loss: 2.4690
Step 350   | Loss: 2.3861
Step 360   | Loss: 2.3831
Step 370   | Loss: 2.3728
Step 380   | Loss: 2.3915
Step 390   | Loss: 2.3968
Step 400   | Loss: 2.4132
Step 410   | Loss: 2.4075
Step 420   | Loss: 2.3164
Step 430   | Loss: 2.3473
Step 440   | Loss: 2.3236
Step 450   | Loss: 2.3763
Step 460   | Loss: 2.3856
Step 470   | Loss: 2.3621
Step 480   | Loss: 2.3448
Step 490   | Loss: 2.3395
Step 500   | Loss: 2.3033
Step 510   | Loss: 2.3917
Step 520   | Loss: 2.3388
Step 530   | Loss: 2.3533
Step 540   | Loss: 2.3662
Step 550   | Loss: 2.3528
Step 560   | Loss: 2.4448
Step 570   | Loss: 2.3439
Step 580   | Loss: 2.3591
Step 590   | Loss: 2.3489
Step 600   | Loss: 2.3099
Step 610   | Loss: 2.3178
Step 620   | Loss: 2.3324
Step 630   | Loss: 2.3659
Step 640   | Loss: 2.3476
Step 650   | Loss: 2.3493
Step 660   | Loss: 2.4009
Step 670   | Loss: 2.2174
Step 680   | Loss: 2.2495
Step 690   | Loss: 2.3134
Step 700   | Loss: 2.3365
Step 710   | Loss: 2.2857
Step 720   | Loss: 2.2664
Step 730   | Loss: 2.2979
Step 740   | Loss: 2.3011
Step 750   | Loss: 2.3881
Step 760   | Loss: 2.2481
Step 770   | Loss: 2.2965
Step 780   | Loss: 2.3013
Step 790   | Loss: 2.2624
Step 800   | Loss: 2.2777
Step 810   | Loss: 2.3072
Step 820   | Loss: 2.3017
Step 830   | Loss: 2.2739
Step 840   | Loss: 2.2901
Step 850   | Loss: 2.2856
Step 860   | Loss: 2.2701
Step 870   | Loss: 2.3474
Step 880   | Loss: 2.3203
Step 890   | Loss: 2.3222
Step 900   | Loss: 2.2792
Step 910   | Loss: 2.3127
Step 920   | Loss: 2.3095
Step 930   | Loss: 2.2886
Step 940   | Loss: 2.2641
Step 950   | Loss: 2.2314
Step 960   | Loss: 2.2815
Step 970   | Loss: 2.3030
Step 980   | Loss: 2.3379
Step 990   | Loss: 2.2643
Step 1000  | Loss: 2.3149
Saved checkpoint: checkpoint_EML-V4_step1000.pt
Step 1010  | Loss: 2.2569
Step 1020  | Loss: 2.3005
Step 1030  | Loss: 2.2497
Step 1040  | Loss: 2.3095
Step 1050  | Loss: 2.3022
Step 1060  | Loss: 2.2500
Step 1070  | Loss: 2.3051
Step 1080  | Loss: 2.3089
Step 1090  | Loss: 2.3036
Step 1100  | Loss: 2.3614
Step 1110  | Loss: 2.2697
Step 1120  | Loss: 2.2485
Step 1130  | Loss: 2.3435
Step 1140  | Loss: 2.3406
Step 1150  | Loss: 8.8011
Step 1160  | Loss: 3.9739
Step 1170  | Loss: 3.4152
Step 1180  | Loss: 2.9224
Step 1190  | Loss: 2.6090
Step 1200  | Loss: 2.5537
Step 1210  | Loss: 2.4024
Step 1220  | Loss: 2.4519
Step 1230  | Loss: 2.4164
Step 1240  | Loss: 2.4281
Step 1250  | Loss: 2.4361
Step 1260  | Loss: 2.3295
Step 1270  | Loss: 2.3461
Step 1280  | Loss: 2.4351
Step 1290  | Loss: 2.3404
Step 1300  | Loss: 2.3770
Step 1310  | Loss: 2.3845
Step 1320  | Loss: 2.3540
Step 1330  | Loss: 2.3570
Step 1340  | Loss: 2.3507
Step 1350  | Loss: 2.4131
Step 1360  | Loss: 2.4026
Step 1370  | Loss: 2.3664
Step 1380  | Loss: 2.4509
Step 1390  | Loss: 2.3700
Step 1400  | Loss: 2.3678
Step 1410  | Loss: 2.3263
Step 1420  | Loss: 2.3699