It seems when wrap 'setup_model_and_optimizer' with context, adam_cls becomes torch.optim.adamw.AdamW instead of transformer_engine.pytorch.optimizers.fused_adam.FusedAdam, while torch AdamW does't have 'adam_w_mode' and some other args when init adam_cls
[rank7]: File "/data/minimax-dialogue/users/jager/Megatron-LM/megatron/training/training.py", line 1209, in setup_model_and_optimizer
[rank7]: optimizer = get_megatron_optimizer(
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/data/minimax-dialogue/users/jager/Megatron-LM/megatron/core/optimizer/init.py", line 604, in get_megatron_optimizer
[rank7]: _get_megatron_optimizer_based_on_param_groups(
[rank7]: File "/data/minimax-dialogue/users/jager/Megatron-LM/megatron/core/optimizer/init.py", line 394, in _get_megatron_optimizer_based_on_param_groups
[rank7]: optimizer = adam_cls(**kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^
[rank7]: TypeError: AdamW.init() got an unexpected keyword argument 'adam_w_mode'
It seems when wrap 'setup_model_and_optimizer' with context, adam_cls becomes
torch.optim.adamw.AdamWinstead oftransformer_engine.pytorch.optimizers.fused_adam.FusedAdam, while torch AdamW does't have 'adam_w_mode' and some other args when init adam_cls[rank7]: File "/data/minimax-dialogue/users/jager/Megatron-LM/megatron/training/training.py", line 1209, in setup_model_and_optimizer
[rank7]: optimizer = get_megatron_optimizer(
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/data/minimax-dialogue/users/jager/Megatron-LM/megatron/core/optimizer/init.py", line 604, in get_megatron_optimizer
[rank7]: _get_megatron_optimizer_based_on_param_groups(
[rank7]: File "/data/minimax-dialogue/users/jager/Megatron-LM/megatron/core/optimizer/init.py", line 394, in _get_megatron_optimizer_based_on_param_groups
[rank7]: optimizer = adam_cls(**kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^
[rank7]: TypeError: AdamW.init() got an unexpected keyword argument 'adam_w_mode'