[Distributed] [CustomDevices] adapt lora sp && polish MC2 APIs

Author: SylarTiaNII

paddlenlp/peft/lora/lora_layers.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import math
-import os
 from typing import List, Optional
 
 import paddle
@@ -24,28 +23,22 @@
     ColumnParallelLinear,
     RowParallelLinear,
 )
-
-from paddlenlp.transformers.sequence_parallel_utils import (
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
     AllGatherOp,
     ColumnSequenceParallelLinear,
-    MC2ColumnSeqParallelLinear,
-    MC2RowSeqParallelLinear,
     ReduceScatterOp,
     RowSequenceParallelLinear,
     mark_as_sequence_parallel_parameter,
 )
 
-from .lora_quick_layers import quick_lora
-
-if "npu" in paddle.device.get_all_custom_device_type():
-    from .mc2_lora_npu import MC2LoRaColumnParallelLinear, MC2LoRaRowParallelLinear
-else:
-    MC2LoRaRowParallelLinear = None
-    MC2LoRaColumnParallelLinear = None
-
+from paddlenlp.transformers.mc2_parallel_linear import (
+    MC2ColumnParallelCoreLinear,
+    MC2ColumnSeqParallelCoreLinear,
+    MC2RowParallelCoreLinear,
+    MC2RowSeqParallelCoreLinear,
+)
 
-def is_mc2_valid():
-    return "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0"))
+from .lora_quick_layers import quick_lora
 
 
 class LoRALinear(nn.Linear):
@@ -280,16 +273,16 @@ def forward(self, x: paddle.Tensor):
             )
         else:
             # x @ W : [bz, in_f / ws] ===> [bz, out_f]
-            if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                output = MC2LoRaRowParallelLinear.apply(input_mp, self.weight, self.model_parallel_group)
-            else:
+            if MC2RowParallelCoreLinear is None:
                 result_mp = F.linear(x=input_mp, weight=self.weight, name=self.name)
                 output = mp_ops._mp_allreduce(
                     result_mp,
                     group=self.model_parallel_group,
                     use_calc_stream=True,
                     use_model_parallel=True,
                 )
+            else:
+                output = MC2RowParallelCoreLinear.apply(input_mp, self.weight, self.model_parallel_group)
 
             if not self.merged:
                 # x @ A: [bz, in_f/ ws] ===> [bz, r]
@@ -402,21 +395,21 @@ def forward(self, x: paddle.Tensor):
         else:
             input_mp = x
 
-        if not is_mc2_valid():
+        if MC2RowSeqParallelCoreLinear is None:
             output_parallel = self.linear(input_mp, self.weight, name=self._name)
             output_ = ReduceScatterOp.apply(output_parallel)
             result_mp = output_ + self.bias if self.bias is not None else output_
         else:
-            output_ = MC2RowSeqParallelLinear.apply(input_mp, self.weight, self.model_parallel_group)
+            output_ = MC2RowSeqParallelCoreLinear.apply(input_mp, self.weight, self.model_parallel_group)
             result_mp = output_ + self.bias if self.bias is not None else output_
 
         if not self.merged:
             input_mp = self.lora_dropout(input_mp)
-            if not is_mc2_valid():
+            if MC2RowSeqParallelCoreLinear is None:
                 input_mp = input_mp @ self.lora_A
                 input_mp = ReduceScatterOp.apply(input_mp)
             else:
-                input_mp = MC2RowSeqParallelLinear.apply(input_mp, self.lora_A, self.model_parallel_group)
+                input_mp = MC2RowSeqParallelCoreLinear.apply(input_mp, self.lora_A, self.model_parallel_group)
             delta_mp = (input_mp @ self.lora_B) * self.scaling
             result_mp += delta_mp
         return result_mp
@@ -528,21 +521,21 @@ def forward(self, input: paddle.Tensor):
                 world_size=self.world_size,
             )
         else:
-            if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                res_mp = MC2LoRaColumnParallelLinear.apply(input, self.weight, self.model_parallel_group)
-                result_mp = res_mp + self.bias
-            else:
+            if MC2ColumnParallelCoreLinear is None:
                 input_mp = mp_ops._c_identity(input, group=self.model_parallel_group)
                 result_mp = F.linear(x=input_mp, weight=self.weight, bias=self.bias, name=self.name)
+            else:
+                res_mp = MC2ColumnParallelCoreLinear.apply(input, self.weight, self.model_parallel_group)
+                result_mp = res_mp + self.bias
 
             if not self.merged:
                 input_a = self.lora_dropout(input) @ self.lora_A
-                if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                    tmp = MC2LoRaColumnParallelLinear.apply(input_a, self.lora_B, self.model_parallel_group)
-                    delta_mp = tmp * self.scaling
-                else:
+                if MC2ColumnParallelCoreLinear is None:
                     input_a_mp = mp_ops._c_identity(input_a, group=self.model_parallel_group)
                     delta_mp = (input_a_mp @ self.lora_B) * self.scaling
+                else:
+                    tmp = MC2ColumnParallelCoreLinear.apply(input_a, self.lora_B, self.model_parallel_group)
+                    delta_mp = tmp * self.scaling
                 result_mp += delta_mp
 
         if self.gather_output and self.is_mp:
@@ -641,24 +634,24 @@ def eval(self):
             self.merged = True
 
     def forward(self, x: paddle.Tensor):
-        if not is_mc2_valid():
+        if MC2ColumnSeqParallelCoreLinear is None:
             if self.is_mp:
                 input_parallel = AllGatherOp.apply(x)
             else:
                 input_parallel = x
             result_mp = self.linear(input_parallel, self.weight, self.bias, name=self._name)
         else:
-            result_mp = MC2ColumnSeqParallelLinear.apply(x, self.weight, self.model_parallel_group)
+            result_mp = MC2ColumnSeqParallelCoreLinear.apply(x, self.weight, self.model_parallel_group)
             if self.bias is not None:
                 result_mp += self.bias
 
         if not self.merged:
             input_a = self.lora_dropout(x) @ self.lora_A
-            if not is_mc2_valid():
+            if MC2ColumnSeqParallelCoreLinear is None:
                 input_a = AllGatherOp.apply(input_a)
                 delta_mp = (input_a @ self.lora_B) * self.scaling
             else:
-                input_a = MC2ColumnSeqParallelLinear.apply(input_a, self.lora_B, self.model_parallel_group)
+                input_a = MC2ColumnSeqParallelCoreLinear.apply(input_a, self.lora_B, self.model_parallel_group)
                 delta_mp = input_a * self.scaling
             result_mp += delta_mp
 

paddlenlp/peft/lora/lora_model.py

@@ -31,8 +31,7 @@
     PipelineLayer,
     RowParallelLinear,
 )
-
-from paddlenlp.transformers.sequence_parallel_utils import (
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
     ColumnSequenceParallelLinear,
     RowSequenceParallelLinear,
 )

paddlenlp/peft/lora/mc2_lora_npu.py

@@ -62,6 +62,10 @@ def swiglu(x, y=None):
     init_name_mappings,
 )
 from paddlenlp.transformers.long_sequence_strategies import LongSequenceStrategies
+from paddlenlp.transformers.mc2_parallel_linear import (
+    MC2ColumnSeqParallelLinear,
+    MC2RowSeqParallelLinear,
+)
 from paddlenlp.transformers.model_outputs import (
     BaseModelOutputWithPastAndCrossAttentions,
     CausalLMOutputWithCrossAttentions,
@@ -96,13 +100,6 @@ def swiglu(x, y=None):
 ]
 
 
-def is_mc2_valid():
-    current_device = get_env_device()
-    if current_device == "npu":
-        return True
-    return False
-
-
 def _get_interleave(n):
     def _get_interleave_power_of_2(n):
         start = 2 ** (-(2 ** -(math.log2(n) - 3)))
@@ -574,12 +571,7 @@ def __init__(self, config):
         self.fuse_attention_ffn = config.fuse_attention_ffn
 
         if config.sequence_parallel:
-            if is_mc2_valid and int(os.getenv("FLAGS_NPU_MC2", 0)):
-                from paddlenlp.transformers.mc2_seqence_parallel_linear import (
-                    MC2ColumnSeqParallelLinear,
-                    MC2RowSeqParallelLinear,
-                )
-
+            if MC2ColumnSeqParallelLinear is not None and MC2RowSeqParallelLinear is not None:
                 ColumnParallelLinear = MC2ColumnSeqParallelLinear
                 RowParallelLinear = MC2RowSeqParallelLinear
             else:
@@ -697,12 +689,7 @@ def __init__(self, config: LlamaConfig, layerwise_recompute: bool = False):
                 self.use_fused_rope = False
 
         if config.sequence_parallel:
-            if is_mc2_valid and int(os.getenv("FLAGS_NPU_MC2", 0)):
-                from paddlenlp.transformers.mc2_seqence_parallel_linear import (
-                    MC2ColumnSeqParallelLinear,
-                    MC2RowSeqParallelLinear,
-                )
-
+            if MC2ColumnSeqParallelLinear is not None and MC2RowSeqParallelLinear is not None:
                 ColumnParallelLinear = MC2ColumnSeqParallelLinear
                 RowParallelLinear = MC2RowSeqParallelLinear
             else: