Make Muon optimizer easier to enable (#7555)

The original Muon optimizer PR
(#7509) requires user to
explicitly set `use_muon` flags in `model.parameters()`, as shown in
test
https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/ops/muon/test_muon.py#L27
.

This PR integrate setting of `use_muon` into DeepSpeed before engine
initialization. This makes Muon optimizer easier to use. User only needs
to change optimizer in `config.json` from `AdamW` to `Muon`, no need to
change code. It will solve the following issue
#7552

---------

Signed-off-by: Ma, Guokai <[email protected]>
Co-authored-by: Olatunji Ruwase <[email protected]>
Co-authored-by: Masahiro Tanaka <[email protected]>

Author: 3 people

deepspeed/__init__.py

@@ -28,7 +28,7 @@
 from .accelerator import get_accelerator
 from .constants import TORCH_DISTRIBUTED_DEFAULT_PORT
 from .runtime.engine import DeepSpeedEngine, DeepSpeedOptimizerCallable, DeepSpeedSchedulerCallable
-from .runtime.engine import ADAM_OPTIMIZER, LAMB_OPTIMIZER
+from .runtime.engine import ADAM_OPTIMIZER, LAMB_OPTIMIZER, MUON_OPTIMIZER
 from .runtime.hybrid_engine import DeepSpeedHybridEngine
 from .runtime.pipe.engine import PipelineEngine
 from .inference.engine import InferenceEngine
@@ -66,6 +66,15 @@ def _parse_version(version_str):
 dist = None
 
 
+def set_optimizer_flags(config_class, model):
+    if config_class.optimizer_name == MUON_OPTIMIZER:
+        for p in model.parameters():
+            if p.ndim >= 2:
+                setattr(p, "use_muon", True)
+            else:
+                setattr(p, "use_muon", False)
+
+
 def initialize(args=None,
                model: torch.nn.Module = None,
                optimizer: Optional[Union[Optimizer, DeepSpeedOptimizerCallable]] = None,
@@ -177,6 +186,7 @@ def initialize(args=None,
     assert config is not None, "DeepSpeed requires --deepspeed_config to specify configuration file"
     if not isinstance(model, PipelineModule):
         config_class = DeepSpeedConfig(config, mpu, mesh_device=mesh_device)
+        set_optimizer_flags(config_class, model)
         if config_class.hybrid_engine.enabled:
             engine = DeepSpeedHybridEngine(args=args,
                                            model=model,
@@ -206,6 +216,7 @@ def initialize(args=None,
         assert mpu is None, "mpu must be None with pipeline parallelism"
         mpu = model.mpu()
         config_class = DeepSpeedConfig(config, mpu)
+        set_optimizer_flags(config_class, model)
         engine = PipelineEngine(args=args,
                                 model=model,
                                 optimizer=optimizer,

tests/unit/ops/muon/test_muon.py

@@ -24,14 +24,6 @@
                     muon_configs.append([optimizer_name, stage, lr, model_dim, nlayer])
 
 
-def set_muon_flag(params):
-    for p in params:
-        if p.ndim >= 2:
-            setattr(p, "use_muon", True)
-        else:
-            setattr(p, "use_muon", False)
-
-
 @pytest.mark.parametrize('optimizer_type, zero_stage, lr, hidden_dim, nlayer', muon_configs)
 class TestMuonConfigs(DistributedTest):
 
@@ -55,8 +47,6 @@ def test(self, optimizer_type, zero_stage, lr, hidden_dim, nlayer):
         # Perform a few training steps to ensure the optimizer works correctly
 
         model = SimpleModel(hidden_dim=hidden_dim, nlayers=nlayer)
-        if 'muon' in optimizer_type:
-            set_muon_flag(model.parameters())
         initial_params = [p.clone().cpu() for p in model.parameters()]
         engine, optimizer, _, _ = deepspeed.initialize(
             config=config_dict,