Cherry-pick useful PRs from fleety (#11051)

* [LLM] fix normal sl (#10609)

* [LLM] support disable monkey patch (#10617)

* [LLM] use safer grad sync method when enabling sp (#10714)

* [RL] disable aistudio download and fix qwen bug (#10819)

* [RL] fix qwen load when fuse is enabled and modify gate presion to fp32 (#10842)

* 修复qwen3moe的一系列报错(justin) (#10818)

* [LLM] zcc support rng states (#10430) (#10485)

* [DLTP-85730] optimize save cost for zcc

* [LLM] fix moe using on tensor parallelism

* cherry-pick useful PRs from fleety and fix some typos

---------

Co-authored-by: aiyinyuedejustin <114650502+aiyinyuedejustin@users.noreply.github.com>

paddlenlp/trainer/trainer.py

@@ -140,6 +140,7 @@ from .trainer_callback import (
     DefaultFlowCallback,
     PrinterCallback,
     ProgressCallback,
+    SPGradSyncCallback,
     TrainerCallback,
     TrainerControl,
     TrainerState,
@@ -444,9 +445,8 @@ class Trainer:
             ), "should_save_sharding_stage1_model should be True when using zero cost checkpoint"
             assert (
                 ShardingOption.FULL_SHARD not in self.args.sharding
-            ), "FULL_SHARD is not supported when using zero cost checkpoint"
-            assert not self.args.save_tokenizer, "save_tokenizer is not supported when using zero cost checkpoint"
-            assert not self.args.save_rng_states, "save_rng_states is not supported when using zero cost checkpoint"
+            ), "FULL_SHARD is not supported when using flash save mode"
+            assert not self.args.save_tokenizer, "save_tokenizer is not supported when using flash save mode"
 
             # init attributes for zero cost checkpoint mode
             self.zcc_manager = None
@@ -2021,34 +2021,18 @@ class Trainer:
         if checkpoint is None:
             return
 
-        # if use distributed training
-        if self.args.world_size > 1:
-            process_index = self.args.process_index
-            rng_file_list = [None for x in range(self.args.world_size)]
-            if self.args.should_save:
-                rng_file = os.path.join(checkpoint, f"rng_state_{self.args.world_size}.pth")
-                if os.path.isfile(rng_file):
-                    rng_file_list = paddle.load(rng_file, return_numpy=True)
-            paddle.distributed.broadcast_object_list(rng_file_list, src=0)
-            # if rng_file_list still empty, not log rng state.
-            if rng_file_list[0] is None:
-                logger.info(
-                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
-                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
-                )
-                return
-            else:
-                checkpoint_rng_state = rng_file_list[process_index]
-        else:
-            rng_file = os.path.join(checkpoint, "rng_state.pth")
-            if not os.path.isfile(rng_file):
-                logger.info(
-                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
-                    "fashion, reproducibility is not guaranteed."
-                )
-                return
+        rng_file = os.path.join(checkpoint, f"rng_state_{dist.get_rank()}.pth")
+        if not os.path.isfile(rng_file):
+            logger.info(
+                "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
+                "fashion, reproducibility is not guaranteed."
+            )
+            return
 
-            checkpoint_rng_state = paddle.load(rng_file, return_numpy=True)
+        checkpoint_rng_state = paddle.load(rng_file, return_numpy=True)
+        if checkpoint_rng_state.get("world_size", None) != self.args.world_size:
+            logger.warn("Cannot load rng states when changing world size of training job.")
+            return
 
         random.setstate(checkpoint_rng_state["python"])
         np.random.set_state(checkpoint_rng_state["numpy"])
@@ -2210,11 +2194,6 @@ class Trainer:
             else:
                 model, self.optimizer = decorated
 
-        if self.args.tensor_parallel_degree > 1 and self.args.sequence_parallel:
-            register_sequence_parallel_allreduce_hooks(
-                model, self.args.gradient_accumulation_steps, self.args.fuse_sequence_parallel_allreduce
-            )
-
         if self.args.world_size == 1:
             if self.args.amp_master_grad:
                 mix_precision_utils.MixPrecisionLayer(model, dtype=self.amp_dtype)
@@ -2403,6 +2382,17 @@ class Trainer:
                 ):
                     self.optimizer._set_broadcast_overlap(True, model)
 
+        # use callback for sp grad sync in case of unexpected behaviour (except sharding stage 2&3)
+        if self.args.tensor_parallel_degree > 1 and self.args.sequence_parallel:
+            if ShardingOption.SHARD_GRAD_OP in self.args.sharding or ShardingOption.FULL_SHARD in self.args.sharding:
+                register_sequence_parallel_allreduce_hooks(
+                    unwrap_model(model),
+                    self.args.gradient_accumulation_steps,
+                    self.args.fuse_sequence_parallel_allreduce,
+                )
+            else:
+                self.add_callback(SPGradSyncCallback(model._layers))
+
         return model
 
     def _prepare_input(self, data: Union[paddle.Tensor, Any]) -> Union[paddle.Tensor, Any]:
@@ -2739,28 +2729,24 @@ class Trainer:
         if self.args.should_save:
             self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
 
-        # Save RNG state in non-distributed training
-        rng_states = {
-            "python": random.getstate(),
-            "numpy": np.random.get_state(),
-            "cuda": paddle.get_rng_state(),
-            "cpu": paddle.framework.core.default_cpu_generator().get_state(),
-        }
-        if self.args.use_hybrid_parallel:
-            rng_states[
-                "hybrid_parallel_rng_state_tracker"
-            ] = fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
+        if self.args.save_rng_states:
+            # Save RNG state in non-distributed training
+            rng_states = {
+                "python": random.getstate(),
+                "numpy": np.random.get_state(),
+                "cuda": paddle.get_rng_state(),
+                "cpu": paddle.framework.core.default_cpu_generator().get_state(),
+                "world_size": self.args.world_size,
+            }
+            if self.args.use_hybrid_parallel:
+                rng_states[
+                    "hybrid_parallel_rng_state_tracker"
+                ] = fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
 
         if self.args.save_rng_states:
-            if self.args.world_size > 1:
-                rng_states_list = []
-                paddle.distributed.all_gather_object(rng_states_list, rng_states)
-                if self.args.should_save:
-                    os.makedirs(output_dir, exist_ok=True)
-                    paddle.save(rng_states_list, os.path.join(output_dir, f"rng_state_{self.args.world_size}.pth"))
-            else:
-                os.makedirs(output_dir, exist_ok=True)
-                paddle.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
+            rng_state_file = os.path.join(output_dir, f"rng_state_{dist.get_rank()}.pth")
+            os.makedirs(output_dir, exist_ok=True)
+            paddle.save(rng_states, rng_state_file)
 
         # only save model state dict, ignore optimizer and scheduler
         if not self.args.ignore_save_lr_and_optim:

paddlenlp/trainer/trainer_callback.py

@@ -20,10 +20,18 @@ Callbacks to use with the Trainer class and customize the training loop.
 """
 import dataclasses
 import json
+import time
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 import numpy as np
+from paddle.distributed.fleet import fleet
+from paddle.distributed.fleet.utils.hybrid_parallel_util import (
+    fused_allreduce_gradients_with_group,
+)
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+    is_sequence_parallel_parameter,
+)
 from tqdm.auto import tqdm
 
 from paddlenlp.utils.log import logger
@@ -609,3 +617,33 @@ class EarlyStoppingCallback(TrainerCallback):
         self.check_metric_value(args, state, control, metric_value)
         if self.early_stopping_patience_counter >= self.early_stopping_patience:
             control.should_training_stop = True
+
+
+class SPGradSyncCallback(TrainerCallback):
+    """
+    SPGradSyncCallback
+    只能在非 sharding stage2 的情况下使用。
+    开启sharding stage2 时，在 `on_optimizer_begin` 的时候 grad 已经被清空了
+    """
+
+    def __init__(self, model):
+        assert hasattr(fleet, "_hcg"), "must use MP when calling this Callback"
+        logger.info("using sp callback")
+        params = []
+        self.model = model
+        for n, p in model.named_parameters():
+            if is_sequence_parallel_parameter(p):
+                logger.info(f"register bw hook for:{n}")
+                params.append(p)
+
+        logger.info(f"#-sp-sync param:{len(params)}")
+        self._sp_params = params
+
+    def on_optimizer_begin(self, args, state, control, **kwargs):
+        """on_optimizer_begin"""
+        if self._sp_params:
+            now = time.time()
+            mp_group = fleet.get_hybrid_communicate_group().get_model_parallel_group()
+            fused_allreduce_gradients_with_group(self._sp_params, group=mp_group, scale=1.0)  # sum not mean
+            another_time = time.time()
+            logger.info(f"sync gradients takes {another_time - now} time")

paddlenlp/trainer/utils/zero_cost_checkpoint.py

@@ -18,10 +18,12 @@ import hashlib
 import json
 import multiprocessing
 import os
+import random
 import time
 from collections import OrderedDict
 from enum import Enum
 
+import numpy as np
 import paddle
 import paddle.autograd as imperative_base
 import paddle.distributed as dist
@@ -414,10 +416,26 @@ class ZeroCostCheckpointCallback(TrainerCallback):
             self.maybe_update_zcc_worker(args, model, optimizer, state.global_step)
             checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
             save_infos = self._get_save_infos_based_on_steps(state, args, checkpoint_folder)
-            non_cached_objects = (lr_scheduler.state_dict(), copy.deepcopy(state))
+            non_cached_objects = (lr_scheduler.state_dict(), state, self.get_rng_states(args))
             self.manager.get_idle_worker_for_saving((save_infos, non_cached_objects))
             self.runtime_timer.stop()
 
+    def get_rng_states(self, args):
+        if not args.save_rng_states:
+            return None
+        rng_states = {
+            "python": random.getstate(),
+            "numpy": np.random.get_state(),
+            "cuda": paddle.get_rng_state(),
+            "cpu": paddle.framework.core.default_cpu_generator().get_state(),
+            "world_size": args.world_size,
+        }
+        if args.use_hybrid_parallel:
+            rng_states[
+                "hybrid_parallel_rng_state_tracker"
+            ] = dist.fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
+        return rng_states
+
     def _get_save_infos_based_on_steps(self, state, args, checkpoint_folder):
         flash_device_checkpoint_dir = None
         persistent_checkpoint_dir = None
@@ -701,6 +719,7 @@ class ZeroCostCheckpointWorker:
         # TODO(@gexiao): remove lr scheduler saves
         self.lr_scheduler = None
         self.trainer_state = None
+        self.rng_state = None
 
         # for dumping
         self.flash_device_save_dir = None
@@ -734,7 +753,7 @@ class ZeroCostCheckpointWorker:
             return
         save_infos, non_cached_objects = prepares
         self.flash_device_save_dir, self.persistent_save_dir = save_infos
-        self.lr_scheduler, self.trainer_state = non_cached_objects
+        self.lr_scheduler, self.trainer_state, self.rng_state = non_cached_objects
 
     def process_offload_task(self, dump, global_step):
         """
@@ -897,6 +916,11 @@ class ZeroCostCheckpointWorker:
         if self.device_id == 0:
             self.trainer_state.save_to_json(trainer_state_name_path)
 
+        # Step2.5: save RNG State
+        if self.rng_state is not None:
+            rng_state_name_path = os.path.join(output_dir, f"rng_state_{dist.get_rank()}.pth")
+            paddle.save(self.rng_state, rng_state_name_path)
+
         # Step3: dump save signals
         saved_signal_path = os.path.join(output_dir, f"saved_signal_{self.global_rank}")
         with open(saved_signal_path, mode="w+") as f:

paddlenlp/transformers/moe_gate.py

@@ -326,6 +326,7 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
             logits += self.gumbel_rsample(logits.shape)
 
         gates = self.gate_score_func(logits=logits)
+
         capacity = self._capacity(gates, self.capacity_factor, self.max_capacity, self.min_capacity)
 
         # Create a mask for 1st's expert per token
@@ -396,6 +397,7 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
         logits: paddle.Tensor,
     ) -> Tuple[int, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:
         # everything is in fp32 in this function
+
         gates = self.gate_score_func(logits=logits)
 
         # Create a mask for 1st's expert per token.

paddlenlp/transformers/moe_layer.py

@@ -24,6 +24,8 @@ import paddle.distributed as dist
 from paddle import Tensor, nn
 from paddle.distributed.communication.group import Group
 
+from paddlenlp.utils.log import logger
+
 from .moe_gate import PretrainedMoEGate
 from .token_dispatcher import MoEFlexTokenDispatcher
 
@@ -155,12 +157,12 @@ class MoELayer(nn.Layer):
     def __init__(
         self,
         config,
-        moe_num_experts: int,
+        moe_num_experts: int,  # 128
         expert_class: nn.Layer,
         expert_kwargs: dict,
         gate: PretrainedMoEGate,
         capacity: int = 1.0,
-        moe_group: str = "data",
+        moe_group: str = "tp",  # will be re-assigned from config
         all_to_all_dropout=0.0,
     ):
         super().__init__()
@@ -169,6 +171,8 @@ class MoELayer(nn.Layer):
 
         self.moe_num_experts = moe_num_experts
         self.capacity = capacity
+        self.is_tp_moe = False
+        self.is_dp_moe = False
 
         try:
             dist.fleet.get_hybrid_communicate_group()
@@ -179,7 +183,7 @@ class MoELayer(nn.Layer):
         if (
             is_fleet_init
             and dist.fleet.get_hybrid_communicate_group().get_data_parallel_world_size() > 1
-            and moe_group == "data"
+            and moe_group == "dp"
         ):
             self.moe_group = dist.fleet.get_hybrid_communicate_group().get_data_parallel_group()
             self.moe_rank = dist.get_rank(self.moe_group)
@@ -190,8 +194,29 @@ class MoELayer(nn.Layer):
                 self.moe_num_experts, self.expert_parallel_degree
             )
             self.is_dummy_moe = False if self.expert_parallel_degree > 1 else True
+            self.is_dp_moe = True
+        elif (
+            is_fleet_init
+            and dist.fleet.get_hybrid_communicate_group().get_model_parallel_world_size() > 1
+            and moe_group == "tp"
+        ):
+            # for qwen3moe,moe_group should be "tp", since dp always == 1
+            self.moe_group = dist.fleet.get_hybrid_communicate_group().get_model_parallel_group()
+
+            self.moe_rank = dist.get_rank(self.moe_group)  # i for num_worker in a TP group
+            self.moe_rank = 0 if self.moe_rank < 0 else self.moe_rank  # 1
+
+            self.expert_parallel_degree = dist.get_world_size(self.moe_group)
+
+            self.expert_parallel_degree = 1 if self.expert_parallel_degree < 0 else self.expert_parallel_degree  # 4
+
+            self.moe_num_experts_per_device = self._parse_moe_expert_parallel(
+                self.moe_num_experts, self.expert_parallel_degree
+            )  # e.g. 单机2路tp， 那么 32  = 128/4
+
+            self.is_dummy_moe = False if self.expert_parallel_degree > 1 else True  # False
+            self.is_tp_moe = True
         else:
-            # when moe_group is dummy, we don't need to use all_to_all
             self.moe_group = None
             self.moe_rank = 0
             self.expert_parallel_degree = 1
@@ -210,6 +235,7 @@ class MoELayer(nn.Layer):
 
         self.gate = gate
         self.gate.group = self.moe_group
+
         self._post_init()
 
     def _parse_moe_expert_parallel(self, moe_num_experts, expert_parallel_degree):
@@ -219,6 +245,7 @@ class MoELayer(nn.Layer):
         assert (
             moe_num_experts % expert_parallel_degree == 0
         ), f"expert moe_num_experts={moe_num_experts} % moe_world_size={expert_parallel_degree} == 0"
+
         moe_num_experts_per_device = moe_num_experts // expert_parallel_degree
         return moe_num_experts_per_device
 
@@ -229,9 +256,11 @@ class MoELayer(nn.Layer):
         for k in self.experts:
             if k is not None:
                 for p in k.parameters():
-                    p.expert = not self.is_dummy_moe
-                    p.no_sync = not self.is_dummy_moe
-                    # logger.info(f"expert param={p.name}, no-sync={p.no_sync}")
+                    p.expert = not (self.is_tp_moe or self.is_dummy_moe)  # type: ignore
+                    p.no_sync = not (self.is_tp_moe or self.is_dummy_moe)
+                    logger.info(f"expert param={p.name}, no-sync={p.no_sync}")
+                    if self.is_tp_moe or self.is_dp_moe:
+                        p.is_distributed = True
 
     def forward(
         self,
@@ -257,10 +286,12 @@ class MoELayer(nn.Layer):
         # topk_ids    : sk
         # token_priority    : se
         # self.exp_counts  :
-        capacity, topk_weight, topk_ids, token_priority, l_aux, l_zloss = self.gate(hidden_state)
+
+        capacity, topk_weight, topk_ids, token_priority, l_aux, l_zloss = self.gate(hidden_state)  # here
 
         """MoE expert dispatch from: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py"""
         cnts = paddle.zeros([topk_ids.shape[0], len(self.experts)], dtype=topk_ids.dtype)
+
         cnts = cnts.put_along_axis(topk_ids, 1, axis=1)
 
         tokens_per_expert = cnts.sum(axis=0)
@@ -271,6 +302,7 @@ class MoELayer(nn.Layer):
 
         if self.expert_parallel_degree > 1:
             tokens_per_ep_rank = tokens_per_expert.reshape([self.expert_parallel_degree, -1]).sum(axis=1)
+
             tokens_per_expert_group = _AllToAll.apply(
                 [tokens_per_expert.shape[0]], tokens_per_expert, group=self.moe_group
             )
@@ -278,6 +310,7 @@ class MoELayer(nn.Layer):
                 tokens_per_expert_group.reshape([self.expert_parallel_degree, -1]).sum(axis=1).cpu().tolist()
             )
             input_split_sizes = tokens_per_ep_rank.cpu().tolist()
+
             gathered_tokens = _AllToAll.apply(
                 [tokens_per_expert_group.sum(axis=0).cpu().item(), sorted_tokens.shape[1]],
                 sorted_tokens,
@@ -290,25 +323,35 @@ class MoELayer(nn.Layer):
                 [self.expert_parallel_degree, self.moe_num_experts_per_device]
             ).sum(axis=0)
             gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)
+
             s = 0
             for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
                 gatherd_idxs[s : s + k] = i % self.moe_num_experts_per_device
                 s += k
             gatherd_idxs = gatherd_idxs.argsort()
             sorted_tokens = gathered_tokens[gatherd_idxs]
+
             tokens_per_expert = tokens_per_expert_post_gather
 
         outputs = []
         start_idx = 0
+
         for i, num_tokens in enumerate(tokens_per_expert):
             end_idx = start_idx + num_tokens
+
             if num_tokens == 0:
                 continue
+
             expert = self.experts[i + self.moe_rank * self.moe_num_experts_per_device]
+
             tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+
             expert_out = expert(tokens_for_this_expert)
+
             outputs.append(expert_out)
+
             start_idx = end_idx
+
         outs = paddle.concat(outputs, axis=0) if len(outputs) > 0 else paddle.to_tensor(0, dtype=sorted_tokens.dtype)
         if self.expert_parallel_degree > 1:
             new_x = paddle.empty_like(outs)
@@ -324,6 +367,7 @@ class MoELayer(nn.Layer):
 
         new_x = paddle.empty_like(outs)
         new_x[idxs] = outs
+
         final_out = (
             new_x.reshape(topk_ids.shape + [-1])
             .astype(topk_weight.dtype)
@@ -339,7 +383,6 @@ class MoELayer(nn.Layer):
 
 class MoEFlexTokenLayer(nn.Layer):
     def __init__(self, config, moe_num_experts, expert_class, expert_kwargs, gate, moe_group):
-
         super().__init__()
         self.config = config
         self.moe_group = moe_group
@@ -357,7 +400,6 @@ class MoEFlexTokenLayer(nn.Layer):
     def expert_forward(self, dispatched_input, tokens_per_expert):
         outputs = []
         tokens_per_expert = tokens_per_expert.tolist()
-        # print(f"all tokens: {sum(tokens_per_expert)}, detail: {tokens_per_expert}")
         chunks = paddle.split(dispatched_input, num_or_sections=tokens_per_expert, axis=0)
         for chunk, expert in zip(chunks, self.experts):
             chunk = chunk.contiguous()

paddlenlp/transformers/qwen2/modeling.py

@@ -18,6 +18,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Paddle Qwen2 model."""
+
 from __future__ import annotations
 
 import math
@@ -740,6 +741,7 @@ class Qwen2Attention(nn.Layer):
                 training=self.training,
                 sequence_parallel=self.sequence_parallel,
             )
+
         if output_attentions:
             attn_output, attn_weights = outputs
         else:
@@ -928,7 +930,6 @@ class Qwen2PretrainedModel(PretrainedModel):
 
     @classmethod
     def _get_tensor_parallel_mappings(cls, config: Qwen2Config, is_split=True):
-
         from paddlenlp.transformers.conversion_utils import split_or_merge_func
 
         fn = split_or_merge_func(
@@ -1042,7 +1043,11 @@ class Qwen2PretrainedModel(PretrainedModel):
                     for fuse_keys in fuse_qkv_keys:
                         keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_keys])
                         final_actions[keys] = partial(
-                            fn, split_nums=3, is_qkv=True, num_heads=num_heads, num_key_value_heads=num_key_value_heads
+                            fn,
+                            split_nums=3,
+                            is_qkv=True,
+                            num_heads=num_heads,
+                            num_key_value_heads=num_key_value_heads,
                         )
             if not fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
@@ -1270,7 +1275,6 @@ class Qwen2Model(Qwen2PretrainedModel):
         return_dict: Optional[bool] = None,
         attn_mask_startend_row_indices=None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
-
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states  # fmt:skip
         use_cache = use_cache if use_cache is not None else self.config.use_cache

paddlenlp/transformers/qwen2_moe/modeling.py

@@ -441,7 +441,15 @@ class Qwen2MoeMLP(nn.Layer):
             ColumnParallelLinear = linear_utils.ColumnParallelLinear
             RowParallelLinear = linear_utils.RowParallelLinear
 
-        if config.tensor_parallel_degree > 1:
+        if config.moe_group == "tp":
+            use_parallel_linear = False
+        elif config.moe_group == "dp":
+            use_parallel_linear = True
+        else:
+            use_parallel_linear = False
+
+        # 只有在张量并行度 > 1 且 *不是* 用于专家并行(EP)的场景下，才使用并行化的Linear层。
+        if use_parallel_linear:
             if self.fuse_attention_ffn:
                 self.gate_up_fused_proj = ColumnParallelLinear(
                     self.hidden_size,
@@ -469,6 +477,8 @@ class Qwen2MoeMLP(nn.Layer):
                 has_bias=False,
             )
         else:
+            # 1. 单卡或非tp (config.tensor_parallel_degree <= 1)
+            # 2. tp，但此MLP被用作一个完整的专家
             if self.fuse_attention_ffn:
                 self.gate_up_fused_proj = Linear(self.hidden_size, self.intermediate_size * 2, bias_attr=False)
             else:
@@ -755,7 +765,7 @@ class Qwen2MoeGate(PretrainedMoEGate):
         # [hidden_size, n_expert]
         self.weight = paddle.create_parameter(
             shape=[expert_hidden_size, num_experts],
-            dtype=paddle.get_default_dtype(),
+            dtype="float32",
             is_bias=False,
             default_initializer=nn.initializer.Constant(1.0),
         )
@@ -772,6 +782,7 @@ class Qwen2MoeGate(PretrainedMoEGate):
 
         with paddle.amp.auto_cast(False):
             scores = self.gate_score_func(logits=logits)
+
             scores = scores.cast(paddle.get_default_dtype())
 
         capacity, combine_weights, dispatch_mask, exp_counts, l_aux, l_zloss = self.topkgating(scores)
@@ -1020,14 +1031,18 @@ class Qwen2MoePretrainedModel(PretrainedModel):
                 base_actions.pop("embed_tokens.weight")
 
             # Column Linear
-            base_actions["layers.0.self_attn.q_proj.weight"] = partial(fn, is_column=True)
-            base_actions["layers.0.self_attn.q_proj.bias"] = partial(fn, is_column=True)
-            # if we have enough num_key_value_heads to split, then split it.
-            if config.num_key_value_heads % config.tensor_parallel_degree == 0:
-                base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True)
-                base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True)
-                base_actions["layers.0.self_attn.k_proj.bias"] = partial(fn, is_column=True)
-                base_actions["layers.0.self_attn.v_proj.bias"] = partial(fn, is_column=True)
+            if config.fuse_attention_qkv:
+                base_actions["layers.0.self_attn.qkv_proj.weight"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.qkv_proj.bias"] = partial(fn, is_column=True)
+            else:
+                base_actions["layers.0.self_attn.q_proj.weight"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.q_proj.bias"] = partial(fn, is_column=True)
+                # if we have enough num_key_value_heads to split, then split it.
+                if config.num_key_value_heads % config.tensor_parallel_degree == 0:
+                    base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True)
+                    base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True)
+                    base_actions["layers.0.self_attn.k_proj.bias"] = partial(fn, is_column=True)
+                    base_actions["layers.0.self_attn.v_proj.bias"] = partial(fn, is_column=True)
 
             for key, action in base_actions.items():
                 if "layers.0." in key:
@@ -1036,11 +1051,20 @@ class Qwen2MoePretrainedModel(PretrainedModel):
                 final_actions[key] = action
 
             # Add tp split for expert params.
-            base_actions = {
-                "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
-                "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
-            }
+            if config.fuse_attention_ffn:
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_up_fused_proj.weight": partial(
+                        fn, is_column=True, is_naive_2fuse=True
+                    ),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
+            else:
+                # Add tp split for expert params.
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
             for key, action in base_actions.items():
                 for i in range(num_layers):
                     newkey = key.replace("layers.0.", f"layers.{i}.")
@@ -1049,11 +1073,19 @@ class Qwen2MoePretrainedModel(PretrainedModel):
                         final_actions[newkey2] = action
 
             # Add tp split for shared expert params.
-            base_actions = {
-                "layers.0.mlp.shared_expert.gate_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.shared_expert.up_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
-            }
+            if config.fuse_attention_ffn:
+                base_actions = {
+                    "layers.0.mlp.shared_expert.gate_up_fused_proj.weight": partial(
+                        fn, is_column=True, is_naive_2fuse=True
+                    ),
+                    "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
+                }
+            else:
+                base_actions = {
+                    "layers.0.mlp.shared_expert.gate_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.shared_expert.up_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
+                }
             for key, action in base_actions.items():
                 if "layers.0." in key:
                     for i in range(num_layers):
@@ -1090,24 +1122,24 @@ class Qwen2MoePretrainedModel(PretrainedModel):
         ]
 
         fuse_gate_up_keys = (
-            "layers.0.mlp.gate_proj.weight",
-            "layers.0.mlp.up_proj.weight",
-            "layers.0.mlp.gate_up_fused_proj.weight",
+            "layers.0.mlp.experts.0.gate_proj.weight",
+            "layers.0.mlp.experts.0.up_proj.weight",
+            "layers.0.mlp.experts.0.gate_up_fused_proj.weight",
         )
         num_heads = config.num_attention_heads
         num_key_value_heads = getattr(config, "num_key_value_heads", num_heads)
         fuse_attention_qkv = getattr(config, "fuse_attention_qkv", False)
         fuse_attention_ffn = getattr(config, "fuse_attention_ffn", False)
+        num_experts = getattr(config, "num_experts", 128)
 
         final_actions = {}
         if is_fuse:
             if fuse_attention_qkv:
                 for i in range(config.num_hidden_layers):
-                    for fuse_keys in fuse_qkv_keys:
-                        keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_keys])
-                        final_actions[keys] = partial(
-                            fn, is_qkv=True, num_heads=num_heads, num_key_value_heads=num_key_value_heads
-                        )
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = fn
             if fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
                     keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
@@ -1118,12 +1150,18 @@ class Qwen2MoePretrainedModel(PretrainedModel):
                     for fuse_keys in fuse_qkv_keys:
                         keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_keys])
                         final_actions[keys] = partial(
-                            fn, split_nums=3, is_qkv=True, num_heads=num_heads, num_key_value_heads=num_key_value_heads
+                            fn,
+                            split_nums=3,
+                            is_qkv=True,
+                            num_heads=num_heads,
+                            num_key_value_heads=num_key_value_heads,
                         )
             if not fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
-                    keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
-                    final_actions[keys] = partial(fn, split_nums=2)
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = partial(fn, split_nums=2)
         return final_actions
 
     def _init_weights(self, layer):
@@ -1569,7 +1607,7 @@ class Qwen2MoeForCausalLM(Qwen2MoePretrainedModel):
         attention_mask=None,
         inputs_embeds=None,
         output_router_logits=False,
-        **kwargs
+        **kwargs,
     ):
         batch_size, seq_length = input_ids.shape
         position_ids = kwargs.get("position_ids", paddle.arange(seq_length).expand((batch_size, seq_length)))

paddlenlp/transformers/qwen3_moe/configuration.py

@@ -174,6 +174,7 @@ class Qwen3MoeConfig(PretrainedConfig):
         output_router_logits=False,
         router_aux_loss_coef=0.001,
         mlp_only_layers=None,
+        moe_group=None,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -209,6 +210,7 @@ class Qwen3MoeConfig(PretrainedConfig):
         self.output_router_logits = output_router_logits
         self.router_aux_loss_coef = router_aux_loss_coef
         self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+        self.moe_group = moe_group
 
         super().__init__(
             tie_word_embeddings=tie_word_embeddings,

paddlenlp/transformers/qwen3_moe/modeling.py

@@ -94,6 +94,7 @@ class ExpertParallelQwen3MoeSparseMoeBlock(MoELayer):
             expert_kwargs={"config": config},
             gate=gate,
             capacity=2.0,
+            moe_group=config.moe_group,
         )
 
         self.top_k = config.num_experts_per_tok
@@ -101,7 +102,6 @@ class ExpertParallelQwen3MoeSparseMoeBlock(MoELayer):
 
     def forward(self, hidden_states):
         final_hidden_states, l_aux, l_zloss = super().forward(hidden_states)
-
         return final_hidden_states, l_aux
 
 
@@ -191,6 +191,7 @@ class Qwen3MoeDecoderLayer(nn.Layer):
         past_key_value: Optional[Tuple[paddle.Tensor]] = None,
         use_cache: Optional[bool] = False,
         attn_mask_startend_row_indices: Optional[paddle.Tensor] = None,
+        batch_size=None,
         **kwargs,
     ) -> Tuple[paddle.Tensor, Optional[Tuple[paddle.Tensor, paddle.Tensor]]]:
         """
@@ -215,6 +216,8 @@ class Qwen3MoeDecoderLayer(nn.Layer):
 
         hidden_states = self.input_layernorm(hidden_states)
 
+        sp_batch_size = kwargs.pop("batch_size", None) if batch_size is None else batch_size
+
         # Self Attention
         has_gradient = not hidden_states.stop_gradient
         if (
@@ -233,16 +236,20 @@ class Qwen3MoeDecoderLayer(nn.Layer):
                 use_cache,
                 attn_mask_startend_row_indices,
                 use_reentrant=self.config.recompute_use_reentrant,
+                batch_size=sp_batch_size,  # Qwen3Attention 有batch_szie这个参数，手动传递
+                **kwargs,
             )
         else:
-            outputs = self.self_attn(
+            outputs = self.self_attn(  # Qwen3Attention
                 hidden_states,
-                position_ids,
-                past_key_value,
-                attention_mask,
-                output_attentions,
-                use_cache,
+                position_ids=position_ids,  # kwargs
+                past_key_value=past_key_value,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
                 attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+                batch_size=sp_batch_size,  # Qwen3Attention 有batch_szie这个参数，手动传递
+                **kwargs,  # 传递剩下的kwargs
             )
 
         if type(outputs) is tuple:
@@ -260,13 +267,35 @@ class Qwen3MoeDecoderLayer(nn.Layer):
 
         # Fully Connected
         residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        if isinstance(hidden_states, tuple):
-            hidden_states, router_logits = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)  # hha
+
+        # ==================== BEGIN FINAL FIX ====================
+        # 为了兼容序列并行 (输入是2D) 和非序列并行 (输入是3D)
+        is_2d_input = len(hidden_states.shape) == 2
+        if is_2d_input:
+            # 如果是序列并行产生的 2D 张量 [num_tokens, hidden_size]
+            # 我们需要将其 reshape 为 MoE 层期望的 3D 格式。
+            # 最安全的做法是将其视为 [num_tokens, 1, hidden_size]
+            # 这样 batch_size=num_tokens, seq_len=1
+            original_shape = hidden_states.shape
+            hidden_states = hidden_states.unsqueeze(1)
+
+        # 现在 hidden_states 保证是 3D 的
+
+        mlp_output = self.mlp(hidden_states)
+
+        if isinstance(mlp_output, tuple):
+            hidden_states, router_logits = mlp_output
         else:
+            hidden_states = mlp_output
             router_logits = None
 
+        if is_2d_input:
+            # 如果原始输入是 2D，我们需要在计算后恢复其形状
+            # [num_tokens, 1, hidden_size] -> [num_tokens, hidden_size]
+            hidden_states = hidden_states.reshape(original_shape)
+        # ===================== END FINAL FIX =====================
+
         hidden_states = residual + hidden_states
 
         outputs = (hidden_states,)
@@ -291,6 +320,33 @@ class Qwen3MoePretrainedModel(PretrainedModel):
     base_model_prefix = "model"
     _keys_to_ignore_on_load_unexpected = [r"self_attn.rotary_emb.inv_freq"]
 
+    @classmethod
+    def get_tensor_parallel_convert_actions(
+        cls, config, loaded_state_dict_keys, is_split=True, ignore_error=False, base_model_prefix=None
+    ):
+        """
+        Get the tensor parallel convert actions for the model.
+        This function is overridden to handle the case where MoE experts are grouped and should not be split across TP ranks.
+        """
+        # Get the default tensor parallel actions from the base class by calling super() with the exact same arguments.
+        tp_actions = super().get_tensor_parallel_convert_actions(
+            config,
+            loaded_state_dict_keys,
+            is_split=is_split,
+            ignore_error=ignore_error,
+            base_model_prefix=base_model_prefix,
+        )
+
+        # If moe_group is set, expert parameters should not be split.
+        # We remove them from the tp_actions dictionary.
+        if "Qwen3MoeForCausalLM" in config.architectures and config.moe_group == "tp":
+            # Iterate over a copy of the keys to safely modify the dictionary
+            for key in list(tp_actions.keys()):
+                if "mlp.experts" in key:
+                    del tp_actions[key]
+
+        return tp_actions
+
     @classmethod
     def _get_name_mappings(cls, config: Qwen3MoeConfig) -> list[StateDictNameMapping]:
         mappings: list[StateDictNameMapping] = []
@@ -358,11 +414,18 @@ class Qwen3MoePretrainedModel(PretrainedModel):
                 base_actions.pop("embed_tokens.weight")
 
             # Column Linear
-            base_actions["layers.0.self_attn.q_proj.weight"] = partial(fn, is_column=True)
-            # if we have enough num_key_value_heads to split, then split it.
-            if config.num_key_value_heads % config.tensor_parallel_degree == 0:
-                base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True)
-                base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True)
+            if config.fuse_attention_qkv:
+                base_actions["layers.0.self_attn.qkv_proj.weight"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.qkv_proj.bias"] = partial(fn, is_column=True)
+            else:
+                base_actions["layers.0.self_attn.q_proj.weight"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.q_proj.bias"] = partial(fn, is_column=True)
+                # if we have enough num_key_value_heads to split, then split it.
+                if config.num_key_value_heads % config.tensor_parallel_degree == 0:
+                    base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True)
+                    base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True)
+                    base_actions["layers.0.self_attn.k_proj.bias"] = partial(fn, is_column=True)
+                    base_actions["layers.0.self_attn.v_proj.bias"] = partial(fn, is_column=True)
 
             for key, action in base_actions.items():
                 if "layers.0." in key:
@@ -371,11 +434,20 @@ class Qwen3MoePretrainedModel(PretrainedModel):
                 final_actions[key] = action
 
             # Add tp split for expert params.
-            base_actions = {
-                "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
-                "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
-            }
+            if config.fuse_attention_ffn:
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_up_fused_proj.weight": partial(
+                        fn, is_column=True, is_naive_2fuse=True
+                    ),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
+            else:
+                # Add tp split for expert params.
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
             for key, action in base_actions.items():
                 for i in range(num_layers):
                     newkey = key.replace("layers.0.", f"layers.{i}.")
@@ -415,14 +487,15 @@ class Qwen3MoePretrainedModel(PretrainedModel):
         ]
 
         fuse_gate_up_keys = (
-            "layers.0.mlp.gate_proj.weight",
-            "layers.0.mlp.up_proj.weight",
-            "layers.0.mlp.gate_up_fused_proj.weight",
+            "layers.0.mlp.experts.0.gate_proj.weight",
+            "layers.0.mlp.experts.0.up_proj.weight",
+            "layers.0.mlp.experts.0.gate_up_fused_proj.weight",
         )
         num_heads = config.num_attention_heads
         num_key_value_heads = getattr(config, "num_key_value_heads", num_heads)
         fuse_attention_qkv = getattr(config, "fuse_attention_qkv", False)
         fuse_attention_ffn = getattr(config, "fuse_attention_ffn", False)
+        num_experts = getattr(config, "num_experts", 128)
 
         final_actions = {}
         if is_fuse:
@@ -435,20 +508,28 @@ class Qwen3MoePretrainedModel(PretrainedModel):
                         )
             if fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
-                    keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
-                    final_actions[keys] = fn
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = fn
         else:
             if not fuse_attention_qkv:
                 for i in range(config.num_hidden_layers):
                     for fuse_keys in fuse_qkv_keys:
                         keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_keys])
                         final_actions[keys] = partial(
-                            fn, split_nums=3, is_qkv=True, num_heads=num_heads, num_key_value_heads=num_key_value_heads
+                            fn,
+                            split_nums=3,
+                            is_qkv=True,
+                            num_heads=num_heads,
+                            num_key_value_heads=num_key_value_heads,
                         )
             if not fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
-                    keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
-                    final_actions[keys] = partial(fn, split_nums=2)
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = partial(fn, split_nums=2)
         return final_actions
 
     def _init_weights(self, layer):
@@ -595,15 +676,31 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
         past_key_value: Tensor,
         use_cache: bool,
         attn_mask_startend_row_indices=None,
+        batch_size=None,
+        **kwargs,
     ):
-        def create_custom_forward(module):
+        # 定义一个闭包，它会捕获所有需要的关键字参数
+        def create_custom_forward(module, **kwargs_to_bind):
             def custom_forward(*inputs):
-                return module(*inputs)
+                # 当 recompute 调用 custom_forward(*inputs) 时，
+                # 它会执行 module(*inputs, **kwargs_to_bind)
+                return module(*inputs, **kwargs_to_bind)
 
             return custom_forward
 
-        hidden_states = recompute(
-            create_custom_forward(layer_module),
+        # 准备好所有要绑定的关键字参数
+        kwargs_for_layer = {
+            "attn_mask_startend_row_indices": attn_mask_startend_row_indices,
+            "batch_size": batch_size,
+            **kwargs,
+        }
+
+        # 创建实例
+        wrapped_layer = create_custom_forward(layer_module, **kwargs_for_layer)
+
+        # 调用 recompute，和方案一完全一样
+        layer_outputs = recompute(
+            wrapped_layer,
             hidden_states,
             position_ids,
             attention_mask,
@@ -611,11 +708,10 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
             output_router_logits,
             past_key_value,
             use_cache,
-            attn_mask_startend_row_indices,
             use_reentrant=self.config.recompute_use_reentrant,
         )
 
-        return hidden_states
+        return layer_outputs
 
     def forward(
         self,
@@ -632,6 +728,8 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
         attn_mask_startend_row_indices=None,
         **kwargs,
     ) -> Union[Tuple, MoEModelOutputWithPast]:
+        # batch_size = kwargs.pop("batch_size", None)
+
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
 
         output_router_logits = (
@@ -648,6 +746,7 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
             raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
         elif input_ids is not None:
             batch_size, seq_length = input_ids.shape
+
         elif inputs_embeds is not None:
             batch_size, seq_length, _ = inputs_embeds.shape
         else:
@@ -679,11 +778,13 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
             attention_mask = None
         else:
             # [bs, seq_len]
+
             attention_mask = (
                 paddle.ones((batch_size, seq_length_with_past), dtype=paddle.bool)
                 if attention_mask is None
                 else attention_mask
             )
+
             attention_mask = self._prepare_decoder_attention_mask(
                 attention_mask, (batch_size, seq_length), cache_length, inputs_embeds.dtype
             )  # [bs, 1, seq_len, seq_len]
@@ -723,9 +824,11 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
                     past_key_value,
                     use_cache,
                     attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+                    batch_size=batch_size,
+                    **kwargs,  # 传递剩下的kwargs
                 )
             else:
-                layer_outputs = decoder_layer(
+                layer_outputs = decoder_layer(  # here
                     hidden_states,
                     position_ids,
                     attention_mask,
@@ -734,6 +837,8 @@ class Qwen3MoeModel(Qwen3MoePretrainedModel):
                     past_key_value,
                     use_cache,
                     attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+                    batch_size=batch_size,  # here add for qwen3moe 传递 batch_size
+                    **kwargs,  # 传递剩下的kwargs
                 )
 
             # NOTE: clear outdate cache after it has been used for memory saving
@@ -792,7 +897,7 @@ class Qwen3MoeForCausalLM(Qwen3MoePretrainedModel):
         self.config = config
 
         self.model = Qwen3MoeModel(config)
-        self.lm_head = Qwen3MoeLMHead(config)
+        self.lm_head = Qwen3MoeLMHead(config)  # Qwen2LMHead
         self.criterion = Qwen3MoePretrainingCriterion(config)
         self.router_aux_loss_coef = config.router_aux_loss_coef
         self.num_experts = config.num_experts
@@ -905,6 +1010,7 @@ class Qwen3MoeForCausalLM(Qwen3MoePretrainedModel):
         output_router_logits: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         attn_mask_startend_row_indices=None,
+        **kwargs,  # here add for qwen3moe
     ):
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -922,9 +1028,8 @@ class Qwen3MoeForCausalLM(Qwen3MoePretrainedModel):
             )
             attention_mask = None
 
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
-            input_ids=input_ids,  # [bs, seq_len]
+            input_ids=input_ids,
             position_ids=position_ids,
             attention_mask=attention_mask,
             inputs_embeds=inputs_embeds,
@@ -935,6 +1040,7 @@ class Qwen3MoeForCausalLM(Qwen3MoePretrainedModel):
             output_router_logits=output_router_logits,
             return_dict=return_dict,
             attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+            **kwargs,  # 传递剩下的kwargs(如有)
         )
 
         hidden_states = outputs[0]  # [bs, seq_len, dim]
@@ -961,8 +1067,11 @@ class Qwen3MoeForCausalLM(Qwen3MoePretrainedModel):
             else:
                 loss = paddle.sum(masked_lm_loss * binary_sequence) / count
             logits = None
+
         else:
-            logits = self.lm_head(hidden_states, tensor_parallel_output=tensor_parallel_output)
+            logits = self.lm_head(
+                hidden_states, tensor_parallel_output=tensor_parallel_output, batch_size=input_ids.shape[0]
+            )
 
             loss = None
             if labels is not None:

paddlenlp/transformers/utils.py

@@ -56,7 +56,12 @@ from paddlenlp.utils.import_utils import import_module
 from paddlenlp.utils.log import logger
 
 from ..utils.download import resolve_file_path
-from .aistudio_utils import aistudio_download
+
+# TODO(@zewu): upgrade aistudio to the newest version
+try:
+    from .aistudio_utils import aistudio_download
+except:
+    aistudio_download = None
 
 HUGGINGFACE_CO_RESOLVE_ENDPOINT = "https://huggingface.co"
 

paddlenlp/utils/paddle_patch.py

@@ -12,11 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 import ml_dtypes
 import numpy as np
 import paddle
 import paddle.nn as nn
 
+enable_paddlenlp_monkey_patch = int(os.getenv("FLAGS_enable_paddlenlp_monkey_patch", "1"))
+
 np.bfloat16 = ml_dtypes.bfloat16
 np.float8_e5m2 = ml_dtypes.float8_e5m2
 np.float8_e4m3fn = ml_dtypes.float8_e4m3fn
@@ -141,8 +145,9 @@ def _numel(self, *args, **kwargs):
     return origin_numel(self, *args, **kwargs)
 
 
-paddle.Tensor.numpy = _numpy
-paddle.Tensor.__call__ = enhance_init
-paddle.to_tensor = enhance_to_tensor
-paddle.core.eager.Tensor.set_value = enhance_set_value
-paddle.Tensor.numel = _numel
+if enable_paddlenlp_monkey_patch:
+    paddle.Tensor.numpy = _numpy
+    paddle.Tensor.__call__ = enhance_init
+    paddle.to_tensor = enhance_to_tensor
+    paddle.core.eager.Tensor.set_value = enhance_set_value
+    paddle.Tensor.numel = _numel