[bugfix] fix train_type full freeze_llm (#6651)

Author: Jintao-Huang

swift/llm/model/model_arch.py

@@ -337,7 +337,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
     register_model_arch(
         MultiModelKeys(
             MLLMModelArch.llava_hf,
-            language_model='model.language_model',
+            language_model=['model.language_model', 'lm_head'],
             aligner='model.multi_modal_projector',
             vision_tower='model.vision_tower',
         ))
@@ -362,7 +362,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
     register_model_arch(
         MultiModelKeys(
             MLLMModelArch.llava_next_video_hf,
-            language_model='model.language_model',
+            language_model=['model.language_model', 'lm_head'],
             aligner=['model.multi_modal_projector'],
             vision_tower='model.vision_tower'))
 else:
@@ -400,7 +400,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.interns1,
-        language_model='model.language_model',
+        language_model=['model.language_model', 'lm_head'],
         aligner='model.multi_modal_projector',
         vision_tower='model.vision_tower',
     ))
@@ -521,31 +521,31 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
     register_model_arch(
         MultiModelKeys(
             MLLMModelArch.qwen2_vl,
-            language_model='model.language_model',
+            language_model=['model.language_model', 'lm_head'],
             aligner='model.visual.merger',
             vision_tower='model.visual',
         ))
 else:
     register_model_arch(
         MultiModelKeys(
             MLLMModelArch.qwen2_vl,
-            language_model='model',
+            language_model=['model', 'lm_head'],
             aligner='visual.merger',
             vision_tower='visual',
         ))
 
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.qwen3_vl,
-        language_model='model.language_model',
+        language_model=['model.language_model', 'lm_head'],
         aligner=['model.visual.merger', 'model.visual.deepstack_merger_list'],
         vision_tower='model.visual',
     ))
 
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.qwen2_5_omni,
-        language_model='thinker.model',
+        language_model=['thinker.model', 'thinker.lm_head'],
         vision_tower=['thinker.audio_tower', 'thinker.visual'],
         aligner=['thinker.audio_tower.proj', 'thinker.visual.merger'],
         generator=['talker', 'token2wav'],
@@ -554,7 +554,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.qwen3_omni,
-        language_model='thinker.model',
+        language_model=['thinker.model', 'thinker.lm_head'],
         vision_tower=['thinker.audio_tower', 'thinker.visual'],
         aligner=[
             'thinker.audio_tower.proj1', 'thinker.audio_tower.proj2', 'thinker.visual.merger',
@@ -574,7 +574,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.step_audio2_mini,
-        language_model='model',
+        language_model=['model', 'lm_head'],
         aligner=['adapter'],
         vision_tower=['encoder'],
     ))
@@ -589,7 +589,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.glm4_1v,
-        language_model='model.language_model',
+        language_model=['model.language_model', 'lm_head'],
         aligner='model.visual.merger',
         vision_tower='model.visual',
     ))
@@ -622,7 +622,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.ernie_vl,
-        language_model='model',
+        language_model=['model', 'lm_head'],
         aligner='model.resampler_model',
         vision_tower='vision_model',
     ))
@@ -631,7 +631,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
     register_model_arch(
         MultiModelKeys(
             MLLMModelArch.llama3_2_vision,
-            language_model='model.language_model',
+            language_model=['model.language_model', 'lm_head'],
             aligner='model.multi_modal_projector',
             vision_tower='model.vision_model',
         ))
@@ -696,15 +696,15 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.gemma3n,
-        language_model='model.language_model',
+        language_model=['model.language_model', 'lm_head'],
         aligner=['model.embed_vision', 'model.embed_audio'],
         vision_tower=['model.vision_tower', 'model.audio_tower'],
     ))
 
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.keye_vl,
-        language_model='model',
+        language_model=['model', 'lm_head'],
         aligner='mlp_AR',
         vision_tower='visual',
     ))
@@ -717,7 +717,7 @@ def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> Non
 register_model_arch(
     MultiModelKeys(
         MLLMModelArch.llava_onevision1_5,
-        language_model='model.language_model',
+        language_model=['model.language_model', 'lm_head'],
         aligner='model.visual.merger',
         vision_tower='model.visual',
     ))

swift/llm/train/tuner.py

@@ -111,13 +111,16 @@ def get_multimodal_target_regex(
     res = []
     for module in modules:
         rejected_modules = []
-        if not freeze_vit:
+        if not freeze_vit or not freeze_llm:
             for aligner in model_arch.aligner:
                 if aligner.startswith(f'{module}.'):
                     rejected_modules.append(aligner)
 
         sub_module = deep_getattr(model, module)
-        target_modules = find_all_linears(sub_module, model_arch, extra_layers)
+        if isinstance(sub_module, nn.Linear) and module.endswith('lm_head'):
+            target_modules = []
+        else:
+            target_modules = find_all_linears(sub_module, model_arch, extra_layers)
         if exclude_router and model.model_info.is_moe_model:
             target_modules = [tm for tm in target_modules if tm not in {'gate'}]
         if not target_modules:

swift/megatron/trainers/utils.py

@@ -111,9 +111,8 @@ def get_batch_on_this_cp_rank(batch: Dict[str, Any]):
     if cp_size > 1:
         args = get_args()
         keys = ['labels', 'attention_mask', 'position_ids', 'loss_scale']
-        if args.is_multimodal:
-            keys.append('decoder_input')
-        else:
+        if not args.is_multimodal:
+            # Multimodal models will handle CP in input_embeds.
             keys.append('input_ids')
 
         packed_seq_params = batch.get('packed_seq_params')