compat qwen3_vl zero3 (#6080)

Author: Jintao-Huang

examples/models/qwen3_vl/zero2.sh renamed to examples/models/qwen3_vl/zero3.sh

@@ -1,4 +1,4 @@
-# zero2: 70GiB
+# 2 * 42GiB
 IMAGE_MAX_TOKEN_NUM=1024 \
 NPROC_PER_NODE=2 \
 CUDA_VISIBLE_DEVICES=0,1 \
@@ -30,7 +30,7 @@ swift sft \
     --max_length 2048 \
     --output_dir output \
     --warmup_ratio 0.05 \
-    --deepspeed zero2 \
+    --deepspeed zero3 \
     --use_liger_kernel true \
     --dataset_num_proc 4 \
     --dataloader_num_workers 4

swift/llm/model/register.py

@@ -208,6 +208,45 @@ def _new_process_model_before_weight_loading(self, model, *args, **kwargs):
         pass
 
 
+def deepspeed_set_z3_leaf_modules(model):
+    if not is_deepspeed_zero3_enabled():
+        return
+    try:
+        architecture = model.config.architectures[0]
+    except Exception:
+        return
+    z3_leaf_modules = None
+    if architecture == 'Qwen3VLMoeForConditionalGeneration':
+        from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock
+        z3_leaf_modules = [Qwen3VLMoeTextSparseMoeBlock]
+    elif architecture == 'Qwen3OmniMoeForConditionalGeneration':
+        from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeThinkerTextSparseMoeBlock
+        z3_leaf_modules = [Qwen3OmniMoeThinkerTextSparseMoeBlock]
+    elif architecture == 'Qwen2MoeForCausalLM':
+        from transformers.models.qwen2_moe.modeling_qwen2_moe import Qwen2MoeSparseMoeBlock
+        z3_leaf_modules = [Qwen2MoeSparseMoeBlock]
+    elif architecture == 'Qwen3MoeForCausalLM':
+        from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock
+        z3_leaf_modules = [Qwen3MoeSparseMoeBlock]
+    elif architecture == 'Glm4MoeForCausalLM':
+        from transformers.models.glm4_moe.modeling_glm4_moe import Glm4MoeMoE
+        z3_leaf_modules = [Glm4MoeMoE]
+    elif architecture == 'Glm4vMoeForConditionalGeneration':
+        from transformers.models.glm4v_moe.modeling_glm4v_moe import Glm4vMoeTextMoE
+        z3_leaf_modules = [Glm4vMoeTextMoE]
+    elif architecture == 'GptOssForCausalLM':
+        from transformers.models.gpt_oss.modeling_gpt_oss import GptOssMLP
+        z3_leaf_modules = [GptOssMLP]
+    elif architecture == 'Llama4ForCausalLM':
+        from transformers.models.llama4.modeling_llama4 import Llama4TextMoe
+        z3_leaf_modules = [Llama4TextMoe]
+
+    if z3_leaf_modules:
+        from deepspeed.utils import set_z3_leaf_modules
+        set_z3_leaf_modules(model, z3_leaf_modules)
+        logger.info(f'Setting z3_leaf_modules: {z3_leaf_modules}')
+
+
 def get_model_tokenizer_from_local(model_dir: str,
                                    model_info: ModelInfo,
                                    model_kwargs: Dict[str, Any],
@@ -329,6 +368,8 @@ def get_model_tokenizer_from_local(model_dir: str,
     if model is not None:
         # fix seq classification task
         HfConfigFactory.set_model_config_attr(model, 'pad_token_id', pad_token)
+        # deepspeed zero3
+        deepspeed_set_z3_leaf_modules(model)
 
     return model, tokenizer