Terrific scaling milestone, congratulations on taking this from 46M to 325M! Having both the NPU checkpoint and the exact B200 GPU reference weights (b200_reference/) hosted together with evaluate_unseen.py makes this an exemplary systems study.
The paired window loss delta on unseen FineWeb-Edu text (-0.0001 with SE 0.0007 across 1,024 windows, winning on 50.1% of windows) is compelling evidence that numerical error did not accumulate or drift over 76k steps of backpropagation.
Did your custom training stack require any special handling for master weight precision or gradient accumulation/reduction across the 32 NPU chips compared to standard BF16/FP32 mixed precision on the B200? Looking forward to the technical paper!