Is there any plan for Maximal Update Parametrization support ?

[starlitsky2010]

Motivation

Support Maximal Update Parametrization as it could help get optimal HPs from small proxy models and apply them on large LLMs
Is there any plan for adding this feature?
More details, please refer to : https://github.com/microsoft/mup

Description

It will acclerate the LLMs development as optimal HPs could be found by small proxy models and apply them on large LLMS. And this will save a lot of training cost. So there are quite a lot benefits for the LLM community.

Alternatives (optional)

No response

Additional context (optional)

No response

[contramundum53]

Thank you for your suggestion!
MuP needs tightly coupled model implementation (i.e. you need MuReadout at appropriate place and also it has a special attention module) and optimizer (i.e. you need MuAdam/MuSGD). It's not something that can be implemented in Optuna and that users can forget everything about, but is more an application.
You need to carefully write your model & your training pipeline (correctly!) with the mup library, and after that Optuna can help with the hyperparameter tuning.