feat: FSDP2 online dynamic weight conversion support-MindSpeed-LLM-AtomGit

已合并

feat: FSDP2 online dynamic weight conversion support #4597

iansheng创建于 6月15日从iansheng/MindSpeed-LLM： dyn_conv 合入到 Ascend/MindSpeed-LLM： master

feat: FSDP2 online dynamic weight conversion support #4597

已合并

iansheng创建于 6月15日

从iansheng/MindSpeed-LLM dyn_conv合入到Ascend/MindSpeed-LLM master

共 13 个文件变更+652-201

Mdocs/zh/pytorch/features/fsdp2/quantization.md+0-1

	@@ -99,7 +99,6 @@ QUANT_ARGS="
99	--parallel.efsdp_shard_placement_fn shard_by_dim_0	99	--parallel.efsdp_shard_placement_fn shard_by_dim_0
100	"	100	"
101		101
102	-bash tests/tools/fsdp2/moe_hf_param_merge_experts.sh
103	torchrun $DISTRIBUTED_ARGS train_fsdp2.py \	102	torchrun $DISTRIBUTED_ARGS train_fsdp2.py \
104	examples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A3.yaml \	103	examples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A3.yaml \
105	$QUANT_ARGS \	104	$QUANT_ARGS \

Mexamples/fsdp2/minimax_m27/chat_minimax_m2p7_fsdp2_A3.sh+0-1

	@@ -17,7 +17,6 @@ DISTRIBUTED_ARGS="
17	"	17	"
18		18
19	mkdir -p ./logs	19	mkdir -p ./logs
20	-bash tests/tools/fsdp2/moe_hf_param_merge_experts.sh
21	torchrun $DISTRIBUTED_ARGS inference_fsdp2.py examples/fsdp2/minimax_m27/pretrain_minimax_m2p7_229b_8K_fsdp2_A3.yaml \	20	torchrun $DISTRIBUTED_ARGS inference_fsdp2.py examples/fsdp2/minimax_m27/pretrain_minimax_m2p7_229b_8K_fsdp2_A3.yaml \
22	--model.model_name_or_path /home/data/MiniMax-M2.7/ \	21	--model.model_name_or_path /home/data/MiniMax-M2.7/ \
23	--parallel.fsdp_size 32 \	22	--parallel.fsdp_size 32 \

Mexamples/fsdp2/minimax_m27/pretrain_minimax_m2p7_229b_8K_fsdp2_A3.sh+0-1

	@@ -16,5 +16,4 @@ DISTRIBUTED_ARGS="
16	--master_port $MASTER_PORT	16	--master_port $MASTER_PORT
17	"	17	"
18	mkdir -p ./logs	18	mkdir -p ./logs
19	-bash tests/tools/fsdp2/moe_hf_param_merge_experts.sh
20	torchrun $DISTRIBUTED_ARGS train_fsdp2.py examples/fsdp2/minimax_m27/pretrain_minimax_m2p7_229b_8K_fsdp2_A3.yaml \| tee logs/pretrain_minimax_m2p7_229b_8K_fsdp2_A3_${TIMESTAMP}.log	19	torchrun $DISTRIBUTED_ARGS train_fsdp2.py examples/fsdp2/minimax_m27/pretrain_minimax_m2p7_229b_8K_fsdp2_A3.yaml \| tee logs/pretrain_minimax_m2p7_229b_8K_fsdp2_A3_${TIMESTAMP}.log

Mexamples/fsdp2/qwen3_moe/pretrain_qwen3_235b_4k_fsdp2_A3.sh+1-2

	@@ -22,5 +22,4 @@ DISTRIBUTED_ARGS="
22	--master_port $MASTER_PORT	22	--master_port $MASTER_PORT
23	"	23	"
24		24
25	-bash tests/tools/fsdp2/moe_hf_param_merge_experts.sh	25	+torchrun $DISTRIBUTED_ARGS train_fsdp2.py examples/fsdp2/qwen3_moe/pretrain_qwen3_235b_4k_fsdp2_A3.yaml \| tee logs/train_fsdp2_qwen3_235b_A3.log
26	-torchrun $DISTRIBUTED_ARGS train_fsdp2.py examples/fsdp2/qwen3_moe/pretrain_qwen3_235b_4k_fsdp2_A3.yaml \| tee logs/train_fsdp2_qwen3_235b_A3.log

Mexamples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A3.sh+1-2

	@@ -18,7 +18,6 @@ DISTRIBUTED_ARGS="
18	"	18	"
19		19
20	mkdir -p ./logs	20	mkdir -p ./logs
21	-bash tests/tools/fsdp2/moe_hf_param_merge_experts.sh
22	torchrun $DISTRIBUTED_ARGS train_fsdp2.py \	21	torchrun $DISTRIBUTED_ARGS train_fsdp2.py \
23	examples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A3.yaml \	22	examples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A3.yaml \
24	- \| tee logs/pretrain_qwen3_moe_30b_a3b_4K_fsdp2_${TIMESTAMP}.log	23	+ \| tee logs/pretrain_qwen3_moe_30b_a3b_4K_fsdp2_${TIMESTAMP}.log

Mexamples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_quant_A5.sh+1-2

	@@ -24,8 +24,7 @@ QUANT_ARGS="
24	--parallel.ep_dispatcher eager \	24	--parallel.ep_dispatcher eager \
25	"	25	"
26	mkdir -p ./logs	26	mkdir -p ./logs
27	-bash tests/tools/fsdp2/moe_hf_param_merge_experts.sh
28	torchrun $DISTRIBUTED_ARGS train_fsdp2.py \	27	torchrun $DISTRIBUTED_ARGS train_fsdp2.py \
29	examples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A5.yaml \	28	examples/fsdp2/qwen3_moe/pretrain_qwen3_30b_4k_fsdp2_A5.yaml \
30	$QUANT_ARGS\	29	$QUANT_ARGS\
31	- \| tee logs/pretrain_qwen3_moe_30b_a3b_4K_fsdp2_${TIMESTAMP}.log	30	+ \| tee logs/pretrain_qwen3_moe_30b_a3b_4K_fsdp2_${TIMESTAMP}.log

Mmindspeed_llm/fsdp2/checkpoint/__init__.py+27-1

	@@ -11,4 +11,30 @@
11	# distributed under the License is distributed on an "AS IS" BASIS,	11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.	12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and	13	# See the License for the specific language governing permissions and
14	-# limitations under the License.	14	+# limitations under the License.
		15	+
		16	+"""
		17	+Checkpoint utilities for FSDP2.
		18	+
		19	+This module provides weight conversion utilities that leverage the
		20	+transformers conversion_mapping API for transforming HuggingFace
		21	+checkpoint weights into MindSpeed-LLM format.
		22	+"""
		23	+
		24	+from mindspeed_llm.fsdp2.checkpoint.weight_conv_adapter import WeightConvAdapter
		25	+from mindspeed_llm.fsdp2.checkpoint.conversion_mappings import (
		26	+ FlattenExperts,
		27	+ UnflattenExperts,
		28	+ MappingRule,
		29	+ MappingRuleRegistry,
		30	+ apply_custom_mappings,
		31	+)
		32	+
		33	+__all__ = [
		34	+ 'WeightConvAdapter',
		35	+ 'FlattenExperts',
		36	+ 'UnflattenExperts',
		37	+ 'MappingRule',
		38	+ 'MappingRuleRegistry',
		39	+ 'apply_custom_mappings',
		40	+]

Amindspeed_llm/fsdp2/checkpoint/conversion_mappings.py+253-0

@@ -0,0 +1,253 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""
	3	+Custom conversion mappings for MindSpeed-LLM.
	4	+
	5	+Defines model-specific conversion rules that extend transformers' native
	6	+conversion_mapping. New rules can be added by subclassing ``MappingRule``
	7	+and registering via ``MappingRuleRegistry.register()``.
	8	+"""
	9	+
	10	+from abc import ABC, abstractmethod
	11	+from importlib.metadata import version
	12	+from typing import Optional, List
	13	+
	14	+import torch
	15	+
	16	+from mindspeed_llm.fsdp2.utils.logging import get_logger
	17	+from mindspeed_llm.fsdp2.utils.global_vars import get_args
	18	+
	19	+logger = get_logger(__name__)
	20	+
	21	+if version("transformers") >= "5.0.0":
	22	+ from transformers.core_model_loading import ConversionOps
	23	+else:
	24	+
	25	+ class ConversionOps:
	26	+ pass
	27	+
	28	+
	29	+# ==============================================================================
	30	+# Conversion Ops
	31	+# ==============================================================================
	32	+
	33	+
	34	+class FlattenExperts(ConversionOps):
	35	+ """
	36	+ Flatten a 3-D stacked-expert tensor into the 2-D merged layout expected
	37	+ by MindSpeed-LLM's fused expert modules.
	38	+
	39	+ Forward: ``(N, d1, d2) → reshape → (N*d1, d2)``
	40	+ Reverse: ``(Nd1, d2) → reshape → (N, d1, d2)`` (requires num_experts*)
	41	+
	42	+ Supports EP (Expert Parallel) sharding: works with both global and local expert counts.
	43	+ """
	44	+
	45	+ def __init__(self, num_experts: Optional[int] = None):
	46	+ self.num_experts = num_experts
	47	+
	48	+ @torch.no_grad()
	49	+ def convert(self, input_dict, source_patterns, target_patterns, **kwargs):
	50	+ target_pattern = self._get_target_pattern(input_dict, source_patterns, target_patterns)
	51	+ tensors = next(iter(input_dict.values()))
	52	+ t = tensors[0] if isinstance(tensors, list) else tensors
	53	+ if t.ndim == 3:
	54	+ t = t.contiguous().view(-1, t.shape[2])
	55	+ return {target_pattern: t}
	56	+
	57	+ @staticmethod
	58	+ def _get_target_pattern(input_dict, source_patterns, target_patterns):
	59	+ if len(target_patterns) > 1:
	60	+ return next(iter(input_dict.keys()))
	61	+ return target_patterns[0]
	62	+
	63	+ @property
	64	+ def reverse_op(self):
	65	+ return UnflattenExperts(num_experts=self.num_experts)
	66	+
	67	+ def __repr__(self):
	68	+ return "FlattenExperts()"
	69	+
	70	+
	71	+class UnflattenExperts(ConversionOps):
	72	+ """
	73	+ Reverse of :class:`FlattenExperts`. Reshapes a 2-D merged tensor back to 3-D.
	74	+
	75	+ ``(N*d1, d2) → reshape → (N, d1, d2)``
	76	+
	77	+ Supports EP (Expert Parallel) sharding: automatically detects num_local_experts
	78	+ when EP is enabled, falling back to num_experts for non-EP cases.
	79	+ """
	80	+
	81	+ def __init__(self, num_experts: Optional[int] = None):
	82	+ self.num_experts = num_experts
	83	+
	84	+ @torch.no_grad()
	85	+ def convert(self, input_dict, source_patterns, target_patterns, **kwargs):
	86	+ n = self.num_experts
	87	+ if n is None:
	88	+ config = kwargs.get("config")
	89	+ if config is None:
	90	+ raise ValueError("UnflattenExperts requires num_experts or config")
	91	+ n = getattr(config, "num_local_experts", None) or getattr(config, "num_experts", None)
	92	+ if n is None:
	93	+ raise ValueError(f"Cannot infer num_experts from config: {config}")
	94	+
	95	+ target_pattern = self._get_target_pattern(input_dict, source_patterns, target_patterns)
	96	+ tensors = next(iter(input_dict.values()))
	97	+ t = tensors[0] if isinstance(tensors, list) else tensors
	98	+ if t.ndim == 2:
	99	+ d2 = t.shape[1]
	100	+ t = t.contiguous().view(n, -1, d2)
	101	+ return {target_pattern: t}
	102	+
	103	+ @staticmethod
	104	+ def _get_target_pattern(input_dict, source_patterns, target_patterns):
	105	+ if len(target_patterns) > 1:
	106	+ return next(iter(input_dict.keys()))
	107	+ return target_patterns[0]
	108	+
	109	+ @property
	110	+ def reverse_op(self):
	111	+ return FlattenExperts(num_experts=self.num_experts)
	112	+
	113	+ def __repr__(self):
	114	+ return f"UnflattenExperts(num_experts={self.num_experts})"
	115	+
	116	+
	117	+# ==============================================================================
	118	+# Mapping Rule Registry
	119	+# ==============================================================================
	120	+
	121	+
	122	+class MappingRule(ABC):
	123	+ """
	124	+ Base class for custom conversion mapping rules.
	125	+
	126	+ Subclass and implement ``condition()`` and ``apply()`` to define new rules.
	127	+ Register via ``MappingRuleRegistry.register(YourRule())``.
	128	+ """
	129	+
	130	+ name: str = ""
	131	+ model_types: Optional[List[str]] = None
	132	+
	133	+ @abstractmethod
	134	+ def condition(self, model_type: str) -> bool:
	135	+ """Return True if this rule should be applied for the given model_type."""
	136	+
	137	+ @abstractmethod
	138	+ def apply(self, conversions: list, model_type: str) -> None:
	139	+ """Mutate conversions in-place to inject custom operations."""
	140	+
	141	+
	142	+class MappingRuleRegistry:
	143	+ """
	144	+ Global registry for ``MappingRule`` instances.
	145	+
	146	+ Rules are evaluated in registration order. Each rule whose ``condition()``
	147	+ returns True will have its ``apply()`` called on the conversions list.
	148	+ """
	149	+
	150	+ _rules: List[MappingRule] = []
	151	+
	152	+ @classmethod
	153	+ def register(cls, rule: MappingRule) -> None:
	154	+ if not isinstance(rule, MappingRule):
	155	+ raise TypeError(f"Expected MappingRule, got {type(rule)}")
	156	+ cls._rules.append(rule)
	157	+
	158	+ @classmethod
	159	+ def unregister(cls, name: str) -> Optional[MappingRule]:
	160	+ for i, rule in enumerate(cls._rules):
	161	+ if rule.name == name:
	162	+ return cls._rules.pop(i)
	163	+ return None
	164	+
	165	+ @classmethod
	166	+ def apply_all(cls, conversions: list, model_type: str) -> None:
	167	+ for rule in cls._rules:
	168	+ if rule.model_types is not None and model_type not in rule.model_types:
	169	+ continue
	170	+ if rule.condition(model_type):
	171	+ rule.apply(conversions, model_type)
	172	+ logger.info_rank0(f"> Applied mapping rule '{rule.name}' for model_type={model_type}")
	173	+
	174	+ @classmethod
	175	+ def list_rules(cls) -> List[str]:
	176	+ return [r.name for r in cls._rules]
	177	+
	178	+ @classmethod
	179	+ def clear(cls) -> None:
	180	+ cls._rules.clear()
	181	+
	182	+
	183	+# ==============================================================================
	184	+# Built-in Rules
	185	+# ==============================================================================
	186	+
	187	+_MODELS_WITH_2D_MERGED_EXPERTS = [
	188	+ "qwen2_moe",
	189	+ "qwen3_moe",
	190	+ "qwen3_next",
	191	+ "qwen3_5_moe",
	192	+ "qwen3_omni_moe",
	193	+ "qwen3_omni_moe_thinker",
	194	+ "minimax",
	195	+ "minimax_m2",
	196	+]
	197	+
	198	+
	199	+class GroupedGemmFlattenRule(MappingRule):
	200	+ """
	201	+ When ``moe_grouped_gemm`` is enabled, append flatten operations to converters
	202	+ containing ``MergeModulelist`` so that expert weights match the offline
	203	+ conversion script layout (moe_hf_param_merge_experts.py).
	204	+
	205	+ Both gate_up_proj and down_proj need Transpose(1,2) before FlattenExperts:
	206	+
	207	+ gate_up_proj (multi-source: gate+up):
	208	+ MergeModulelist → Concatenate → Transpose(1,2) → FlattenExperts
	209	+ [E,I,H] → [E,2I,H] → [E,H,2I] → [E*H, 2I]
	210	+
	211	+ down_proj (single-source):
	212	+ MergeModulelist → Transpose(1,2) → FlattenExperts
	213	+ [E,I,H] → [E,H,I] → [E*H, I]
	214	+ """
	215	+
	216	+ name = "grouped_gemm_flatten"
	217	+ model_types = _MODELS_WITH_2D_MERGED_EXPERTS
	218	+
	219	+ def condition(self, model_type: str) -> bool:
	220	+ return getattr(get_args(), "moe_grouped_gemm", False)
	221	+
	222	+ def apply(self, conversions: list, model_type: str) -> None:
	223	+ from transformers.core_model_loading import MergeModulelist, Transpose
	224	+
	225	+ for entry in conversions:
	226	+ if not hasattr(entry, "operations"):
	227	+ continue
	228	+ if not any(isinstance(op, MergeModulelist) for op in entry.operations):
	229	+ continue
	230	+ entry.operations.extend([Transpose(1, 2), FlattenExperts()])
	231	+
	232	+
	233	+# ==============================================================================
	234	+# Register built-in rules
	235	+# ==============================================================================
	236	+
	237	+MappingRuleRegistry.register(GroupedGemmFlattenRule())
	238	+
	239	+
	240	+# ==============================================================================
	241	+# Public API
	242	+# ==============================================================================
	243	+
	244	+
	245	+def apply_custom_mappings(conversions: list, model_type: str) -> None:
	246	+ """
	247	+ Apply all registered MindSpeed-LLM custom conversion rules.
	248	+
	249	+ Args:
	250	+ conversions: List of WeightConverter/WeightRenaming from transformers.
	251	+ model_type: HuggingFace model_type identifier.
	252	+ """
	253	+ MappingRuleRegistry.apply_all(conversions, model_type)

Amindspeed_llm/fsdp2/checkpoint/weight_conv_adapter.py+118-0

@@ -0,0 +1,118 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""
	3	+Weight conversion execution engine for MindSpeed-LLM.
	4	+
	5	+Provides a thin wrapper over transformers' conversion_mapping API to execute
	6	+weight conversions during checkpoint loading. Mapping rules are defined
	7	+separately in conversion_mappings.py.
	8	+"""
	9	+
	10	+import re
	11	+from copy import deepcopy
	12	+from importlib.metadata import version
	13	+from typing import Optional, Tuple
	14	+
	15	+from mindspeed_llm.fsdp2.utils.logging import get_logger
	16	+
	17	+logger = get_logger(__name__)
	18	+
	19	+
	20	+class WeightConvAdapter:
	21	+ """
	22	+ Thin wrapper over transformers.conversion_mapping.
	23	+
	24	+ Exposes native WeightRenaming / WeightConverter objects and delegates
	25	+ key-matching / tensor conversion to transformers' own implementations.
	26	+ """
	27	+
	28	+ def __init__(self, model_type: Optional[str] = None):
	29	+ self.renamings: list = []
	30	+ self.converters: list = []
	31	+
	32	+ if not model_type:
	33	+ return
	34	+
	35	+ if version("transformers") < "5.0.0":
	36	+ logger.info_rank0(
	37	+ f"> Online weight conversion requires transformers >= 5.0.0, "
	38	+ f"current version: {version('transformers')}. Skipping."
	39	+ )
	40	+ return
	41	+
	42	+ from transformers.conversion_mapping import get_checkpoint_conversion_mapping
	43	+ from transformers.core_model_loading import WeightConverter, WeightRenaming
	44	+ from mindspeed_llm.fsdp2.checkpoint.conversion_mappings import apply_custom_mappings
	45	+
	46	+ conversions = get_checkpoint_conversion_mapping(model_type)
	47	+ if not conversions:
	48	+ return
	49	+
	50	+ apply_custom_mappings(conversions, model_type)
	51	+
	52	+ for entry in conversions:
	53	+ if isinstance(entry, WeightRenaming):
	54	+ self.renamings.append(entry)
	55	+ elif isinstance(entry, WeightConverter):
	56	+ self.converters.append(entry)
	57	+
	58	+ if conversions:
	59	+ logger.info_rank0(
	60	+ f"> Weight conversion mapping for model_type={model_type}: "
	61	+ f"{len(self.renamings)} renamings, {len(self.converters)} converters"
	62	+ )
	63	+
	64	+ @property
	65	+ def has_conversions(self) -> bool:
	66	+ return bool(self.renamings) or bool(self.converters)
	67	+
	68	+ def rename_key(self, key: str) -> Tuple[str, Optional[str]]:
	69	+ """
	70	+ Rename checkpoint key via transformers rename_source_key.
	71	+
	72	+ Returns:
	73	+ (renamed_key, source_pattern_or_None)
	74	+ """
	75	+ from transformers.core_model_loading import rename_source_key
	76	+
	77	+ return rename_source_key(key, self.renamings, self.converters)
	78	+
	79	+ def match_converter(self, source_pattern: str):
	80	+ """Find the converter template that owns source_pattern."""
	81	+ for c in self.converters:
	82	+ if source_pattern in c.source_patterns:
	83	+ return c
	84	+ return None
	85	+
	86	+ @staticmethod
	87	+ def dispatch_converted(converter, target_name: str, collected: dict, original_keys: list = None):
	88	+ """
	89	+ Run the native WeightConverter.convert() pipeline on raw tensors.
	90	+
	91	+ Args:
	92	+ converter: Converter template from match_converter().
	93	+ target_name: Full model parameter name (e.g. model.layers.0.mlp.experts.gate_up_proj).
	94	+ collected: {source_pattern: [tensor, ...]} grouped for one layer.
	95	+ original_keys: Original checkpoint keys for the first source pattern's
	96	+ tensors, used to sort by expert index when weights span multiple
	97	+ safetensors files.
	98	+
	99	+ Yields:
	100	+ (full_name, tensor) pairs ready for dispatch.
	101	+ """
	102	+ if original_keys and len(original_keys) > 1:
	103	+ m = re.search(r'\.experts\.(\d+)\.', original_keys[0])
	104	+ if m:
	105	+ indexed = []
	106	+ for sp, tensors in collected.items():
	107	+ pairs = list(zip(original_keys, tensors))
	108	+ pairs.sort(key=lambda p: int(re.search(r'\.experts\.(\d+)\.', p[0]).group(1)))
	109	+ indexed.append((sp, [t for _, t in pairs]))
	110	+ collected = dict(indexed)
	111	+
	112	+ fresh = deepcopy(converter)
	113	+ fresh.collected_tensors = collected
	114	+ result = fresh.convert(target_name)
	115	+ for name, tensor in result.items():
	116	+ if isinstance(tensor, list):
	117	+ tensor = tensor[0]
	118	+ yield name, tensor

Mmindspeed_llm/fsdp2/distributed/mindspeed_parallel_engine.py+23-16

	@@ -1,9 +1,8 @@
1	# Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved.	1	# Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved.
2	-from typing import Optional	2	+from typing import Optional, Any
3	import torch	3	import torch
4		4
5	-from mindspeed.fsdp.distributed.fully_shard_parallel.fully_shard_parallel import \	5	+from mindspeed.fsdp.distributed.fully_shard_parallel.fully_shard_parallel import fully_shard_parallel_modules
6	- fully_shard_parallel_modules
7	from mindspeed.fsdp.distributed.tensor_parallel.tensor_parallel import tensor_parallel_modules	6	from mindspeed.fsdp.distributed.tensor_parallel.tensor_parallel import tensor_parallel_modules
8	from mindspeed.fsdp.memory.recompute.recompute import recompute_modules	7	from mindspeed.fsdp.memory.recompute.recompute import recompute_modules
9	from mindspeed_llm.fsdp2.distributed.parallel_state import init_parallel_state	8	from mindspeed_llm.fsdp2.distributed.parallel_state import init_parallel_state
	@@ -18,12 +17,20 @@ logger = get_logger(__name__)
18		17
19		18
20	class MindSpeedParallelEngine(torch.nn.Module):	19	class MindSpeedParallelEngine(torch.nn.Module):
21	- def __init__(self, config: ParallelEngineConfig, model: torch.nn.Module, init_device: str = "cpu", weights_path: Optional[str] = None):	20	+ def __init__(
22	- super(MindSpeedParallelEngine, self).__init__()	21	+ self,
		22	+ config: ParallelEngineConfig,
		23	+ model: torch.nn.Module,
		24	+ init_device: str = "cpu",
		25	+ weights_path: Optional[str] = None,
		26	+ hf_config: Optional[Any] = None,
		27	+ ):
		28	+ super().__init__()
23	self.config = config	29	self.config = config
24	self.model = model	30	self.model = model
25	self.init_device = init_device	31	self.init_device = init_device
26	self.weights_path = weights_path	32	self.weights_path = weights_path
		33	+ self.hf_config = hf_config
27		34
28	self.parallel_state = init_parallel_state(self.config)	35	self.parallel_state = init_parallel_state(self.config)
29	self.apply_quantization_modules()	36	self.apply_quantization_modules()
	@@ -36,14 +43,12 @@ class MindSpeedParallelEngine(torch.nn.Module):
36	# For meta device: load weights after fsdp wrapping	43	# For meta device: load weights after fsdp wrapping
37	if self.init_device == "meta":	44	if self.init_device == "meta":
38	logger.info_rank0("> Loading weights after FSDP wrapping...")	45	logger.info_rank0("> Loading weights after FSDP wrapping...")
39	- WeightLoader.load(	46	+ WeightLoader.load(model=self.model, weights_path=self.weights_path, device=None, hf_config=self.hf_config)
40	- model=self.model,
41	- weights_path=self.weights_path,
42	- device=None # Auto-detect device
43	- )
44		47
45	def apply_fsdp_modules(self):	48	def apply_fsdp_modules(self):
46	- self.model = fully_shard_parallel_modules(self.model, self.parallel_state.get_fsdp_device_mesh(), self.config.fsdp_plan)	49	+ self.model = fully_shard_parallel_modules(
		50	+ self.model, self.parallel_state.get_fsdp_device_mesh(), self.config.fsdp_plan
		51	+ )
47		52
48	def apply_tp_modules(self):	53	def apply_tp_modules(self):
49	if self.config.tensor_parallel_size == 1:	54	if self.config.tensor_parallel_size == 1:
	@@ -52,12 +57,15 @@ class MindSpeedParallelEngine(torch.nn.Module):
52		57
53	def apply_ep_modules(self):	58	def apply_ep_modules(self):
54	if self.config.expert_parallel_size > 1:	59	if self.config.expert_parallel_size > 1:
55	- self.model = expert_parallelize_modules(self.model, self.parallel_state.get_ep_device_mesh(), self.config.ep_plan)	60	+ self.model = expert_parallelize_modules(
		61	+ self.model, self.parallel_state.get_ep_device_mesh(), self.config.ep_plan
		62	+ )
56	if self.config.expert_fully_shard_parallel_size > 1:	63	if self.config.expert_fully_shard_parallel_size > 1:
57	- self.model = expert_fully_shard_modules(self.model, self.parallel_state.get_efsdp_device_mesh(), self.config.ep_plan)	64	+ self.model = expert_fully_shard_modules(
		65	+ self.model, self.parallel_state.get_efsdp_device_mesh(), self.config.ep_plan
		66	+ )
58		67
59	def apply_cp_modules(self):	68	def apply_cp_modules(self):
60	-
61	VALID_CP_TYPES = ("ulysses", "ring")	69	VALID_CP_TYPES = ("ulysses", "ring")
62	cp_size = self.config.context_parallel_size	70	cp_size = self.config.context_parallel_size
63	cp_type = self.config.context_parallel_type	71	cp_type = self.config.context_parallel_type
	@@ -69,7 +77,6 @@ class MindSpeedParallelEngine(torch.nn.Module):
69	raise ValueError(f"num_attention_heads must be divisible by context_parallel_size (current: {cp_size})")	77	raise ValueError(f"num_attention_heads must be divisible by context_parallel_size (current: {cp_size})")
70	apply_context_parallelize_modules(self.model, self.config.cp_plan)	78	apply_context_parallelize_modules(self.model, self.config.cp_plan)
71		79
72	-
73	def apply_recompute_modules(self):	80	def apply_recompute_modules(self):
74	if not self.config.recompute:	81	if not self.config.recompute:
75	return	82	return
	@@ -88,7 +95,7 @@ class MindSpeedParallelEngine(torch.nn.Module):
88	model_converters = build_model_converter(self.config.quantization_plan)	95	model_converters = build_model_converter(self.config.quantization_plan)
89	model_converters.convert(self.model)	96	model_converters.convert(self.model)
90	except Exception as e:	97	except Exception as e:
91	- raise RuntimeError(f"Failed to convert quantization plan") from e	98	+ raise RuntimeError("Failed to convert quantization plan") from e
92		99
93	def apply_optimizer_hook(self, optimizer: torch.optim.Optimizer):	100	def apply_optimizer_hook(self, optimizer: torch.optim.Optimizer):
94	if not self.config.quantization_plan.recipe_name:	101	if not self.config.quantization_plan.recipe_name:

Mmindspeed_llm/fsdp2/models/model_factory.py+27-40

	@@ -1,8 +1,7 @@
1	import os	1	import os
2	import torch	2	import torch
3	import torch.distributed as dist	3	import torch.distributed as dist
4	-from typing import Any, Type	4	+from transformers import AutoConfig
5	-from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig
6		5
7	from mindspeed_llm.fsdp2.models.model_registry import ModelRegistry	6	from mindspeed_llm.fsdp2.models.model_registry import ModelRegistry
8	from mindspeed_llm.fsdp2.distributed.mindspeed_parallel_engine import MindSpeedParallelEngine	7	from mindspeed_llm.fsdp2.distributed.mindspeed_parallel_engine import MindSpeedParallelEngine
	@@ -12,7 +11,7 @@ from mindspeed_llm.fsdp2.distributed.parallel_engine_config import (
12	TPPlanConfig,	11	TPPlanConfig,
13	EPPlanConfig,	12	EPPlanConfig,
14	CPPlanConfig,	13	CPPlanConfig,
15	- QuantizeConfig	14	+ QuantizeConfig,
16	)	15	)
17		16
18	from mindspeed_llm.fsdp2.utils.logging import get_logger	17	from mindspeed_llm.fsdp2.utils.logging import get_logger
	@@ -26,9 +25,9 @@ logger = get_logger(__name__)
26	# ==============================================================================	25	# ==============================================================================
27	class ModelFactory:	26	class ModelFactory:
28	"""	27	"""
29	- Responsible for building HuggingFace native models and wrapping them	28	+ Responsible for building HuggingFace native models and wrapping them
30	as MindSpeed FSDP instances based on parallelization arguments.	29	as MindSpeed FSDP instances based on parallelization arguments.
31	-	30	+
32	Supports two initialization modes controlled by model_args.init_model_with_meta_device:	31	Supports two initialization modes controlled by model_args.init_model_with_meta_device:
33	- False: Load model fully on CPU (original behavior)	32	- False: Load model fully on CPU (original behavior)
34	- True: Create empty model on meta device, load weights after FSDP wrapping	33	- True: Create empty model on meta device, load weights after FSDP wrapping
	@@ -38,9 +37,9 @@ class ModelFactory:
38	def create(model_args, parallel_args):	37	def create(model_args, parallel_args):
39	"""	38	"""
40	Creates a MindSpeed FSDP wrapped model.	39	Creates a MindSpeed FSDP wrapped model.
41	-	40	+
42	Args:	41	Args:
43	- model_args: Contains model_name_or_path, trust_remote_code, train_from_scratch,	42	+ model_args: Contains model_name_or_path, trust_remote_code, train_from_scratch,
44	init_model_with_meta_device, etc.	43	init_model_with_meta_device, etc.
45	parallel_args: Contains tp_size, fsdp_size, recompute, ep_size, etc.	44	parallel_args: Contains tp_size, fsdp_size, recompute, ep_size, etc.
46	"""	45	"""
	@@ -56,15 +55,14 @@ class ModelFactory:
56	# 2. Determine initialization device based on init_model_with_meta_device flag	55	# 2. Determine initialization device based on init_model_with_meta_device flag
57	use_meta_device = getattr(model_args, 'init_model_with_meta_device', False)	56	use_meta_device = getattr(model_args, 'init_model_with_meta_device', False)
58	init_device = "meta" if use_meta_device else "cpu"	57	init_device = "meta" if use_meta_device else "cpu"
59	- logger.info_rank0(f"> Model initialization device: {init_device} (init_model_with_meta_device={use_meta_device})")	58	+ logger.info_rank0(
		59	+ f"> Model initialization device: {init_device} (init_model_with_meta_device={use_meta_device})"
		60	+ )
60		61
61	# 3. Load HF Config	62	# 3. Load HF Config
62	logger.info_rank0(f"> Loading AutoConfig from {model_args.model_name_or_path}...")	63	logger.info_rank0(f"> Loading AutoConfig from {model_args.model_name_or_path}...")
63	trust_remote_code = model_args.trust_remote_code	64	trust_remote_code = model_args.trust_remote_code
64	- hf_config = AutoConfig.from_pretrained(	65	+ hf_config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=trust_remote_code)
65	- model_args.model_name_or_path,
66	- trust_remote_code=trust_remote_code
67	- )
68		66
69	# 4. Load HF Model	67	# 4. Load HF Model
70	# Decide loading method based on init_device and whether training from scratch or fine-tuning.	68	# Decide loading method based on init_device and whether training from scratch or fine-tuning.
	@@ -79,11 +77,12 @@ class ModelFactory:
79	# Use ModelLoader to create model based on init_device	77	# Use ModelLoader to create model based on init_device
80	loader = ModelLoader(model_args, init_device=init_device)	78	loader = ModelLoader(model_args, init_device=init_device)
81	model, weights_path = loader.create_model(model_cls=model_cls)	79	model, weights_path = loader.create_model(model_cls=model_cls)
		80	+ hf_config = loader.hf_config
82		81
83	# 5. Build MindSpeed FSDP Configuration	82	# 5. Build MindSpeed FSDP Configuration
84	# Dynamically calculate Data Parallel (DP) Size	83	# Dynamically calculate Data Parallel (DP) Size
85	world_size = dist.get_world_size() if dist.is_initialized() else 1	84	world_size = dist.get_world_size() if dist.is_initialized() else 1
86	-	85	+
87	# Guard against division by zero if args are not set correctly	86	# Guard against division by zero if args are not set correctly
88	tp_size = parallel_args.tp_size	87	tp_size = parallel_args.tp_size
89	fsdp_size = parallel_args.fsdp_size	88	fsdp_size = parallel_args.fsdp_size
	@@ -97,12 +96,9 @@ class ModelFactory:
97		96
98	# MindSpeed FSDP will shard and wrap the CPU model based on the config.	97	# MindSpeed FSDP will shard and wrap the CPU model based on the config.
99	# The wrapped model automatically handles forward/backward communication.	98	# The wrapped model automatically handles forward/backward communication.
100	- # Pass init_device and weights_path for meta device support.	99	+ # Pass init_device, weights_path, and hf_config for meta device support.
101	model = MindSpeedParallelEngine(	100	model = MindSpeedParallelEngine(
102	- config=parallel_config,	101	+ config=parallel_config, model=model, init_device=init_device, weights_path=weights_path, hf_config=hf_config
103	- model=model,
104	- init_device=init_device,
105	- weights_path=weights_path
106	)	102	)
107		103
108	# 7. Move to target device	104	# 7. Move to target device
	@@ -123,26 +119,27 @@ class ModelFactory:
123	# --- 1. FSDP Plan ---	119	# --- 1. FSDP Plan ---
124	# Requirement: Apply FSDP to transformer layers	120	# Requirement: Apply FSDP to transformer layers
125	apply_modules = {	121	apply_modules = {
126	- parallel_args.fsdp_modules[0]: {'reshard_after_forward': parallel_args.reshard_after_forward,	122	+ parallel_args.fsdp_modules[0]: {
127	- 'shard_placement_fn': parallel_args.shard_placement_fn},	123	+ 'reshard_after_forward': parallel_args.reshard_after_forward,
		124	+ 'shard_placement_fn': parallel_args.shard_placement_fn,
		125	+ },
128	}	126	}
129	for modules in parallel_args.fsdp_modules[1:]:	127	for modules in parallel_args.fsdp_modules[1:]:
130	- apply_modules[modules] = {'reshard_after_forward': parallel_args.reshard_after_forward,}	128	+ apply_modules[modules] = {
		129	+ 'reshard_after_forward': parallel_args.reshard_after_forward,
		130	+ }
131	fsdp_plan = FSDPPlanConfig(	131	fsdp_plan = FSDPPlanConfig(
132	ignored_modules=parallel_args.ignored_modules if parallel_args.ignored_modules else [],	132	ignored_modules=parallel_args.ignored_modules if parallel_args.ignored_modules else [],
133	- apply_modules= apply_modules,	133	+ apply_modules=apply_modules,
134	param_dtype=parallel_args.param_dtype,	134	param_dtype=parallel_args.param_dtype,
135	reduce_dtype=parallel_args.reduce_dtype,	135	reduce_dtype=parallel_args.reduce_dtype,
136	num_to_forward_prefetch=parallel_args.num_to_forward_prefetch,	136	num_to_forward_prefetch=parallel_args.num_to_forward_prefetch,
137	- num_to_backward_prefetch=parallel_args.num_to_backward_prefetch	137	+ num_to_backward_prefetch=parallel_args.num_to_backward_prefetch,
138	)	138	)
139		139
140	# --- 2. Tensor Parallel Plan ---	140	# --- 2. Tensor Parallel Plan ---
141	# Requirement: Column Parallel for Q/K/V/Gate/Up, Row Parallel for O/Down	141	# Requirement: Column Parallel for Q/K/V/Gate/Up, Row Parallel for O/Down
142	- tp_plan = TPPlanConfig(	142	+ tp_plan = TPPlanConfig(colwise_parallel=parallel_args.tp_colwise, rowwise_parallel=parallel_args.tp_rowwise)
143	- colwise_parallel=parallel_args.tp_colwise,
144	- rowwise_parallel=parallel_args.tp_rowwise
145	- )
146		143
147	# --- 3. Expert Parallel Plan ---	144	# --- 3. Expert Parallel Plan ---
148	# For Mixture-of-Experts (MoE) models	145	# For Mixture-of-Experts (MoE) models
	@@ -155,11 +152,7 @@ class ModelFactory:
155	dispatcher=parallel_args.ep_dispatcher,	152	dispatcher=parallel_args.ep_dispatcher,
156	)	153	)
157		154
158	-	155	+ cp_plan = CPPlanConfig(context_parallel_type=parallel_args.cp_type, is_pack=getattr(model_args, "pack", False))
159	- cp_plan = CPPlanConfig(
160	- context_parallel_type=parallel_args.cp_type,
161	- is_pack=getattr(model_args, "pack", False)
162	- )
163		156
164	# --- 4. Recompute Plan ---	157	# --- 4. Recompute Plan ---
165	# Activation Checkpointing	158	# Activation Checkpointing
	@@ -183,30 +176,24 @@ class ModelFactory:
183	config = ParallelEngineConfig(	176	config = ParallelEngineConfig(
184	# Parallelism parameters	177	# Parallelism parameters
185	data_parallel_size=dp_size,	178	data_parallel_size=dp_size,
186	-
187	fully_shard_parallel_size=fsdp_size,	179	fully_shard_parallel_size=fsdp_size,
188	fsdp_plan=fsdp_plan,	180	fsdp_plan=fsdp_plan,
189	-
190	tensor_parallel_size=tp_size,	181	tensor_parallel_size=tp_size,
191	tp_plan=tp_plan,	182	tp_plan=tp_plan,
192	-
193	# Expert Parallelism	183	# Expert Parallelism
194	expert_parallel_size=ep_size,	184	expert_parallel_size=ep_size,
195	expert_fully_shard_parallel_size=ep_fsdp_size,	185	expert_fully_shard_parallel_size=ep_fsdp_size,
196	expert_data_parallel_size=dp_size, # Usually EP data parallel size matches global or has specific logic	186	expert_data_parallel_size=dp_size, # Usually EP data parallel size matches global or has specific logic
197	ep_plan=ep_plan,	187	ep_plan=ep_plan,
198	-
199	# Context Parallelism	188	# Context Parallelism
200	context_parallel_size=parallel_args.cp_size,	189	context_parallel_size=parallel_args.cp_size,
201	context_parallel_type=parallel_args.cp_type,	190	context_parallel_type=parallel_args.cp_type,
202	cp_plan=cp_plan,	191	cp_plan=cp_plan,
203	-
204	# Recomputation	192	# Recomputation
205	recompute=parallel_args.recompute,	193	recompute=parallel_args.recompute,
206	recompute_plan=recompute_plan,	194	recompute_plan=recompute_plan,
207	-
208	# Quantization	195	# Quantization
209	- quantization_plan = quantization_plan	196	+ quantization_plan=quantization_plan,
210	)	197	)
211		198
212	return config	199	return config
	@@ -227,4 +214,4 @@ class AutoModelFactory:
227	"""	214	"""
228	Factory method that forwards arguments to the specific implementation.	215	Factory method that forwards arguments to the specific implementation.
229	"""	216	"""
230	- return ModelFactory.create(args, *kwargs)	217	+ return ModelFactory.create(args, *kwargs)

Mmindspeed_llm/fsdp2/models/model_loader.py+138-80

	@@ -5,11 +5,12 @@ import json
5	from contextlib import contextmanager	5	from contextlib import contextmanager
6	from typing import Optional, Dict, Tuple, Set	6	from typing import Optional, Dict, Tuple, Set
7		7
		8	+from safetensors import safe_open
8	import torch	9	import torch
9	-import torch.nn as nn	10	+from torch import nn
10	from torch.distributed.tensor import distribute_tensor	11	from torch.distributed.tensor import distribute_tensor
11	from transformers import AutoConfig, AutoModelForCausalLM	12	from transformers import AutoConfig, AutoModelForCausalLM
12	-from safetensors import safe_open	13	+
13	try:	14	try:
14	from transformers.modeling_utils import no_init_weights	15	from transformers.modeling_utils import no_init_weights
15	except ImportError:	16	except ImportError:
	@@ -18,6 +19,8 @@ except ImportError:
18		19
19	from mindspeed_llm.fsdp2.utils.logging import get_logger	20	from mindspeed_llm.fsdp2.utils.logging import get_logger
20	from mindspeed_llm.fsdp2.utils.global_vars import get_args	21	from mindspeed_llm.fsdp2.utils.global_vars import get_args
		22	+from mindspeed_llm.fsdp2.checkpoint.weight_conv_adapter import WeightConvAdapter
		23	+
21	logger = get_logger(__name__)	24	logger = get_logger(__name__)
22		25
23		26
	@@ -42,6 +45,7 @@ def init_empty_weights():
42	if param.device == torch.device("meta")	45	if param.device == torch.device("meta")
43	else param_cls(module._parameters[name].to("meta"), **kwargs)	46	else param_cls(module._parameters[name].to("meta"), **kwargs)
44	)	47	)
		48	+
45	try:	49	try:
46	nn.Module.register_parameter = register_empty_parameter	50	nn.Module.register_parameter = register_empty_parameter
47	yield	51	yield
	@@ -73,7 +77,7 @@ def _find_submodule(module: nn.Module, name: str) -> Tuple[nn.Module, str]:
73	# ==============================================================================	77	# ==============================================================================
74	class ModelLoader:	78	class ModelLoader:
75	"""Load model on CPU or meta device."""	79	"""Load model on CPU or meta device."""
76	-	80	+
77	def __init__(self, model_args, init_device: str = "cpu"):	81	def __init__(self, model_args, init_device: str = "cpu"):
78	self.model_args = model_args	82	self.model_args = model_args
79	self.init_device = init_device	83	self.init_device = init_device
	@@ -81,7 +85,7 @@ class ModelLoader:
81	self.model_path = model_args.model_name_or_path	85	self.model_path = model_args.model_name_or_path
82	self.train_from_scratch = getattr(model_args, 'train_from_scratch', False)	86	self.train_from_scratch = getattr(model_args, 'train_from_scratch', False)
83	self.hf_config = None	87	self.hf_config = None
84	-	88	+
85	def load_config(self) -> AutoConfig:	89	def load_config(self) -> AutoConfig:
86	"""Load HuggingFace model config."""	90	"""Load HuggingFace model config."""
87	logger.info_rank0(f"> Loading config from {self.model_path}...")	91	logger.info_rank0(f"> Loading config from {self.model_path}...")
	@@ -90,20 +94,20 @@ class ModelLoader:
90	trust_remote_code=self.trust_remote_code,	94	trust_remote_code=self.trust_remote_code,
91	# Context parallelism requires uniformly applying a patch to the attention component,	95	# Context parallelism requires uniformly applying a patch to the attention component,
92	# which is unified here into the `eager` implementation part	96	# which is unified here into the `eager` implementation part
93	- attn_implementation="eager" if get_args().cp_size >1 else None,	97	+ attn_implementation="eager" if get_args().cp_size > 1 else None,
94	)	98	)
95	return self.hf_config	99	return self.hf_config
96	-	100	+
97	def create_model(self, model_cls=None) -> Tuple[nn.Module, Optional[str]]:	101	def create_model(self, model_cls=None) -> Tuple[nn.Module, Optional[str]]:
98	"""Create model based on init_device."""	102	"""Create model based on init_device."""
99	if self.hf_config is None:	103	if self.hf_config is None:
100	self.load_config()	104	self.load_config()
101	-	105	+
102	if self.init_device == "meta":	106	if self.init_device == "meta":
103	return self._create_on_meta(model_cls)	107	return self._create_on_meta(model_cls)
104	else:	108	else:
105	return self._create_on_cpu(model_cls)	109	return self._create_on_cpu(model_cls)
106	-	110	+
107	def _create_on_cpu(self, model_cls=None) -> Tuple[nn.Module, None]:	111	def _create_on_cpu(self, model_cls=None) -> Tuple[nn.Module, None]:
108	"""Create and load model on CPU."""	112	"""Create and load model on CPU."""
109	if model_cls is not None:	113	if model_cls is not None:
	@@ -113,14 +117,12 @@ class ModelLoader:
113	config=self.hf_config,	117	config=self.hf_config,
114	low_cpu_mem_usage=True,	118	low_cpu_mem_usage=True,
115	device_map="cpu",	119	device_map="cpu",
116	- torch_dtype=torch.float32	120	+ torch_dtype=torch.float32,
117	)	121	)
118	elif self.train_from_scratch:	122	elif self.train_from_scratch:
119	logger.info_rank0("> Creating model with random weights on CPU...")	123	logger.info_rank0("> Creating model with random weights on CPU...")
120	model = AutoModelForCausalLM.from_config(	124	model = AutoModelForCausalLM.from_config(
121	- self.hf_config,	125	+ self.hf_config, trust_remote_code=self.trust_remote_code, torch_dtype=torch.float32
122	- trust_remote_code=self.trust_remote_code,
123	- torch_dtype=torch.float32
124	)	126	)
125	else:	127	else:
126	logger.info_rank0(f"> Loading pretrained model on CPU from {self.model_path}...")	128	logger.info_rank0(f"> Loading pretrained model on CPU from {self.model_path}...")
	@@ -130,11 +132,11 @@ class ModelLoader:
130	trust_remote_code=self.trust_remote_code,	132	trust_remote_code=self.trust_remote_code,
131	torch_dtype=torch.float32,	133	torch_dtype=torch.float32,
132	low_cpu_mem_usage=True,	134	low_cpu_mem_usage=True,
133	- device_map="cpu"	135	+ device_map="cpu",
134	)	136	)
135	-	137	+
136	return model, None	138	return model, None
137	-	139	+
138	def _create_on_meta(self, model_cls=None) -> Tuple[nn.Module, Optional[str]]:	140	def _create_on_meta(self, model_cls=None) -> Tuple[nn.Module, Optional[str]]:
139	"""Create empty model on meta device."""	141	"""Create empty model on meta device."""
140	weights_path = None if self.train_from_scratch else self.model_path	142	weights_path = None if self.train_from_scratch else self.model_path
	@@ -150,19 +152,15 @@ class ModelLoader:
150	logger.info_rank0("> Creating empty model on meta device for random init...")	152	logger.info_rank0("> Creating empty model on meta device for random init...")
151	with init_empty_weights():	153	with init_empty_weights():
152	model = AutoModelForCausalLM.from_config(	154	model = AutoModelForCausalLM.from_config(
153	- self.hf_config,	155	+ self.hf_config, trust_remote_code=self.trust_remote_code, torch_dtype=torch.float32
154	- trust_remote_code=self.trust_remote_code,
155	- torch_dtype=torch.float32
156	)	156	)
157	else:	157	else:
158	logger.info_rank0(f"> Creating empty model on meta device (weights: {self.model_path})...")	158	logger.info_rank0(f"> Creating empty model on meta device (weights: {self.model_path})...")
159	with init_empty_weights(), no_init_weights():	159	with init_empty_weights(), no_init_weights():
160	model = AutoModelForCausalLM.from_config(	160	model = AutoModelForCausalLM.from_config(
161	- self.hf_config,	161	+ self.hf_config, trust_remote_code=self.trust_remote_code, torch_dtype=torch.float32
162	- trust_remote_code=self.trust_remote_code,
163	- torch_dtype=torch.float32
164	)	162	)
165	-	163	+
166	logger.info_rank0(f"> Model created on meta device. Weights path: {weights_path}")	164	logger.info_rank0(f"> Model created on meta device. Weights path: {weights_path}")
167	return model, weights_path	165	return model, weights_path
168		166
	@@ -173,108 +171,166 @@ class ModelLoader:
173	class WeightLoader:	171	class WeightLoader:
174	"""	172	"""
175	Load weights into FSDP-wrapped model.	173	Load weights into FSDP-wrapped model.
		174	+ Supports weight conversion via transformers conversion_mapping API.
176	"""	175	"""
177	-	176	+
178	@staticmethod	177	@staticmethod
179	def load(	178	def load(
180	- model: nn.Module,	179	+ model: nn.Module,
181	- weights_path: Optional[str],	180	+ weights_path: Optional[str],
182	- device: Optional[str] = None	181	+ device: Optional[str] = None,
		182	+ hf_config: Optional[AutoConfig] = None,
183	) -> None:	183	) -> None:
184	"""Load or initialize weights after FSDP wrapping."""	184	"""Load or initialize weights after FSDP wrapping."""
185	if device is None:	185	if device is None:
186	device = torch.accelerator.current_accelerator().type	186	device = torch.accelerator.current_accelerator().type
187	-	187	+
188	if weights_path is None:	188	if weights_path is None:
189	WeightLoader._init_random(model, device)	189	WeightLoader._init_random(model, device)
190	else:	190	else:
191	- WeightLoader._load_pretrained(model, weights_path, device)	191	+ WeightLoader._load_pretrained(model, weights_path, device, hf_config)
192	-	192	+
193	@staticmethod	193	@staticmethod
194	def _init_random(model: nn.Module, device: str) -> None:	194	def _init_random(model: nn.Module, device: str) -> None:
195	"""Initialize model with random weights."""	195	"""Initialize model with random weights."""
196	logger.info_rank0(f"> Initializing random weights on {device}...")	196	logger.info_rank0(f"> Initializing random weights on {device}...")
197	-	197	+
198	model.to_empty(device=device)	198	model.to_empty(device=device)
199	model = model.float()	199	model = model.float()
200	reset_hf_initialized_flag(model)	200	reset_hf_initialized_flag(model)
201	-	201	+
202	if hasattr(model, 'init_weights'):	202	if hasattr(model, 'init_weights'):
203	model.init_weights()	203	model.init_weights()
204	-	204	+
205	logger.info_rank0("> Random initialization done")	205	logger.info_rank0("> Random initialization done")
206	-	206	+
207	@staticmethod	207	@staticmethod
208	@torch.no_grad()	208	@torch.no_grad()
209	- def _load_pretrained(model: nn.Module, weights_path: str, device: str) -> None:	209	+ def _load_pretrained(
		210	+ model: nn.Module, weights_path: str, device: str, hf_config: Optional[AutoConfig] = None
		211	+ ) -> None:
210	"""	212	"""
211	- Load pretrained weights.	213	+ Load pretrained weights with dynamic weight conversion support.
212	- """	214	+ Uses transformers conversion_mapping API for model-type-specific weight transformations.
		215	+ """
		216	+
213	logger.info_rank0(f"> Loading pretrained weights from {weights_path}...")	217	logger.info_rank0(f"> Loading pretrained weights from {weights_path}...")
214	-	218	+
		219	+ model_type = getattr(hf_config, 'model_type', None) if hf_config else None
		220	+ adapter = WeightConvAdapter(model_type=model_type) if model_type else None
		221	+
215	# Step 1: Save buffers before to_empty	222	# Step 1: Save buffers before to_empty
216	buffer_dict = {name: buffer.clone() for name, buffer in model.named_buffers()}	223	buffer_dict = {name: buffer.clone() for name, buffer in model.named_buffers()}
217	parameter_names_to_load = {name for name, _ in model.named_parameters()}	224	parameter_names_to_load = {name for name, _ in model.named_parameters()}
218	-	225	+
219	logger.info_rank0(f"> Saved {len(buffer_dict)} buffers, {len(parameter_names_to_load)} parameters to load")	226	logger.info_rank0(f"> Saved {len(buffer_dict)} buffers, {len(parameter_names_to_load)} parameters to load")
220	-	227	+
221	# Step 2: Materialize model to device	228	# Step 2: Materialize model to device
222	model.to_empty(device=device)	229	model.to_empty(device=device)
223	model = model.float()	230	model = model.float()
224	logger.info_rank0(f"> Model materialized to {device}")	231	logger.info_rank0(f"> Model materialized to {device}")
225	-	232	+
226	# Step 3: Load state dict and dispatch parameters	233	# Step 3: Load state dict and dispatch parameters
227	state_dict_files = WeightLoader._get_state_dict_files(weights_path)	234	state_dict_files = WeightLoader._get_state_dict_files(weights_path)
228	-	235	+
		236	+ pending_weights: Dict[str, torch.Tensor] = {}
		237	+ collected_keys: Set[str] = set()
		238	+ converted_groups: Dict[str, Dict[str, list]] = {}
		239	+
229	for state_dict_file in state_dict_files:	240	for state_dict_file in state_dict_files:
230	for name, tensor in WeightLoader._iterate_state_dict(state_dict_file):	241	for name, tensor in WeightLoader._iterate_state_dict(state_dict_file):
		242	+ if adapter and adapter.has_conversions:
		243	+ renamed_key, source_pattern = adapter.rename_key(name)
		244	+
		245	+ if source_pattern is not None:
		246	+ converter = adapter.match_converter(source_pattern)
		247	+ if converter:
		248	+ converted_groups.setdefault(renamed_key, {}).setdefault(source_pattern, []).append(
		249	+ (name, tensor)
		250	+ )
		251	+ collected_keys.add(name)
		252	+ continue
		253	+
		254	+ name = renamed_key
		255	+
231	if name in buffer_dict:	256	if name in buffer_dict:
232	- # Update buffer in buffer_dict
233	buffer_dict[name] = tensor.clone()	257	buffer_dict[name] = tensor.clone()
234	elif name in parameter_names_to_load:	258	elif name in parameter_names_to_load:
235	parameter_names_to_load.remove(name)	259	parameter_names_to_load.remove(name)
236	WeightLoader._dispatch_parameter(model, name, tensor)	260	WeightLoader._dispatch_parameter(model, name, tensor)
237	else:	261	else:
238	- logger.debug(f"> Unexpected key in state dict: {name}")	262	+ pending_weights[name] = tensor
239	-	263	+
		264	+ # Dispatch pending weights from first pass
		265	+ for name in list(pending_weights.keys()):
		266	+ if name in buffer_dict:
		267	+ buffer_dict[name] = pending_weights.pop(name).clone()
		268	+ elif name in parameter_names_to_load:
		269	+ parameter_names_to_load.remove(name)
		270	+ WeightLoader._dispatch_parameter(model, name, pending_weights.pop(name))
		271	+
		272	+ # Dispatch converted weights (streaming via native WeightConverter.convert)
		273	+ for target_name, collected in converted_groups.items():
		274	+ source_pattern = next(iter(collected))
		275	+ converter = adapter.match_converter(source_pattern)
		276	+ if not converter:
		277	+ continue
		278	+
		279	+ full_name = WeightLoader._find_model_param(parameter_names_to_load, target_name)
		280	+ if not full_name:
		281	+ logger.debug(f"> No model param matching converted target: {target_name}")
		282	+ continue
		283	+
		284	+ parameter_names_to_load.discard(full_name)
		285	+ original_keys = [item[0] for item in collected[source_pattern]]
		286	+ tensors_only = {sp: [item[1] for item in items] for sp, items in collected.items()}
		287	+ for name, tensor in WeightConvAdapter.dispatch_converted(converter, full_name, tensors_only, original_keys):
		288	+ WeightLoader._dispatch_parameter(model, name, tensor)
		289	+
		290	+ # Log remaining unexpected keys
		291	+ for name in pending_weights:
		292	+ if name not in collected_keys:
		293	+ logger.debug(f"> Unexpected key in state dict: {name}")
		294	+ del pending_weights
		295	+
240	# Step 4: Post-process (restore buffers, handle missing params)	296	# Step 4: Post-process (restore buffers, handle missing params)
241	WeightLoader._post_process(model, buffer_dict, parameter_names_to_load, distribute_tensor)	297	WeightLoader._post_process(model, buffer_dict, parameter_names_to_load, distribute_tensor)
242	if torch.distributed.is_initialized():	298	if torch.distributed.is_initialized():
243	torch.distributed.barrier()	299	torch.distributed.barrier()
244	logger.info_rank0("> Pretrained weights loaded successfully")	300	logger.info_rank0("> Pretrained weights loaded successfully")
245	-	301	+
246	@staticmethod	302	@staticmethod
247	def _get_state_dict_files(weights_path: str):	303	def _get_state_dict_files(weights_path: str):
248	"""Get list of state dict files."""	304	"""Get list of state dict files."""
249	# Check for safetensors index	305	# Check for safetensors index
250	index_file = os.path.join(weights_path, "model.safetensors.index.json")	306	index_file = os.path.join(weights_path, "model.safetensors.index.json")
251	if os.path.exists(index_file):	307	if os.path.exists(index_file):
252	- with open(index_file, 'r') as f:	308	+ with open(index_file, 'r', encoding='utf-8') as f:
253	index = json.load(f)	309	index = json.load(f)
254	files = set(index["weight_map"].values())	310	files = set(index["weight_map"].values())
255	return [os.path.join(weights_path, f) for f in sorted(files)]	311	return [os.path.join(weights_path, f) for f in sorted(files)]
256	-	312	+
257	# Check for single safetensors file	313	# Check for single safetensors file
258	single_safetensor = os.path.join(weights_path, "model.safetensors")	314	single_safetensor = os.path.join(weights_path, "model.safetensors")
259	if os.path.exists(single_safetensor):	315	if os.path.exists(single_safetensor):
260	return [single_safetensor]	316	return [single_safetensor]
261	-	317	+
262	# Check for multiple safetensors files	318	# Check for multiple safetensors files
263	safetensor_files = sorted(glob.glob(os.path.join(weights_path, "*.safetensors")))	319	safetensor_files = sorted(glob.glob(os.path.join(weights_path, "*.safetensors")))
264	if safetensor_files:	320	if safetensor_files:
265	return safetensor_files	321	return safetensor_files
266	-	322	+
267	# Check for pytorch files	323	# Check for pytorch files
268	pytorch_files = sorted(glob.glob(os.path.join(weights_path, "*.bin")))	324	pytorch_files = sorted(glob.glob(os.path.join(weights_path, "*.bin")))
269	if pytorch_files:	325	if pytorch_files:
270	return pytorch_files	326	return pytorch_files
271	-	327	+
272	pytorch_files = sorted(glob.glob(os.path.join(weights_path, "*.pt")))	328	pytorch_files = sorted(glob.glob(os.path.join(weights_path, "*.pt")))
273	if pytorch_files:	329	if pytorch_files:
274	return pytorch_files	330	return pytorch_files
275	-	331	+
276	raise FileNotFoundError(f"No weight files found in {weights_path}")	332	raise FileNotFoundError(f"No weight files found in {weights_path}")
277	-	333	+
278	@staticmethod	334	@staticmethod
279	def _iterate_state_dict(filepath: str):	335	def _iterate_state_dict(filepath: str):
280	"""Iterate over state dict file, yielding (key, tensor) pairs."""	336	"""Iterate over state dict file, yielding (key, tensor) pairs."""
	@@ -286,7 +342,17 @@ class WeightLoader:
286	state_dict = torch.load(filepath, map_location="cpu", weights_only=True)	342	state_dict = torch.load(filepath, map_location="cpu", weights_only=True)
287	for key, tensor in state_dict.items():	343	for key, tensor in state_dict.items():
288	yield key, tensor	344	yield key, tensor
289	-	345	+
		346	+ @staticmethod
		347	+ def _find_model_param(param_names: Set[str], target_name: str) -> Optional[str]:
		348	+ """Find full model param name that ends with target_name."""
		349	+ if target_name in param_names:
		350	+ return target_name
		351	+ for name in param_names:
		352	+ if name.endswith(target_name):
		353	+ return name
		354	+ return None
		355	+
290	@staticmethod	356	@staticmethod
291	def _dispatch_parameter(	357	def _dispatch_parameter(
292	model: nn.Module,	358	model: nn.Module,
	@@ -343,31 +409,26 @@ class WeightLoader:
343	non_blocking=True,	409	non_blocking=True,
344	)	410	)
345	local_target.copy_(local_dev, non_blocking=True)	411	local_target.copy_(local_dev, non_blocking=True)
346	-	412	+
347	@staticmethod	413	@staticmethod
348	- def _dispatch_buffer(	414	+ def _dispatch_buffer(model: nn.Module, name: str, buffer: torch.Tensor, dtensor_factory) -> None:
349	- model: nn.Module,
350	- name: str,
351	- buffer: torch.Tensor,
352	- dtensor_factory
353	- ) -> None:
354	"""	415	"""
355	Assign buffer to model.	416	Assign buffer to model.
356	"""	417	"""
357	module, local_name = _find_submodule(model, name)	418	module, local_name = _find_submodule(model, name)
358	orig_buffer = dict(module.named_buffers(recurse=False))[local_name]	419	orig_buffer = dict(module.named_buffers(recurse=False))[local_name]
359	-	420	+
360	if hasattr(orig_buffer, "device_mesh"):	421	if hasattr(orig_buffer, "device_mesh"):
361	device_mesh = orig_buffer.device_mesh	422	device_mesh = orig_buffer.device_mesh
362	placements = orig_buffer.placements	423	placements = orig_buffer.placements
363	- module.register_buffer(local_name, dtensor_factory(	424	+ module.register_buffer(
364	- buffer.to(dtype=orig_buffer.dtype),	425	+ local_name, dtensor_factory(buffer.to(dtype=orig_buffer.dtype), device_mesh, placements)
365	- device_mesh,	426	+ )
366	- placements
367	- ))
368	else:	427	else:
369	- dict(module.named_buffers(recurse=False))[local_name].copy_(buffer.to(device=orig_buffer.device, dtype=orig_buffer.dtype))	428	+ dict(module.named_buffers(recurse=False))[local_name].copy_(
370	-	429	+ buffer.to(device=orig_buffer.device, dtype=orig_buffer.dtype)
		430	+ )
		431	+
371	@staticmethod	432	@staticmethod
372	def _init_parameter(model: nn.Module, name: str) -> None:	433	def _init_parameter(model: nn.Module, name: str) -> None:
373	"""	434	"""
	@@ -376,23 +437,20 @@ class WeightLoader:
376	pieces = name.split(".")	437	pieces = name.split(".")
377	init_func = None	438	init_func = None
378	module = model	439	module = model
379	-	440	+
380	for piece in pieces[:-1]:	441	for piece in pieces[:-1]:
381	if hasattr(module, "_init_weights"):	442	if hasattr(module, "_init_weights"):
382	init_func = getattr(module, "_init_weights", None)	443	init_func = getattr(module, "_init_weights", None)
383	module = getattr(module, piece)	444	module = getattr(module, piece)
384	-	445	+
385	if init_func is not None:	446	if init_func is not None:
386	module.apply(init_func)	447	module.apply(init_func)
387	else:	448	else:
388	logger.warning(f"> Cannot find _init_weights for {name}, skipping initialization")	449	logger.warning(f"> Cannot find _init_weights for {name}, skipping initialization")
389	-	450	+
390	@staticmethod	451	@staticmethod
391	def _post_process(	452	def _post_process(
392	- model: nn.Module,	453	+ model: nn.Module, buffer_dict: Dict[str, torch.Tensor], parameter_names_left: Set[str], dtensor_factory
393	- buffer_dict: Dict[str, torch.Tensor],
394	- parameter_names_left: Set[str],
395	- dtensor_factory
396	) -> None:	454	) -> None:
397	"""	455	"""
398	Post-process after weight loading.	456	Post-process after weight loading.
	@@ -403,9 +461,9 @@ class WeightLoader:
403	WeightLoader._dispatch_buffer(model, name, buffer, dtensor_factory)	461	WeightLoader._dispatch_buffer(model, name, buffer, dtensor_factory)
404	except Exception as e:	462	except Exception as e:
405	logger.warning(f"> Failed to restore buffer {name}: {e}")	463	logger.warning(f"> Failed to restore buffer {name}: {e}")
406	-	464	+
407	logger.info_rank0(f"> Restored {len(buffer_dict)} buffers")	465	logger.info_rank0(f"> Restored {len(buffer_dict)} buffers")
408	-	466	+
409	# Initialize missing parameters	467	# Initialize missing parameters
410	if parameter_names_left:	468	if parameter_names_left:
411	logger.info_rank0(f"> Missing {parameter_names_left} parameters, initializing them...")	469	logger.info_rank0(f"> Missing {parameter_names_left} parameters, initializing them...")
	@@ -414,7 +472,7 @@ class WeightLoader:
414	WeightLoader._init_parameter(model, name)	472	WeightLoader._init_parameter(model, name)
415	except Exception as e:	473	except Exception as e:
416	logger.warning(f"> Failed to initialize {name}: {e}")	474	logger.warning(f"> Failed to initialize {name}: {e}")
417	-	475	+
418	# Tie embeddings if needed	476	# Tie embeddings if needed
419	if getattr(model.config, "tie_word_embeddings", True):	477	if getattr(model.config, "tie_word_embeddings", True):
420	try:	478	try:
	@@ -424,7 +482,7 @@ class WeightLoader:
424	output_embeddings.register_parameter(	482	output_embeddings.register_parameter(
425	"weight",	483	"weight",
426	input_embeddings.weight,	484	input_embeddings.weight,
427	- )	485	+ )
428	logger.info_rank0("> Tied input/output embeddings")	486	logger.info_rank0("> Tied input/output embeddings")
429	except Exception as e:	487	except Exception as e:
430	- logger.warning(f"> Failed to tie embeddings: {e}")	488	+ logger.warning(f"> Failed to tie embeddings: {e}")

Mmindspeed_llm/fsdp2/models/qwen3/qwen3_moe.py+63-55

	@@ -1,10 +1,15 @@
1	# Copyright (c) 2025, HUAWEI CORPORATION. All rights reserved.	1	# Copyright (c) 2025, HUAWEI CORPORATION. All rights reserved.
2	# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.	2	# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
3	-import os
4	from typing import Optional, Union	3	from typing import Optional, Union
		4	+from importlib.metadata import version
5		5
6	import torch	6	import torch
7	import torch.nn.functional as F	7	import torch.nn.functional as F
		8	+
		9	+try:
		10	+ import torch_npu
		11	+except ImportError:
		12	+ pass
8	import transformers	13	import transformers
9	from torch import nn	14	from torch import nn
10	from torch.distributed.tensor import DTensor	15	from torch.distributed.tensor import DTensor
	@@ -17,19 +22,18 @@ from transformers.utils import can_return_tuple
17		22
18	from mindspeed.core.fusions.grouped_matmul import Ops	23	from mindspeed.core.fusions.grouped_matmul import Ops
19	from mindspeed.patch_utils import MindSpeedPatchesManager as pm	24	from mindspeed.patch_utils import MindSpeedPatchesManager as pm
		25	+
		26	+try:
		27	+ from mindspeed.ops.npu_moe_token_permute import npu_moe_token_permute
		28	+ from mindspeed.ops.npu_moe_token_unpermute import npu_moe_token_unpermute
		29	+except ImportError:
		30	+ pass
		31	+
20	from mindspeed_llm.fsdp2.features.async_offload import async_save_on_cpu	32	from mindspeed_llm.fsdp2.features.async_offload import async_save_on_cpu
21	from mindspeed_llm.fsdp2.models.common.fusions import fused_rmsnorm_forward, apply_rotary_pos_emb	33	from mindspeed_llm.fsdp2.models.common.fusions import fused_rmsnorm_forward, apply_rotary_pos_emb
22	from mindspeed_llm.fsdp2.models.common.modules import LMHead	34	from mindspeed_llm.fsdp2.models.common.modules import LMHead
23	from mindspeed_llm.fsdp2.utils.global_vars import get_args	35	from mindspeed_llm.fsdp2.utils.global_vars import get_args
24		36
25	-try:
26	- import torch_npu
27	- from mindspeed.ops.npu_moe_token_permute import npu_moe_token_permute
28	- from mindspeed.ops.npu_moe_token_unpermute import npu_moe_token_unpermute
29	- from mindspeed.ops.gmm_mxfp8 import npu_quant_group_gemm
30	-except ImportError:
31	- pass
32	-
33		37
34	class Qwen3MoEForCausalLM(transformers.Qwen3MoePreTrainedModel):	38	class Qwen3MoEForCausalLM(transformers.Qwen3MoePreTrainedModel):
35	_tied_weights_keys = ["lm_head.weight"]	39	_tied_weights_keys = ["lm_head.weight"]
	@@ -51,19 +55,20 @@ class Qwen3MoEForCausalLM(transformers.Qwen3MoePreTrainedModel):
51		55
52	@can_return_tuple	56	@can_return_tuple
53	def forward(	57	def forward(
54	- self,	58	+ self,
55	- input_ids: Optional[torch.LongTensor] = None,	59	+ input_ids: Optional[torch.LongTensor] = None,
56	- attention_mask: Optional[torch.Tensor] = None,	60	+ attention_mask: Optional[torch.Tensor] = None,
57	- position_ids: Optional[torch.LongTensor] = None,	61	+ position_ids: Optional[torch.LongTensor] = None,
58	- past_key_values: Optional[Cache] = None,	62	+ past_key_values: Optional[Cache] = None,
59	- inputs_embeds: Optional[torch.FloatTensor] = None,	63	+ inputs_embeds: Optional[torch.FloatTensor] = None,
60	- labels: Optional[torch.LongTensor] = None,	64	+ labels: Optional[torch.LongTensor] = None,
61	- use_cache: Optional[bool] = None,	65	+ use_cache: Optional[bool] = None,
62	- output_router_logits: Optional[bool] = None,	66	+ output_router_logits: Optional[bool] = None,
63	- cache_position: Optional[torch.LongTensor] = None,	67	+ cache_position: Optional[torch.LongTensor] = None,
64	- logits_to_keep: Union[int, torch.Tensor] = 0,	68	+ logits_to_keep: Union[int, torch.Tensor] = 0,
65	- loss_ctx: Optional[callable] = None,	69	+ loss_ctx: Optional[callable] = None,
66	- **kwargs) -> MoeCausalLMOutputWithPast:	70	+ **kwargs,
		71	+ ) -> MoeCausalLMOutputWithPast:
67	r"""	72	r"""
68	labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, optional):	73	labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, optional):
69	Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,	74	Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
	@@ -85,7 +90,8 @@ class Qwen3MoEForCausalLM(transformers.Qwen3MoePreTrainedModel):
85	>>> generate_ids = model.generate(inputs.input_ids, max_length=30)	90	>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
86	>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]	91	>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
87	"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."	92	"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
88	- ```"""	93	+ ```
		94	+ """
89		95
90	output_router_logits = (	96	output_router_logits = (
91	output_router_logits if output_router_logits is not None else self.config.output_router_logits	97	output_router_logits if output_router_logits is not None else self.config.output_router_logits
	@@ -126,8 +132,7 @@ class Qwen3MoEForCausalLM(transformers.Qwen3MoePreTrainedModel):
126	attention_mask,	132	attention_mask,
127	)	133	)
128	if labels is not None:	134	if labels is not None:
129	- loss += self.router_aux_loss_coef * aux_loss.to(	135	+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
130	- loss.device) # make sure to reside in the same device
131		136
132	return MoeCausalLMOutputWithPast(	137	return MoeCausalLMOutputWithPast(
133	loss=loss,	138	loss=loss,
	@@ -144,20 +149,24 @@ class Qwen3MoEForCausalLM(transformers.Qwen3MoePreTrainedModel):
144	"""patching the transformers model."""	149	"""patching the transformers model."""
145	args = get_args()	150	args = get_args()
146	if getattr(args, "moe_grouped_gemm", False):	151	if getattr(args, "moe_grouped_gemm", False):
147	- pm.register_patch("transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock",	152	+ pm.register_patch(
148	- Qwen3MoeSparseFusedMoeBlock)	153	+ "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock", Qwen3MoeSparseFusedMoeBlock
		154	+ )
149		155
150	if getattr(args, "activation_offload", False):	156	if getattr(args, "activation_offload", False):
151	- pm.register_patch("transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeModel.forward",	157	+ pm.register_patch(
152	- qwen3_moe_model_forward)	158	+ "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeModel.forward", qwen3_moe_model_forward
		159	+ )
153		160
154	if getattr(args, "use_fused_rmsnorm", False):	161	if getattr(args, "use_fused_rmsnorm", False):
155	- pm.register_patch("transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeRMSNorm.forward",	162	+ pm.register_patch(
156	- fused_rmsnorm_forward)	163	+ "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeRMSNorm.forward", fused_rmsnorm_forward
		164	+ )
157		165
158	if getattr(args, "use_fused_rotary_pos_emb", False):	166	if getattr(args, "use_fused_rotary_pos_emb", False):
159	- pm.register_patch("transformers.models.qwen3_moe.modeling_qwen3_moe.apply_rotary_pos_emb",	167	+ pm.register_patch(
160	- apply_rotary_pos_emb)	168	+ "transformers.models.qwen3_moe.modeling_qwen3_moe.apply_rotary_pos_emb", apply_rotary_pos_emb
		169	+ )
161		170
162	pm.apply_patches()	171	pm.apply_patches()
163		172
	@@ -169,10 +178,10 @@ class Qwen3MoeExperts(nn.Module):
169	self.hidden_dim = config.hidden_size	178	self.hidden_dim = config.hidden_size
170	self.intermediate_size = config.moe_intermediate_size	179	self.intermediate_size = config.moe_intermediate_size
171	self.gate_up_proj = torch.nn.Parameter(	180	self.gate_up_proj = torch.nn.Parameter(
172	- torch.empty(self.num_experts * self.hidden_dim, 2 * self.intermediate_size))	181	+ torch.empty(self.num_experts * self.hidden_dim, 2 * self.intermediate_size)
		182	+ )
173		183
174	- self.down_proj = torch.nn.Parameter(	184	+ self.down_proj = torch.nn.Parameter(torch.empty(self.num_experts * self.intermediate_size, self.hidden_dim))
175	- torch.empty(self.num_experts * self.intermediate_size, self.hidden_dim))
176	self.act_fn = ACT2FN[config.hidden_act]	185	self.act_fn = ACT2FN[config.hidden_act]
177		186
178	def forward(self, hidden_states, routing_weights=None, selected_experts=None):	187	def forward(self, hidden_states, routing_weights=None, selected_experts=None):
	@@ -239,29 +248,28 @@ class Qwen3MoeSparseFusedMoeBlock(nn.Module):
239		248
240	args = get_args()	249	args = get_args()
241	if getattr(args, 'ep_size') > 1:	250	if getattr(args, 'ep_size') > 1:
242	- final_hidden_states = self.experts(	251	+ final_hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
243	- hidden_states, selected_experts, routing_weights
244	- )
245	else:	252	else:
246	-
247	final_hidden_states = self.experts(	253	final_hidden_states = self.experts(
248	hidden_states, routing_weights=routing_weights, selected_experts=selected_experts	254	hidden_states, routing_weights=routing_weights, selected_experts=selected_experts
249	)	255	)
250		256
251	final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)	257	final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
252	- return final_hidden_states, router_logits	258	+ if version("transformers") < "5.0.0":
		259	+ return final_hidden_states, router_logits
		260	+ return final_hidden_states
253		261
254		262
255	def qwen3_moe_model_forward(	263	def qwen3_moe_model_forward(
256	- self,	264	+ self,
257	- input_ids: Optional[torch.LongTensor] = None,	265	+ input_ids: Optional[torch.LongTensor] = None,
258	- attention_mask: Optional[torch.Tensor] = None,	266	+ attention_mask: Optional[torch.Tensor] = None,
259	- position_ids: Optional[torch.LongTensor] = None,	267	+ position_ids: Optional[torch.LongTensor] = None,
260	- past_key_values: Optional[Cache] = None,	268	+ past_key_values: Optional[Cache] = None,
261	- inputs_embeds: Optional[torch.FloatTensor] = None,	269	+ inputs_embeds: Optional[torch.FloatTensor] = None,
262	- use_cache: Optional[bool] = None,	270	+ use_cache: Optional[bool] = None,
263	- cache_position: Optional[torch.LongTensor] = None,	271	+ cache_position: Optional[torch.LongTensor] = None,
264	- **kwargs,	272	+ **kwargs,
265	) -> MoeModelOutputWithPast:	273	) -> MoeModelOutputWithPast:
266	if (input_ids is None) ^ (inputs_embeds is not None):	274	if (input_ids is None) ^ (inputs_embeds is not None):
267	raise ValueError("You must specify exactly one of input_ids or inputs_embeds")	275	raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
	@@ -300,11 +308,11 @@ def qwen3_moe_model_forward(
300	for layer_id, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):	308	for layer_id, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
301	if self.training:	309	if self.training:
302	with async_save_on_cpu(	310	with async_save_on_cpu(
303	- h2d_stream=offload_stream,	311	+ h2d_stream=offload_stream,
304	- d2h_stream=offload_stream,	312	+ d2h_stream=offload_stream,
305	- block_idx=int(layer_id),	313	+ block_idx=int(layer_id),
306	- depth=len(self.layers),	314	+ depth=len(self.layers),
307	- custom_check_fn=lambda x: x.data_ptr() == hidden_states.data_ptr()	315	+ custom_check_fn=lambda x: x.data_ptr() == hidden_states.data_ptr(), # pylint: disable=W0640
308	):	316	):
309	hidden_states = decoder_layer(	317	hidden_states = decoder_layer(
310	hidden_states,	318	hidden_states,