[megatron] Fix SP & LoRA (#5704)

Author: Jintao-Huang

swift/llm/model/utils.py

@@ -324,10 +324,10 @@ def git_clone_github(github_url: str,
         local_repo_name = github_url.rsplit('/', 1)[1]
     github_url = f'{github_url}.git'
     local_repo_path = os.path.join(git_cache_dir, local_repo_name)
-    with safe_ddp_context(None, use_barrier=True):
-        if not is_local_master():
-            return local_repo_path
+    with safe_ddp_context('git_clone', use_barrier=True):
         repo_existed = os.path.exists(local_repo_path)
+        if not is_local_master() and repo_existed:
+            return local_repo_path
         if repo_existed:
             command = ['git', '-C', local_repo_path, 'fetch']
             subprocess_run(command)

swift/megatron/tuners/lora.py

@@ -15,6 +15,7 @@
                                                          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.transformer.mlp import apply_swiglu_sharded_factory
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.router import TopKRouter
@@ -58,6 +59,7 @@ def __init__(
         self.fan_in_fan_out = fan_in_fan_out
         self._active_adapter = adapter_name
         self.is_expert = getattr(base_layer, 'is_expert', False)
+        self.sequence_parallel = getattr(base_layer, 'sequence_parallel', False)
         if self.is_expert:
             self.tp_size = get_expert_tensor_parallel_world_size()
         else:
@@ -189,6 +191,8 @@ def update_layer(self, adapter_name, r, *, lora_alpha, lora_dropout, init_lora_w
                 lora.ub_overlap_ag_dgrad = False
                 lora.ub_overlap_ag_fprop = False
                 lora.ub_overlap_rs_dgrad = False
+        lora_a.sequence_parallel = False
+        lora_b.sequence_parallel = False
         self.lora_A[adapter_name] = lora_a
         self.lora_B[adapter_name] = lora_b
         if hasattr(self, 'lora_bias'):
@@ -287,6 +291,8 @@ 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
@@ -306,7 +312,8 @@ 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)