[Feature] Plan to add model_register

[XinDongol]

Plan to add model_register? so users can add their own models

[wanchaol]

@XinDongol I think that's a good idea, the way this could work is that register_model can insert new entries into various module's module level dict (i.e. this), there are couple of things required from user:

1. user registered model must be a meta init ready model (with init_weights, from_model_args, etc.)
2. user need to provide at least the parallelization method simiar to parallelize_llama so that torchtitan knows how to shard the model.

In this way any plugged model should be able to run with torchtitan. We don't have time at the moment to implement such a feature, but feel free to submit a PR and we can help review it :)