Enable LoRA for TELinear layers (#13929)

* Enable LoRA for TELinear layers

Signed-off-by: Chen Cui <chcui@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: cuichenx <cuichenx@users.noreply.github.com>

---------

Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: cuichenx <cuichenx@users.noreply.github.com>
Co-authored-by: cuichenx <cuichenx@users.noreply.github.com>

Author: cuichenx

nemo/collections/llm/peft/canonical_lora.py

@@ -230,7 +230,9 @@ def transform(self, m: nn.Module, name=None, prefix=None):
                     m, dim=self.dim, alpha=self.alpha, dropout=self.dropout, lora_A_init_method=self.lora_A_init_method
                 )
 
-            input_is_parallel, in_features, out_features, disable_sp_comm = get_adapter_attributes_from_linear(m)
+            input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = (
+                get_adapter_attributes_from_linear(m)
+            )
 
             adapter_kwargs = dict(
                 dim=self.dim,
@@ -247,6 +249,7 @@ def transform(self, m: nn.Module, name=None, prefix=None):
                 alpha=self.alpha,
                 is_expert=is_expert_linear(full_name),
                 disable_sequence_parallel_comm=disable_sp_comm,
+                base_linear_is_parallel=base_linear_is_parallel,
             )
             if name in ['linear_proj', 'linear_fc2']:
                 adapter = ParallelLinearAdapter(in_features, out_features, **adapter_kwargs)

nemo/collections/llm/peft/dora.py

@@ -180,7 +180,9 @@ def transform(self, m: nn.Module, name=None, prefix=None):
         """
         if (ans := self.match(m, name, prefix)) is not None:
             (match, full_name) = ans
-            input_is_parallel, in_features, out_features, disable_sp_comm = get_adapter_attributes_from_linear(m)
+            input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = (
+                get_adapter_attributes_from_linear(m)
+            )
             logging.info(f"Adding DoRA to: {full_name}")
             adapter = ParallelLinearDoRAAdapter(
                 in_features,
@@ -198,6 +200,7 @@ def transform(self, m: nn.Module, name=None, prefix=None):
                 model_parallel_config=getattr(m, "config", None),
                 alpha=self.alpha,
                 disable_sequence_parallel_comm=disable_sp_comm,
+                base_linear_is_parallel=base_linear_is_parallel,
             )
             return DoRALinear(m, adapter)
         return m

nemo/collections/llm/peft/lora.py

@@ -462,7 +462,9 @@ def transform(self, m: nn.Module, name=None, prefix=None):
                     lora_dtype=self.lora_dtype,
                 )
 
-            input_is_parallel, in_features, out_features, disable_sp_comm = get_adapter_attributes_from_linear(m)
+            input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = (
+                get_adapter_attributes_from_linear(m)
+            )
             logging.info(f"Adding lora to: {full_name}")
             adapter = ParallelLinearAdapter(
                 in_features,
@@ -483,6 +485,7 @@ def transform(self, m: nn.Module, name=None, prefix=None):
                 a2a_experimental=self.a2a_experimental,
                 disable_sequence_parallel_comm=disable_sp_comm,
                 dropout_recompute=self.dropout_recompute,
+                base_linear_is_parallel=base_linear_is_parallel,
             )
             return LoRALinear(m, adapter)
         return m

nemo/collections/llm/peft/utils.py

@@ -73,7 +73,7 @@ def get_adapter_attributes_from_linear(m: nn.Module):
     Return input_is_parallel, in_features, out_feature attributes based on implementation of the base layer.
     """
     disable_sequence_parallel_comm = not m.config.sequence_parallel
-
+    base_linear_is_parallel = True
     if HAVE_TE and any(isinstance(m, te_column_parallel) for te_column_parallel in TECL):
         input_is_parallel = False
         # m.in_features and m.out_features are divided by tp_size already,
@@ -112,6 +112,7 @@ def get_adapter_attributes_from_linear(m: nn.Module):
         input_is_parallel = False
         in_features = m.in_features
         out_features = m.out_features
+        base_linear_is_parallel = False
     elif isinstance(m, ColumnParallelLinear):
         input_is_parallel = False
         in_features = m.input_size
@@ -123,7 +124,7 @@ def get_adapter_attributes_from_linear(m: nn.Module):
     else:
         raise NotImplementedError(f"Layer type is unrecognized for LoRA: {type(m)}")
 
-    return input_is_parallel, in_features, out_features, disable_sequence_parallel_comm
+    return input_is_parallel, in_features, out_features, disable_sequence_parallel_comm, base_linear_is_parallel
 
 
 def is_expert_linear(fqn):
@@ -262,6 +263,7 @@ def __init__(
         is_expert: bool = False,
         disable_sequence_parallel_comm: bool = True,
         dropout_recompute: bool = False,
+        base_linear_is_parallel: bool = True,
         **kwargs,
     ):
         super().__init__()
@@ -310,6 +312,10 @@ def __init__(
         lin_out_gather_output = True if input_is_parallel else False
         if self.use_a2a and input_is_parallel and _sequence_parallel:
             lin_out_gather_output = False
+
+        if not base_linear_is_parallel:
+            lin_out_gather_output = True
+
         self.linear_out = ColumnParallelLinear(
             dim,
             out_features,
@@ -344,6 +350,9 @@ def __init__(
         if not _sequence_parallel:
             self.disable_sequence_parallel_comm = True
 
+        if not base_linear_is_parallel:
+            self.disable_sequence_parallel_comm = True
+
     def _get_init_fn(self, init_method: str):
         if init_method == 'xavier':
             init_fn = nn.init.xavier_normal_