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 model-side changes.

• Introduce extra muP params to training config
• Add Llama-194M config from @divya-kumari32
• Calculate base values that allow us to mimic the training behavior of default Llama-194M config
• Adjust param_init_fn to call reset_parameters with appropriate scale terms
• Adjust optimizer to handle multiple param_groups (0d, 1d, 2d with different LR scaling on each)
• Save singlefile checkpoint at end of training run
• Add reset_stepcount option, and enable taking 0 steps (i.e. single-file checkpoint model conversion)

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.