Minimal implementation of muP scaling for Llama

[daviswer]

Implement muP scaling for Llama models. Model follows muP scaling laws but introduces the minimal set of extra tunable hyperparameters that allows us to recover prior behavior - thus may not be compatible (yet) with existing muP configs. See here for training script changes.

• Introduce extra muP params to Llama config
• Calculate base values that allow us to mimic the training behavior of default Llama-194M config
• Swap out trunc_normal_ init for normal_ (redundant without specifying the clamp values, plus we noticed some FSDP-related issues in the speculator setting)
• Adjust reset_parameters so that each module initializes itself but no submodules (following FSDP init_fn contract)

Note that this is currently only implemented for Llama models, and does not support the old constant-range Llama init scheme. Additional work will be required to make these compatible; should we decide to support MuP then this is just a starting off / reference point.

[daviswer]

Noting that model equivalency checks are failing because muP introduces extra scaling terms into the forward pass, and the default mup scaling param values only correct for this at 194m scale.