[bugfix] fix megatron lora TP all-reduce (modelscope#7911)

Author: Jintao-Huang

swift/megatron/model/gpt_bridge.py

@@ -469,8 +469,9 @@ def _set_state_dict(self,
             if to_mcore:
                 assert mg_param is not None, f'mg_module: {mg_module}, mg_key: {mg_key}'
                 hf_weight = hf_state_dict[hf_key].load()
-                if module_key in {'embedding.word_embeddings', 'output_layer'
-                                  } and hf_weight.shape[0] < self.args.padded_vocab_size:
+                if module_key in {
+                        'embedding.word_embeddings', 'output_layer'
+                } and hf_weight.shape[0] < self.args.padded_vocab_size and self.args.task_type != 'seq_cls':
                     hf_weight = F.pad(hf_weight, (0, 0, 0, self.args.padded_vocab_size - hf_weight.shape[0]))
                 hf_scale_inv = None
                 if f'{hf_key}_scale_inv' in hf_state_dict:
@@ -1295,10 +1296,10 @@ def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcor
         lm_model = getattr(mg_model, 'language_model') if self.args.is_multimodal else mg_model
         if self.args.task_type != 'embedding':
             if self.args.untie_embeddings_and_output_weights:
-                if not to_mcore or self.args.task_type in {'causal_lm', 'generative_reranker'}:
-                    hf_lm_head_key = self.hf_lm_head_key
-                    if self.args.task_type == 'seq_cls':
-                        hf_lm_head_key = self.hf_score_key
+                hf_lm_head_key = self.hf_lm_head_key
+                if self.args.task_type == 'seq_cls':
+                    hf_lm_head_key = self.hf_score_key
+                if not to_mcore or hf_lm_head_key in hf_state_dict:
                     self._set_state_dict(lm_model, 'output_layer.weight', hf_state_dict, hf_lm_head_key, to_mcore)
             elif to_mcore and lm_model.output_layer.weight is not None:
                 self._set_state_dict(lm_model, 'output_layer.weight', hf_state_dict, self.hf_embed_key, to_mcore)

swift/megatron/model/gpt_model.py

@@ -142,6 +142,7 @@ def __init__(
                 parallel_mode=None,
                 skip_weight_param_allocation=False,
             )
+            self.output_layer.weight.average_gradients_across_tp_domain = True
         elif args.task_type == 'embedding' and self.post_process:
             self.output_layer = None
 

swift/megatron/model/mm_gpt_model.py

@@ -5,7 +5,7 @@
 import torch
 from megatron.core import InferenceParams
 from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.tensor_parallel import VocabParallelEmbedding, scatter_to_sequence_parallel_region
+from megatron.core.tensor_parallel import VocabParallelEmbedding, reduce_scatter_to_sequence_parallel_region
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -70,7 +70,7 @@ def forward(_self, input_):
             if reduce_scatter_embeddings:
                 res = res.transpose(0, 1).contiguous()
                 group_kwargs = {'group': _self.tp_group} if mcore_013 else {}
-                res = scatter_to_sequence_parallel_region(res, **group_kwargs)
+                res = reduce_scatter_to_sequence_parallel_region(res, **group_kwargs) / args.tensor_model_parallel_size
             return res
 
         VocabParallelEmbedding.forward = forward

swift/megatron/tuners/lora.py

@@ -16,7 +16,6 @@
                                                          TERowParallelGroupedLinear, TERowParallelLinear)
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.parallel_state import get_expert_tensor_parallel_world_size, get_tensor_model_parallel_world_size
-from megatron.core.tensor_parallel import gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region
 from megatron.core.tensor_parallel.random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
 from megatron.core.transformer.mlp import apply_swiglu_sharded_factory
 from megatron.core.transformer.module import MegatronModule
@@ -201,8 +200,10 @@ def update_layer(self, adapter_name, r, *, lora_alpha, lora_dropout, init_lora_w
                     **kwargs,
                 )
                 lora_b.parallel_mode = self.base_layer.parallel_mode  # fix moe_shared_expert_overlap
-        lora_a.sequence_parallel = False
-        lora_b.sequence_parallel = False
+        for lora in [lora_a, lora_b]:
+            if getattr(lora, 'parallel_mode', None) is None and hasattr(lora, 'weight'):  # TODO: experts
+                sequence_parallel = True if isinstance(self.base_layer, TopKRouter) else self.sequence_parallel
+                lora.weight.sequence_parallel = sequence_parallel
         self.lora_A[adapter_name] = lora_a
         self.lora_B[adapter_name] = lora_b
         if hasattr(self, 'lora_bias'):
@@ -341,8 +342,6 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
         else:
             raise ValueError(f'Unsupported base layer type: {type(self.base_layer)}')
         if not isinstance(self.base_layer, TopKRouter) and not self.disable_adapters and not self.merged:
-            if self.sequence_parallel and self.base_layer.parallel_mode == 'column':
-                x = gather_from_sequence_parallel_region(x)
             for active_adapter in self.active_adapters:
                 if active_adapter not in self.lora_A.keys():
                     continue
@@ -362,8 +361,6 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
                 if isinstance(lora_result, tuple):
                     lora_result = lora_result[0]
                 lora_result = lora_result * scaling
-                if self.sequence_parallel and self.base_layer.parallel_mode == 'row':
-                    lora_result = scatter_to_sequence_parallel_region(lora_result)
                 result = result + lora_result
 
         result = result.to(previous_dtype)

swift/megatron/utils/convert_utils.py

@@ -1,8 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 
 import math
-import os
-import shutil
 from contextlib import contextmanager
 from typing import Any, Dict
 
@@ -25,7 +23,7 @@ def _test_params_sum(model):
     for n, p in model.named_parameters():
         n_parameter += 1
         sum_ = p.to(device='cuda', dtype=torch.float32).abs().sum().cpu().item()
-        if sum_ == 0:
+        if sum_ == 0 and '.lora_B.' not in n:
             zero_count += 1
             logger.warning(f'n: {n}, sum: {sum_}')
         elif math.isnan(sum_) or math.isinf(sum_) or sum_ > 1e10:
@@ -200,7 +198,7 @@ def test_convert_precision(hf_model, mg_model, template, torch_dtype=torch.float
     with torch.inference_mode(), _model_cpu_forward_context(
             mg_modules, torch_dtype, 'cuda', share_embedding=share_embedding, target_device=mg_device):
         mg_logits = forward_step_helper(mg_model, mg_inputs, dtype=torch_dtype)
-        if args.tensor_model_parallel_size > 1:
+        if args.tensor_model_parallel_size > 1 and args.task_type != 'seq_cls':
             from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
             if mg_logits is not None:
                 mg_logits = gather_from_tensor_model_parallel_region(mg_logits)

swift/megatron/utils/utils.py

@@ -15,6 +15,7 @@
 from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint, sharded_state_dict_default
 from megatron.training import checkpointing, get_args
 from packaging import version
+from peft.tuners.lora import Linear as LoraLinear
 from peft.utils.other import ModulesToSaveWrapper
 from torch import nn
 
@@ -156,6 +157,7 @@ def new_deepcopy(x, *args, **kwargs):
 
 
 def prepare_adapter(model):
+    from swift.megatron.tuners import LoraParallelLinear
     args = get_args()
     set_linear_is_expert(model)
     target_modules = get_target_modules(args, model)
@@ -179,6 +181,15 @@ def prepare_adapter(model):
         for n, p in model.named_parameters():
             if '.ref_adapter.' in n:
                 p.requires_grad = False
+    # setting average_gradients_across_tp_domain
+    for m in model.modules():
+        if isinstance(m, LoraLinear):
+            # just check
+            assert args.is_multimodal or args.hf_model_type == 'qwen3_next'
+            assert not isinstance(m, LoraParallelLinear)
+            for p in m.parameters():
+                if p.requires_grad:
+                    p.average_gradients_across_tp_domain = True
     return model