Function set_optimizer_flags

deepspeed/__init__.py:71–77

(config_class, model)

def set_optimizer_flags(config_class, model):
    if config_class.optimizer_name == MUON_OPTIMIZER:
        for name, p in model.named_parameters():
            if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")):
                setattr(p, "use_muon", True)
            else:
                setattr(p, "use_muon", False)


def initialize(args=None,