Support MuParametrization and MuTransfer

[NZ99]

This PR adds initial modeling-related changes to support MuP. It is incomplete, as the interaction between some modeling-related techniques is not yet clear to me (e.g. do we have to include any modifications as to make MuP compatible with AliBi?) but is shared for now as to facilitate discussion and collaboration, as agreed during community project meetings.

I'm uploading my code as is after leaving for NeurIPS -- I'll be checking whether anything important is missing or needs fixing in the code so far over the holidays. Apologies in case there are indeed issues with the code as pushed.

Important TODOs:

• check overall MuP implementation for correctness
• implement coordinate checking
• include the use of Optuna (or any alternative, desirable framework) to handle running the actual hparam optimization experiments

[NZ99]

Hey @othertea, thank you very much for these comments, they are really appreciated! Also sorry for the late reply, being sick with COVID over the holidays I have only got

• Re: coordinate checking and integrating Optuna (and general PR organization), I am happy to go either of the two ways. I have a similar sense w.r.t. the advantages of either approaches as you mention above, so maybe we can be satisfied with coordinate checking results for this PR and have a separate one for the full hparam search pipeline. I maybe err on the side of this approach, but I'm happy to go with any really.
• My feeling with how the code is organized in the codebases I used as inspiration is that (and this is just a guess to be clear) they wanted to maintain the modeling code as close to the original HF implementation as possible, leading to a separate global initialization function. I chose to go with handling initialization via each module's __init__ function mostly because I wanted to maintain flexibility and avoid stuff potentially breaking should we introduce new modules that we want to test. On the other hand, I'm unsure as to just how much flexibility this approach does provide in reality, mostly because if a new module is introduced without making sure it conforms to muP's rules we won't see any breaking change immediately, rather we'll just have hparam results stop transfering as e.g. model width is increased. I am happy to refactor the code either way in case a global initialization function is preferred.
• I'm unsure about relying on the MuP repo, since it does not seem to have received updates recently. I'm happy to use that too if that is the general preference though. One other option that I'd suggest (and that was not available back when I opened this PR) is ezmup, which is up to date w.r.t. the latest developments in MuP's general area (namely https://arxiv.org/abs/2310.17813, which I've not yet taken a look at), seems to be self-contained (e.g. it also contains MuP-compatible optimizers) and already implements coordinate checking. It would also allow us to keep modeling edits at a minimum, since all it seemingly requires is to set our width-delated hparams to a multiple of some large, prime "width_basis". I'm going to test this as IMO this would be a good solution, I just wonder about consequences for scaling both in this case and in the one integrating stuff from the mup repo (which, for example, does not seem to support FSDP yet).

Also, thank you very much for the detailed inline comments! I will take a look soon.

[NZ99]

Made a mistake and pushed from an older account. Anyway, I tried ezmup, but without much success. It runs into errors because it doesn't have support for embedding biases and transformers Conv1Ds out of the box, but even when fixing that it still errors out. Pity because it would have been an easy way to add support for MuP. I will keep testing it. In the meantime, I pushed some updates based on your extremely helpful review. I've kept organization as is for now, but happy to refactor depending on your (and other folks') preference. Re: MuReadout, there is also a MuSharedReadout (see also the dedicated comment mentioned above). I also want to add more documentation w.r.t. configuration options. Next up is the coordinate checking, looking into that today.

[NZ99]

Added the mup refactor and the coordinate checking integration discussed previously on Discord. You can find a colab showcasing mup coordinate checking results on the model here.

[NZ99]

Given that MuP now passes coordinate checking in my tests, maybe we can review (cc @othertea? would love to have a second pair of eyes given that I really don't trust myself much) and consider merging, with follow up work going to a to-be-opened second PR for optuna.

Let me know what you think :)