Updates to support u-muP, as the new default behaviour

[DouglasOrr]

Demo notebook.

README header.

Changes:

• Add mults to nonlinear U.* ops and uu.* modules
• Add U.readout_linear, uu.ReadoutLinear to apply 1/fan_in scaling
• Add U.rms_norm, uu.RMSNorm, U.silu, uu.SiLU, U.silu_glu, U.mse_loss
• Add uu.Parameter, uu.optim.* to provide LR scaling
• Add uu.Trunk, uu.TransformerStack to set mup_scaling_depth and (for TransformerStack) apply appropriate residual mults
• Add demo notebook examples/demo.ipynb
• Change the default constraint from "gmean" -> "to_output_scale"
• Change default uu.Linear to bias=False & uu.*Norm to elementwise_affine=False
• Change U.softmax and U.scaled_dot_product_attention scaling rules to use the empirical fit
• Change definition of tau in residual_split/residual_add
• Change uu.MHSA to use U.scaled_dot_product_attention
• Change uu.MLP to use SwiGLU
• Change uu.TransformerLayer, uu.TransformerDecoder to use RMSNorm and various other tweaks
• Update top of README (other docs updates are lagging)
• +probably more

[DouglasOrr]

@lyprince, thank you for agreeing to review!

@thecharlieblake FYI

This is still WIP, but (I think/hope) the major changes are in.

[lyprince]

Without a linear specialisation (like linear_readout), this will have an extra sqrt(3) factor in the scale.

Our default implementation does not fuse the qkv matmul accordingly

[DouglasOrr]

Ah, yes, I wondered about this; the cheeky thing I was thinking is that when using the new-default constraint "to_output_scale" it shouldn't matter, as the scale just depends on fan_in. But maybe this is a bit cheeky as it can be overridden.

[lyprince]

the scale just depends on fan_in

Won't the output scale of this op be 3 * fan_in rather than fan_in though? I see in the demo notebook that attn_qkv.weight.grad.std = 0.62, which is pretty close to 1/sqrt(3) = 0.58.

[lyprince]

Looks good! I've done my best to check for correctness. I have some uncertainty around residual scaling. A test for unit scale preserved across depth with the residual scaling scheme would give me more confidence.

Other than that, my main comments are around

1. Documentation that points to relevant parts of paper for justifying design choices / changes.
2. Doubts over whether use of nn.Sequential is the right choice for keeping track of depth, given common patterns for implementing mask and positional embeddings.
3. Clarification on whether it is intentional to fuse qkv projection in MHSA.
4. Testing compatibility of Parameter class with torch.compile (does FakeTensor trigger deepcopy failure).

[lyprince]

Demo notebook looks good! Only thing that broke my attention was searching for show_layer_stats. Was easy to miss when colocated with imports. My preference would be to have it defined next to first use, or at least be in its own cell.