MuP for scaling laws

[blahBlahhhJ]

Description

Support MuP in Levanter so that it's easier to do scaling laws analysis.

Reference:
MuP https://arxiv.org/pdf/2203.03466

WandB Link: https://wandb.ai/marin-community/marin/reports/621-MuP--VmlldzoxMTIxMTUzNQ?accessToken=h5qjejzau65v7bab5bau94hu8cltm1q9a0v6tdabocd398wagpnr6rjk4u8yc41a

Definition of Done

Implementation of MuP in Levanter merged. Could open another issue for experiments in marin once this is done.

[blahBlahhhJ]

Follow cerebasGPT's recipe (because it's the cleanest and has the most details)

[blahBlahhhJ]

This means:

Learning rate:

• all linear weights in transformer should be divided by scaling factor.
• Other layers remain unchanged.

Init var:

• all linear weights in transformer should be divided by scaling factor.
• Other layers remain unchanged.
• Note that O_proj and ffn2 should also be divided by 2*n_layers. I don't think this is currently implemented in levanter as well.

Forward:

• embedding output should be scaled by a constant m_emb
• attention should be scaled by 1/d_head instead of 1/sqrt(d_head).
• output logits shouuld be divided by scaling factor.

[dlwh]

in general I'd like to avoid changes that take us away from the llama architecture unless there's a very good reason, though I guess we can absorb things most things into the weights if we need to (for export)

[blahBlahhhJ]

Levanter PR: marin-community/levanter#863

[blahBlahhhJ]

Haliax PR: marin-community/haliax#116

[blahBlahhhJ]

Wandb report: https://wandb.ai/stanford-mercury/marin/reports/621-MuP--VmlldzoxMTIxMTUzNQ

Conclusion: not worth it. Our heuristic LR performs pretty well.