Optimizer refactor: clean up public get_megatron_optimizer interface and provide a more general API to support passing in different hyperparameters to subsets of parameters#2047