feat: add Muon optimizer integration and ShardingV3 support

Muon optimizer integration:
- Create Muon optimizer in trainer when `optim=muon`, with per-head
  QKV metadata annotation for fused QKV weight orthogonalisation
- Handle Muon's `_moment_acc_str` (vs AdamW's `_moment1_acc_str`)
  in optimizer state save/restore
- Add Muon `_muon_update`/`_apply_optimize` offload support in
  `offload_optimizer.py`

ShardingV3 support:
- Add `sharding_v3` training argument and `FLAGS_sharding_v3`
  environment variable dispatch
- Implement `DygraphShardingOptimizerV3` init path in
  `trainer_utils.py`
- Add V3 reshard logic (`reshard/sharding_v3.py`) for checkpoint
  save/restore
- Adapt `sharding_io.py`, `zero_cost_checkpoint.py`, and
  `moe_hybrid_parallel_optimizer.py` for V3 optimizer unwrapping

Tests:
- Add Muon smoke tests (`tests/muon/`) exercising both V2 and V3
  sharding paths on 2 GPUs with AMP O2

Author: xxyux

paddleformers/trainer/trainer.py

@@ -87,6 +87,13 @@
 from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
     DygraphShardingOptimizerV2,
 )
+
+try:
+    from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer_v3 import (
+        DygraphShardingOptimizerV3,
+    )
+except ImportError:
+    DygraphShardingOptimizerV3 = None
 from paddle.distributed.fleet.utils.hybrid_parallel_util import (
     fused_allreduce_gradients,
 )
@@ -211,7 +218,7 @@
 from .utils import reshard as reshard_util
 from .utils.async_save import AsyncSaver
 from .utils.ckpt_converter import CheckpointConverter
-from .utils.reshard import SHARDING_STRATEGY_V1, split_opt_state
+from .utils.reshard import SHARDING_STRATEGY_V1, SHARDING_STRATEGY_V3, split_opt_state
 from .utils.sharding_io import GroupGetter, to_device
 
 try:
@@ -1215,7 +1222,10 @@ def get_metadata_file_name(path):
         enable_bf16_opt = (
             not isinstance(self.model, LoRAModel)
             and self.args.bf16
-            and isinstance(self.optimizer._inner_opt, DygraphShardingOptimizerV2)
+            and isinstance(
+                self.optimizer._inner_opt,
+                (DygraphShardingOptimizerV2,) + ((DygraphShardingOptimizerV3,) if DygraphShardingOptimizerV3 else ()),
+            )
         )
         logger.debug(f"sharded_model_from_ema: {self.args.sharded_model_from_ema}")
         logger.debug(f"enable_bf16_opt: {enable_bf16_opt}")
@@ -1277,11 +1287,12 @@ def recover_params_from_master_weight(opt_state_dict, group):
                 del node_model_state_tmp
                 sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
                 logger.debug(f"sharding_strategy: {sharding_strategy}")
-                restore_func = (
-                    reshard_util.sharding_v1.restore
-                    if sharding_strategy == SHARDING_STRATEGY_V1
-                    else reshard_util.sharding_v2.restore
-                )
+                if sharding_strategy == SHARDING_STRATEGY_V1:
+                    restore_func = reshard_util.sharding_v1.restore
+                elif sharding_strategy == SHARDING_STRATEGY_V3:
+                    restore_func = reshard_util.sharding_v3.restore
+                else:
+                    restore_func = reshard_util.sharding_v2.restore
                 node_model_state = restore_func(node_model_state, self.model, self.optimizer)
                 node_model_state.unpack_keys()
                 master_weights = node_model_state.master_weights
@@ -1993,7 +2004,8 @@ def _inner_training_loop(
                         steps_trained_progress_bar.update(1)
                     if steps_trained_in_current_epoch == 0:
                         self._load_rng_state(resume_from_checkpoint)
-                    self.timers and self.timers("read-data").start()
+                    if self.args.ignore_data_skip:
+                        self.timers and self.timers("read-data").start()
                     # Reset data loading timer for skipped steps
                     _data_load_start_time = time.time()
                     continue
@@ -2930,6 +2942,15 @@ def apply_decay_param_fun(x):
             if hasattr(optimizer_cls, "_create_master_weight") and self.args.fp16_opt_level == "O2":
                 optimizer_kwargs["multi_precision"] = True
 
+            if self.args.optim.value == "muon":
+                # Attach per-head metadata to fused QKV weights so the Muon
+                # optimizer can orthogonalise each head independently.
+                for name, param in self.model.named_parameters():
+                    if "qkv_proj.weight" in name and len(param.shape) == 2:
+                        param.needs_qkv_split = True
+                        param.head_num = self.model.config.num_attention_heads
+                        param.kv_head_num = self.model.config.num_key_value_heads
+
             self.optimizer = optimizer_cls(
                 learning_rate=self.lr_scheduler if lr_scheduler is None else lr_scheduler,
                 apply_decay_param_fun=apply_decay_param_fun,
@@ -2947,6 +2968,7 @@ def apply_decay_param_fun(x):
     def _apply_to_optimizer(self, action):
         attributes = [
             ("_accumulators", "_moment1_acc_str"),
+            ("_accumulators", "_moment_acc_str"),  # Muon uses _moment_acc_str instead of _moment1_acc_str
             ("_accumulators", "_moment2_acc_str"),
             ("_master_weights",),
             ("_accumulators_holder",),
@@ -3070,6 +3092,18 @@ def get_optimizer_cls_and_kwargs(args: TrainingArguments) -> Tuple[Any, Any]:
 
             optimizer_cls = AdamWCustom
             optimizer_kwargs.update(adam_kwargs)
+        elif args.optim == OptimizerNames.MUON:
+            from paddle.optimizer import Muon
+
+            logger.info("Creating Muon optimizer")
+            muon_kwargs = {
+                **adam_kwargs,
+                "momentum": 0.95,
+                "muon_version": 3,
+                "is_split_qkv": True,
+            }
+            optimizer_cls = Muon
+            optimizer_kwargs.update(muon_kwargs)
         else:
             raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}")
 
@@ -4031,9 +4065,7 @@ def _save_checkpoint(self, model, metrics=None):
                                     global_rank, os.path.join(signal_dir, f".master_weight.done.{global_rank}")
                                 )
 
-                if self.args.save_checkpoint_format == "unified_checkpoint" and (
-                    self.args.offload_optim or self.args.tensorwise_offload_optimizer
-                ):
+                if self.args.offload_optim or self.args.tensorwise_offload_optimizer:
                     self._offload_optimizer()
             self.runtime_timer.stop()
 

paddleformers/trainer/trainer_utils.py

@@ -45,6 +45,13 @@
     DygraphShardingOptimizer,
     DygraphShardingOptimizerV2,
 )
+
+try:
+    from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer_v3 import (
+        DygraphShardingOptimizerV3,
+    )
+except ImportError:
+    DygraphShardingOptimizerV3 = None
 from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
 from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_optimizer_stage2 import (
     GroupShardedOptimizerStage2,
@@ -498,6 +505,7 @@ class OptimizerNames(ExplicitEnum):
     ADAFACTOR = "adafactor"
     ADAMW_MINI = "adamw_mini"
     ADAMW_CUSTOM = "adamw_custom"
+    MUON = "muon"
 
 
 class ShardingOption(ExplicitEnum):
@@ -1502,6 +1510,12 @@ def init_optimizer(optimizer, model_sharded_state_dict, state_dict_metadata):
         return
 
     elif DygraphShardingOptimizerV2 is not None and isinstance(inner_opt, DygraphShardingOptimizerV2):
+        # Unwrap to the innermost optimizer (e.g. Muon inside a sharding wrapper).
+        core_opt = optimizer._inner_opt
+        while hasattr(core_opt, "_inner_opt"):
+            core_opt = core_opt._inner_opt
+        is_muon_opt = type(core_opt).__name__ == "Muon"
+
         parameter_list = []
         for buffer in optimizer._comm_buffer_list:
             for param_name, grad_view in buffer._sharding_param_grad_view.items():
@@ -1515,11 +1529,77 @@ def init_optimizer(optimizer, model_sharded_state_dict, state_dict_metadata):
                     slice_param = paddle.slice(param_buffer, axes=[0], starts=[param_begin], ends=[param_end])
                     assert slice_param.numel().item() > 0
                     slice_param.name = param_name
+                    # Preserve original shape so Muon's should_use_muon() can identify 2-D weights.
+                    if is_muon_opt and hasattr(grad_view, "_param") and grad_view._param is not None:
+                        slice_param.original_shape = grad_view._param.shape
                     parameter_list.append(slice_param)
 
         optimizer._create_accumulators(paddle.base.framework.default_main_program().global_block(), parameter_list)
         return
 
+    elif DygraphShardingOptimizerV3 is not None and isinstance(inner_opt, DygraphShardingOptimizerV3):
+        # Unwrap to the innermost optimizer (e.g. Muon inside a V3 sharding wrapper).
+        core_opt = inner_opt._inner_opt
+        while hasattr(core_opt, "_inner_opt"):
+            core_opt = core_opt._inner_opt
+        is_muon_opt = type(core_opt).__name__ == "Muon"
+
+        parameter_list = []
+
+        # --- 1D params: build shard-sized slice params from FusedCommBuffer ---
+        # (same logic as V2 branch above, using _comm_buffer_list)
+        # IMPORTANT: set slice_param.name = "slice@" + param_name so that the
+        # accumulator key matches what V3's sharded_state_dict expects via
+        # _split_state_name (it strips the "_moment1_0" suffix to get static_name,
+        # which must match param_slice_info keys = original param names after
+        # removing the "slice@" prefix added back in sharded_state_dict).
+        for buffer in optimizer._comm_buffer_list:
+            for param_name, grad_view in buffer._sharding_param_grad_view.items():
+                if param_name not in static_to_struct_mapping:
+                    continue
+                struct_name = static_to_struct_mapping[param_name]
+                if not any(struct_name + state_name in state_dict_metadata for state_name in optimizer_state_names):
+                    continue
+                param_buffer = grad_view._param_buffer
+                param_begin = grad_view._param_begin
+                param_end = grad_view._param_end
+                if param_begin >= 0 and param_end > 0 and param_end > param_begin:
+                    slice_param = paddle.slice(param_buffer, axes=[0], starts=[param_begin], ends=[param_end])
+                    assert slice_param.numel().item() > 0
+                    # Use the original param name (no "slice@" prefix), consistent
+                    # with V3's _create_slice_param and V2's init_optimizer branch.
+                    slice_param.name = param_name
+                    parameter_list.append(slice_param)
+
+        # --- 2D non-MoE params: local rank's full tensors (Muon) ---
+        local_2d = optimizer._rank2params_2d.get(optimizer._sharding_rank, [])
+        for param in local_2d:
+            param_name = param.name
+            if param_name not in static_to_struct_mapping:
+                continue
+            struct_name = static_to_struct_mapping[param_name]
+            if not any(struct_name + state_name in state_dict_metadata for state_name in optimizer_state_names):
+                continue
+            parameter_list.append(param)
+
+        # --- 2D MoE expert params: local rank's full tensors (Muon) ---
+        if optimizer._moe_sharding_world_size > 1:
+            moe_rank = optimizer._moe_sharding_rank
+        else:
+            moe_rank = 0
+        local_2d_moe = optimizer._rank2params_2d_moe.get(moe_rank, [])
+        for param in local_2d_moe:
+            param_name = param.name
+            if param_name not in static_to_struct_mapping:
+                continue
+            struct_name = static_to_struct_mapping[param_name]
+            if not any(struct_name + state_name in state_dict_metadata for state_name in optimizer_state_names):
+                continue
+            parameter_list.append(param)
+
+        optimizer._create_accumulators(paddle.base.framework.default_main_program().global_block(), parameter_list)
+        return
+
     elif isinstance(optimizer, GroupShardedOptimizerStage2):
         local_params = optimizer._segment_params()[optimizer._rank]
         for p in local_params:

paddleformers/trainer/training_args.py

@@ -1528,6 +1528,17 @@ class TrainingArguments:
             "help": "Enable parameter sharding to distribute model parameters across devices, reducing memory footprint per GPU (ZeRO-style optimization)."
         },
     )
+    sharding_v3: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Enable ShardingV3 (hybrid tensor-wise + element-wise) for Muon optimizer. "
+                "2D Muon parameters are assigned as whole tensors to ranks (no sharding gather), "
+                "while non-2D AdamW parameters use element-wise splitting for memory balance. "
+                "Requires split_param=True and Muon optimizer. Set FLAGS_sharding_v3=1."
+            )
+        },
+    )
     sd_sharding_comm_overlap: bool = field(
         default=False,
         metadata={
@@ -2095,6 +2106,16 @@ def is_context_parallel_supported():
                             strategy.hybrid_configs["sharding_configs"].split_param = True
                             assert self.amp_master_grad, "Currently sharding stage1 v2 only support amp_master_grad"
 
+                        if self.sharding_v3:
+                            os.environ["FLAGS_sharding_v3"] = "1"
+                            assert self.split_param, "sharding_v3 requires split_param=True"
+                            logger.info("ShardingV3 enabled via sharding_v3=True")
+                        else:
+                            os.environ["FLAGS_sharding_v3"] = "0"
+
+                        if self.tensorwise_offload_optimizer:
+                            os.environ["FLAGS_tensorwise_offload_optimizer"] = "1"
+
                         if self.sd_release_grads:
                             strategy.hybrid_configs["sharding_configs"].release_gradients = True
 

paddleformers/trainer/utils/offload_optimizer.py

@@ -92,6 +92,60 @@ def new_insert_sync(self, sync_var, *args, **kwargs):
 
     setattr(opt_type, "_insert_sync", new_insert_sync)
 
+    # Step 4: mock Muon._muon_update and Muon._apply_optimize
+    # Muon's _muon_update is pure Python (paddle.lerp + paddle.assign),
+    # so it bypasses the _C_ops.adamw_ patch above. We need explicit
+    # reload/offload for Muon's momentum_buffer and master_weights.
+    try:
+        from paddle.optimizer.muon import Muon
+
+        # 4a: Patch _muon_update (staticmethod) — per-param momentum offload
+        origin_muon_update = Muon._muon_update
+
+        def new_muon_update(param, grad, lr, momentum_buffer, *args, **kwargs):
+            reload(momentum_buffer)
+            ret = origin_muon_update(param, grad, lr, momentum_buffer, *args, **kwargs)
+            is_offload_opt = getattr(param, "is_offload_opt", True)
+            if is_offload_opt:
+                offload(momentum_buffer)
+            return ret
+
+        Muon._muon_update = staticmethod(new_muon_update)
+
+        # 4b: Patch _apply_optimize — reload/offload master_weights around Muon updates
+        origin_muon_apply = Muon._apply_optimize
+
+        def new_muon_apply(self, loss, startup_program, params_grads):
+            # Reload master_weights to GPU before Muon update
+            # (needed after checkpoint restore where master_weights may be on CPU/pinned)
+            mw_dict = getattr(self, "_master_weights", None)
+            if mw_dict:
+                for param, grad in params_grads:
+                    if grad is None:
+                        continue
+                    mw = mw_dict.get(param.name)
+                    if mw is not None and isinstance(mw, paddle.Tensor):
+                        reload(mw)
+
+            ret = origin_muon_apply(self, loss, startup_program, params_grads)
+
+            # Offload master_weights back to CPU pinned after Muon update
+            if mw_dict:
+                for param, grad in params_grads:
+                    if grad is None:
+                        continue
+                    mw = mw_dict.get(param.name)
+                    if mw is not None and isinstance(mw, paddle.Tensor):
+                        is_offload_opt = getattr(param, "is_offload_opt", True)
+                        if is_offload_opt:
+                            offload(mw)
+            return ret
+
+        Muon._apply_optimize = new_muon_apply
+
+    except ImportError:
+        pass
+
 
 def hack_offload_optimizer_eb5():
     # Step 1: mock _add_accumulator

paddleformers/trainer/utils/reshard/__init__.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import pp_reshard, sharding_v1, sharding_v2
+from . import pp_reshard, sharding_v1, sharding_v2, sharding_v3
 from .common import (
     SHARDING_STRATEGY_V1,
     SHARDING_STRATEGY_V2,
+    SHARDING_STRATEGY_V3,
     NodeModelState,
     all_gather_state_dict,
     convert_opt_name_to_tname,

paddleformers/trainer/utils/reshard/common.py

@@ -21,6 +21,13 @@
     DygraphShardingOptimizer,
     DygraphShardingOptimizerV2,
 )
+
+try:
+    from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer_v3 import (
+        DygraphShardingOptimizerV3,
+    )
+except ImportError:
+    DygraphShardingOptimizerV3 = None
 from paddle.distributed.fleet.utils.log_util import logger
 
 from paddleformers.utils.tools import get_env_device
@@ -29,6 +36,7 @@
 
 SHARDING_STRATEGY_V1 = "ShardingV1"
 SHARDING_STRATEGY_V2 = "ShardingV2"
+SHARDING_STRATEGY_V3 = "ShardingV3"
 
 
 def is_sharding_opt(optimizer):
@@ -45,10 +53,18 @@ def check(cls):
         if check(DygraphShardingOptimizerV2):
             return True
 
+    if DygraphShardingOptimizerV3 is not None:
+        if check(DygraphShardingOptimizerV3):
+            return True
+
     return False
 
 
 def get_sharding_strategy(optimizer):
+    if DygraphShardingOptimizerV3 is not None:
+        tmp = unwrap_optimizer(optimizer, DygraphShardingOptimizerV3)
+        if tmp is not None:
+            return SHARDING_STRATEGY_V3
     if DygraphShardingOptimizerV2 is not None:
         tmp = unwrap_optimizer(optimizer, DygraphShardingOptimizerV2)
         if tmp is not None: