OSP

[shahbuland]

OSP paper claims with a few modifications to standard transformer training you can enable 4bit quantization with less performance degradation. The proposed changes in the paper are:

1. Switch to Muon everywhere except embeddings
2. Use RMSNorm instead of LayerNorm, and with only a scale parameter
3. Single Scale RMS Norm: don't learn vector gain/scale, just learn a single scalar (i.e. not channel-wise)
4. All projections (i.e. timestep embedding, action embedding, patch embedding, text embedding) have an extra EMBPROJ layer added
5. EMBPROJ is normal linear, but no bias, and using orthogonal init
6. Monitor kurtosis of attn/mlp outputs