feat: block-diagonal preconditioning, scao_step tracking, distributed sync

- preconditioner.py: implement use_block_diagonal for max(m,n)>max_precond_dim
  each block gets its own Kronecker preconditioner (matches paper Algorithm 1)
  handles update_curvature, precondition, natural_grad_norm, state_dict, memory_bytes
- optimizer.py: track scao_step (t_s) in Phase 2; shared-moment bias correction
  uses global step (correct for non-reset design); blend ramp = min(1, t_s/50)
- distributed.py: sync_preconditioners handles use_block_diagonal sub-blocks
- tests: +4 block-diagonal tests, +3 distributed tests (7 new, all pass, 67 total)
- README.md: accurate Innovation 2 description, corrected algorithm pseudocode

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: SCAO Authors

.gitignore

@@ -81,3 +81,4 @@ paper/
 secrets.yml
 *.pem
 *.key
+SCAO_Technical_Paper.docx

README.md

@@ -61,7 +61,7 @@ k* = argmin k such that  Σᵢ₌₁ᵏ λᵢ / Σⱼ λⱼ ≥ 1 − ε
 This reduces memory from `O(m² + n²)` to `O((m+n)·k)`. At GPT-2 scale (`d=768`), typical `k ≤ 32–64`, giving a **16–32× reduction** over full-rank Kronecker factors.
 
 ### Innovation 2 — Sparse Block-Diagonal FIM
-For layers where `max(m, n) > max_precond_dim`, SCAO falls back to a **diagonal curvature approximation** rather than storing any matrix at all. This prevents memory blow-up at large scales while preserving per-element adaptivity.
+For layers where `max(m, n) > max_precond_dim`, SCAO applies **sparse block-diagonal preconditioning**: the gradient matrix is partitioned into contiguous blocks of size ≤ `max_precond_dim` along the larger dimension, and an independent low-rank Kronecker preconditioner is applied per block. This bounds eigendecomp cost at `O(max_precond_dim³)` while preserving full curvature information across all blocks — unlike a diagonal fallback, which discards all inter-parameter correlation.
 
 ### Innovation 3 — Phase-Transition Stability
 The transition from Adam (Phase 1) to SCAO preconditioning (Phase 2) is the most dangerous moment in training. Three guards prevent instability:
@@ -93,18 +93,22 @@ Phase 2 — SCAO preconditioning (steps T_w + 1 onwards):
     Store: (U_L[:, :k], S_L[:k], U_R[:, :k], S_R[:k])
   
   Every step:
-    Project gradient: G_mid = U_L^T · G · U_R                    (k×k)
-    Scale:            G_scaled = diag(S_L^{-1/4}) · G_mid · diag(S_R^{-1/4})
-    Reconstruct:      g_eff = blend · (U_L · G_scaled · U_R^T) + (1-blend) · g
-                      where blend = min(1, (t - T_w) / 50)
+    Preconditioned gradient (identity + low-rank correction):
+      G_proj   = U_L^T · G · U_R                          (k×k)
+      G_scaled = diag(S_L^{-1/4}) · G_proj · diag(S_R^{-1/4})
+      G_precond = G + U_L · (G_scaled - G_proj) · U_R^T   (m×n)
     
-    Apply Adam update on g_eff:
+    50-step linear blend ramp (t_s = Phase-2 step count):
+      blend = min(1.0, t_s / 50)
+      g_eff = blend · G_precond + (1 - blend) · g
+    
+    Apply Adam update on g_eff (shared moments, warm-started from Phase 1):
       m_t ← β₁ · m_{t-1} + (1-β₁) · g_eff
-      v_t ← β₂ · v_t-1  + (1-β₂) · g_eff²
+      v_t ← β₂ · v_{t-1} + (1-β₂) · g_eff²
       θ_t ← θ_{t-1} - α · (m_t / (1-β₁^t)) / (√(v_t/(1-β₂^t)) + ε)
 ```
 
-**Key insight:** SCAO applies Adam *on top of* the preconditioned gradient. Both the first moment `m_t` and second moment `v_t` track `g_eff`, keeping numerator and denominator on the same scale throughout training.
+**Key insight:** SCAO uses **shared momentum tensors** that warm-start from Phase 1 into Phase 2. The blend ramp ensures a smooth transition — at `t_s=1`, `g_eff ≈ 0.98 · g_raw`, so the preconditioner's influence grows gradually without disrupting the accumulated momentum. Both moments track `g_eff`, keeping the numerator and denominator on the same scale throughout training.
 
 ---
 

results_benchmark.csv

@@ -0,0 +1,101 @@
+step,adamw,scao,diag_shampoo
+1,5.5728559494018555,5.5728559494018555,5.5728559494018555
+2,5.578701019287109,5.578701019287109,5.5962300300598145
+3,5.586780548095703,5.586779594421387,5.629923343658447
+4,5.573197364807129,5.573197364807129,5.654507160186768
+5,5.55900239944458,5.55900239944458,5.687943935394287
+6,5.579440593719482,5.579440593719482,5.665504455566406
+7,5.565978050231934,5.565978050231934,5.662998199462891
+8,5.568857192993164,5.568857192993164,5.635212421417236
+9,5.5764312744140625,5.576431751251221,5.631796836853027
+10,5.573879718780518,5.573879718780518,5.6374006271362305
+11,5.5734405517578125,5.5734405517578125,5.630044460296631
+12,5.573633193969727,5.573633193969727,5.607842922210693
+13,5.564875602722168,5.564875602722168,5.60982084274292
+14,5.560469627380371,5.560469627380371,5.59499979019165
+15,5.584121227264404,5.584121227264404,5.5788140296936035
+16,5.569767951965332,5.569767951965332,5.593851089477539
+17,5.562833786010742,5.562833309173584,5.580272674560547
+18,5.558739185333252,5.558739185333252,5.5749406814575195
+19,5.569887638092041,5.569887638092041,5.59522008895874
+20,5.548027515411377,5.548027515411377,5.575991153717041
+21,5.552359580993652,5.552359580993652,5.572174072265625
+22,5.557245254516602,5.557244777679443,5.568301677703857
+23,5.5531721115112305,5.5531721115112305,5.575419902801514
+24,5.551608562469482,5.551608562469482,5.5859551429748535
+25,5.547980785369873,5.547980785369873,5.567434787750244
+26,5.558878421783447,5.558878421783447,5.573901653289795
+27,5.572822570800781,5.572822570800781,5.58870792388916
+28,5.558602333068848,5.558602333068848,5.563601016998291
+29,5.554257869720459,5.554257392883301,5.5545220375061035
+30,5.559041500091553,5.559041500091553,5.554906845092773
+31,5.564431190490723,5.564431190490723,5.592169761657715
+32,5.570614337921143,5.570614337921143,5.566442966461182
+33,5.541515350341797,5.541515350341797,5.56029748916626
+34,5.551905632019043,5.551905632019043,5.564456939697266
+35,5.560442924499512,5.560442924499512,5.5753045082092285
+36,5.552741527557373,5.552741050720215,5.565302848815918
+37,5.552472114562988,5.552472114562988,5.564217567443848
+38,5.564489364624023,5.564489364624023,5.570488929748535
+39,5.560131549835205,5.560131549835205,5.577801704406738
+40,5.56781005859375,5.56781005859375,5.580154895782471
+41,5.559487342834473,5.559488296508789,5.557094097137451
+42,5.549378395080566,5.549378395080566,5.5669989585876465
+43,5.554733753204346,5.554733753204346,5.551562786102295
+44,5.542470455169678,5.5424699783325195,5.5577898025512695
+45,5.553576946258545,5.553576946258545,5.570135116577148
+46,5.560152530670166,5.560152530670166,5.565868854522705
+47,5.557215213775635,5.557215213775635,5.561673164367676
+48,5.553524017333984,5.553524017333984,5.55897331237793
+49,5.558567523956299,5.558567523956299,5.5570068359375
+50,5.561322212219238,5.561322212219238,5.563434600830078
+51,5.562047481536865,5.562047481536865,5.559027194976807
+52,5.554262161254883,5.5553483963012695,5.551726341247559
+53,5.554727077484131,5.554988861083984,5.562111854553223
+54,5.564721584320068,5.566233158111572,5.5689263343811035
+55,5.55735969543457,5.562169075012207,5.564178943634033
+56,5.548841953277588,5.54956579208374,5.55482292175293
+57,5.553441524505615,5.5556817054748535,5.549459934234619
+58,5.551535606384277,5.55885648727417,5.55819845199585
+59,5.548044204711914,5.556419372558594,5.549111843109131
+60,5.554605007171631,5.5576372146606445,5.55856466293335
+61,5.541609287261963,5.5446882247924805,5.542795181274414
+62,5.551147937774658,5.548646926879883,5.5556793212890625
+63,5.546630859375,5.547421932220459,5.549322128295898
+64,5.554196834564209,5.549775123596191,5.561946868896484
+65,5.555569171905518,5.558575630187988,5.5518670082092285
+66,5.5655622482299805,5.562289237976074,5.576994895935059
+67,5.556685924530029,5.547102928161621,5.552685737609863
+68,5.539889812469482,5.554402828216553,5.555147171020508
+69,5.552992820739746,5.554600715637207,5.5526347160339355
+70,5.563039302825928,5.557827472686768,5.560840129852295
+71,5.564231872558594,5.546955585479736,5.567209720611572
+72,5.550198554992676,5.556711673736572,5.553574562072754
+73,5.54953145980835,5.559920787811279,5.551657676696777
+74,5.554469108581543,5.542832851409912,5.54835844039917
+75,5.542092323303223,5.5447916984558105,5.555337429046631
+76,5.54931640625,5.550760269165039,5.557024955749512
+77,5.539743900299072,5.549052715301514,5.546274185180664
+78,5.5560760498046875,5.549284934997559,5.564888954162598
+79,5.550052165985107,5.553570747375488,5.553678035736084
+80,5.551743507385254,5.549341201782227,5.54948616027832
+81,5.550701141357422,5.562289237976074,5.554744720458984
+82,5.545189380645752,5.559484481811523,5.55905294418335
+83,5.5510478019714355,5.564089298248291,5.553728103637695
+84,5.5521464347839355,5.561048984527588,5.546876907348633
+85,5.550338268280029,5.563986301422119,5.553805351257324
+86,5.561927318572998,5.557210922241211,5.559905052185059
+87,5.551202774047852,5.563952922821045,5.550068378448486
+88,5.554379940032959,5.563688278198242,5.5552849769592285
+89,5.549014091491699,5.549198627471924,5.5551958084106445
+90,5.560193061828613,5.565088748931885,5.562371253967285
+91,5.557939052581787,5.5522332191467285,5.554843425750732
+92,5.562680244445801,5.553189277648926,5.557914733886719
+93,5.550804138183594,5.550830364227295,5.545635223388672
+94,5.540013790130615,5.566745281219482,5.540511608123779
+95,5.56000280380249,5.5555830001831055,5.565472602844238
+96,5.549357891082764,5.544004440307617,5.5490899085998535
+97,5.5422773361206055,5.561646461486816,5.544616222381592
+98,5.547672271728516,5.546198844909668,5.546009540557861
+99,5.547053337097168,5.553074359893799,5.554990291595459
+100,5.562369346618652,5.549615859985352,5.562082290649414

results_multiscale.csv

@@ -0,0 +1,4 @@
+optimizer,scale,n_params,steps,seed,final_train,final_val,final_ppl,avg_last_20,auc,total_time_s,tokens_per_sec,peak_mem_gb
+adamw,1M,837888,200,42,2.609443187713623,2.6979486656188967,14.84923970496087,2.6730387449264525,3.289305282831192,249.54487310000695,820.6940798095454,0.0
+scao,1M,837888,200,42,2.7116637229919434,2.798486475944519,16.41977622845508,2.768521749973297,3.3539043509960176,190.83978559996467,1073.1514885962956,0.0
+diag_shampoo,1M,837888,200,42,4.459746360778809,4.489192714691162,89.0495284775349,4.477831196784973,4.808437628746033,190.79926289990544,1073.3794087424722,0.0

results_multiscale_curves.csv

@@ -0,0 +1,61 @@
+optimizer,seed,wall_clock_s,ppl
+adamw,42,3.18,149.4621
+adamw,42,11.18,113.0973
+adamw,42,19.58,82.9878
+adamw,42,27.86,54.7381
+adamw,42,36.02,37.6566
+adamw,42,44.59,28.6875
+adamw,42,53.14,23.9316
+adamw,42,61.41,21.2133
+adamw,42,69.89,19.2509
+adamw,42,78.34,17.9698
+adamw,42,87.19,17.0501
+adamw,42,95.47,16.4875
+adamw,42,104.39,16.0428
+adamw,42,116.65,15.6604
+adamw,42,157.08,15.4163
+adamw,42,189.74,15.2580
+adamw,42,204.84,15.1107
+adamw,42,223.29,15.0021
+adamw,42,235.74,14.9159
+adamw,42,244.24,14.8492
+scao,42,3.47,150.3750
+scao,42,11.79,114.1543
+scao,42,20.34,84.9395
+scao,42,29.06,57.7818
+scao,42,37.52,39.7365
+scao,42,46.08,30.3128
+scao,42,54.66,25.3575
+scao,42,63.22,22.4175
+scao,42,71.86,20.3986
+scao,42,80.28,19.0522
+scao,42,88.86,18.4219
+scao,42,97.52,17.8990
+scao,42,106.12,17.4777
+scao,42,115.52,17.2133
+scao,42,124.05,16.9788
+scao,42,133.15,16.8036
+scao,42,141.87,16.6641
+scao,42,150.55,16.5728
+scao,42,162.65,16.4775
+scao,42,180.19,16.4198
+diag_shampoo,42,6.87,258.2252
+diag_shampoo,42,25.45,232.8457
+diag_shampoo,42,38.61,200.6362
+diag_shampoo,42,46.85,170.7445
+diag_shampoo,42,55.46,149.4159
+diag_shampoo,42,64.14,135.2993
+diag_shampoo,42,72.42,125.2282
+diag_shampoo,42,81.22,117.5615
+diag_shampoo,42,90.68,111.4450
+diag_shampoo,42,98.89,106.5383
+diag_shampoo,42,107.09,102.4727
+diag_shampoo,42,115.54,99.1991
+diag_shampoo,42,123.59,96.5694
+diag_shampoo,42,132.07,94.4825
+diag_shampoo,42,140.33,92.8692
+diag_shampoo,42,148.76,91.6371
+diag_shampoo,42,156.99,90.7070
+diag_shampoo,42,165.30,90.0230
+diag_shampoo,42,173.38,89.4939
+diag_shampoo,42,181.59,89.0495

results_multiscale_v2.csv

@@ -0,0 +1,3 @@
+optimizer,scale,n_params,steps,seed,final_train,final_val,final_ppl,avg_last_20,auc,total_time_s,tokens_per_sec,peak_mem_gb
+adamw,1M,837888,500,42,2.467369318008423,2.4726580572128296,11.85391339318089,2.4515873908996584,2.87281511592865,812.2171206999337,1260.7466327692784,0.01248563826084137
+scao,1M,837888,500,42,2.5562660694122314,2.5774046182632446,13.162930957876387,2.543643128871918,2.96584121799469,1292.7824406999862,792.0899663871888,0.012485504150390625

results_multiscale_v2_curves.csv

@@ -0,0 +1,41 @@
+optimizer,seed,wall_clock_s,ppl
+adamw,42,15.98,116.5229
+adamw,42,42.38,66.5564
+adamw,42,68.01,34.7519
+adamw,42,93.25,21.7346
+adamw,42,118.97,16.9728
+adamw,42,154.26,15.0195
+adamw,42,206.78,14.0978
+adamw,42,270.43,13.4651
+adamw,42,319.93,13.2326
+adamw,42,369.99,12.7835
+adamw,42,396.38,12.6075
+adamw,42,421.70,12.4974
+adamw,42,462.48,12.2746
+adamw,42,513.26,12.2202
+adamw,42,563.49,12.1451
+adamw,42,616.59,12.0276
+adamw,42,667.08,11.9785
+adamw,42,716.76,11.9383
+adamw,42,767.41,11.8942
+adamw,42,801.80,11.8539
+scao,42,16.33,117.5468
+scao,42,66.01,68.5892
+scao,42,120.64,36.3092
+scao,42,174.24,23.0127
+scao,42,248.61,19.3651
+scao,42,315.58,17.1757
+scao,42,388.50,16.1121
+scao,42,462.22,15.3712
+scao,42,526.31,14.7936
+scao,42,598.67,14.3948
+scao,42,675.21,14.1201
+scao,42,738.67,13.8888
+scao,42,796.70,13.7077
+scao,42,872.42,13.5892
+scao,42,942.70,13.5458
+scao,42,1007.91,13.3735
+scao,42,1075.33,13.3204
+scao,42,1142.23,13.2639
+scao,42,1204.05,13.2167
+scao,42,1267.80,13.1629

results_multiscale_v3.csv

@@ -0,0 +1,3 @@
+optimizer,scale,n_params,steps,seed,final_train,final_val,final_ppl,avg_last_20,auc,total_time_s,tokens_per_sec,peak_mem_gb
+adamw,1M,837888,200,42,2.6237757205963135,2.680860481262207,14.597648896524271,2.651871347427368,3.267444581985474,544.8934243000112,751.7066305694295,0.01248563826084137
+scao,1M,837888,200,42,2.755004644393921,2.8124120807647706,16.650031041261926,2.7797602295875548,3.355783064365387,719.2764292000793,569.4611742741574,0.02484196424484253

results_multiscale_v3_curves.csv

@@ -0,0 +1,41 @@
+optimizer,seed,wall_clock_s,ppl
+adamw,42,6.05,148.7067
+adamw,42,22.44,112.6560
+adamw,42,38.52,82.5069
+adamw,42,54.67,54.6896
+adamw,42,72.72,37.4680
+adamw,42,104.85,28.5375
+adamw,42,138.20,23.6325
+adamw,42,169.92,20.6589
+adamw,42,202.55,18.7545
+adamw,42,236.78,17.6057
+adamw,42,270.20,16.7743
+adamw,42,310.94,16.2468
+adamw,42,345.56,15.8245
+adamw,42,378.86,15.4520
+adamw,42,411.10,15.1534
+adamw,42,442.98,14.9748
+adamw,42,462.65,14.8452
+adamw,42,478.92,14.7617
+adamw,42,495.01,14.6753
+adamw,42,524.07,14.5976
+scao,42,11.56,149.6049
+scao,42,44.15,113.7046
+scao,42,81.87,84.5125
+scao,42,128.89,57.1945
+scao,42,180.01,39.5233
+scao,42,220.11,30.2746
+scao,42,254.97,25.1753
+scao,42,287.96,22.0837
+scao,42,321.88,20.8752
+scao,42,355.28,19.8731
+scao,42,388.20,19.0707
+scao,42,421.67,18.4424
+scao,42,453.84,17.9741
+scao,42,487.58,17.6508
+scao,42,521.59,17.3120
+scao,42,555.05,17.0866
+scao,42,587.47,16.9289
+scao,42,624.79,16.8306
+scao,42,662.70,16.7322
+scao,42,696.66,16.6500

results_multiscale_v4.csv

@@ -0,0 +1,3 @@
+optimizer,scale,n_params,steps,seed,final_train,final_val,final_ppl,avg_last_20,auc,total_time_s,tokens_per_sec,peak_mem_gb
+adamw,1M,837888,200,42,2.6237757205963135,2.680860481262207,14.597648896524271,2.651871347427368,3.267444581985474,701.1919704000466,584.1481609755365,0.01248563826084137
+scao,1M,837888,200,42,2.7310643196105957,2.792778310775757,16.32631642946134,2.759758996963501,3.2999178767204285,801.5647062000353,511.0005428529707,0.02490319311618805