From 02cf85e82a16ac8d1a047306d74ae463bbed8b2f Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 4 Aug 2025 17:31:38 +0000 Subject: [PATCH 01/24] Initial implementation of fused LoRA Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 225 +++++++++++++++++++++++++++++- 1 file changed, 223 insertions(+), 2 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index d48876dd2f55..d50d79cd0141 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass, field -from typing import List, Literal +from typing import List, Literal, Tuple import torch @@ -47,13 +47,234 @@ class LoRALinear(AdapterWrapper): class to provide a specific implementation of the forward method. """ - def forward(self, x, *args, **kwargs): + def forward(self, x: torch.Tensor, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: # pylint: disable=C0115,C0116 + if getattr(self, "_enable_fused_impl", False): + return self._fused_forward(x) linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs) adapter_output = self.adapter(layernorm_output.contiguous()) adapter_output = adapter_output.reshape(linear_output.shape) return linear_output + adapter_output, bias + def _fused_forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass with Transformer Engine operation fuser + + The fused implementation is a PyTorch module that shares + params with this module. Since it owns no state, there is no + need for extra checkpointing logic. + + """ + + # Construct fused impl if needed + fused_impl = getattr(self, "_fused_impl", (None,))[0] + if fused_impl is None: + if not HAVE_TE: + raise RuntimeError("Fused LoRALinear implementation requires Transformer Engine") + fused_impl = TEFusedLoRALinear.make_from_lora_linear(self) + self._fused_impl = (fused_impl,) # Wrap in tuple to avoid registering submodule + + # Apply fused impl + return fused_impl(x) + + +if HAVE_TE: + + class TEFusedLoRALinear(nn.Module): + """A LoRA adapter wrapper using Transformer Engine operation fuser + + Its compute is equivalent to LoRALinear. + + There are no guarantees on the checkpoint structure, either + for compatibility with other modules or for backward + compatibility. LoRALinear works around this by treating + TEFusedLoRALinear as stateless. + + """ + + def __init__( + self, + in_features: int, + out_features: int, + lora_dim: int, + *, + bias: bool = True, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + norm_type: Optional[str] = None, + norm_eps: float = 1e-5, + norm_zero_centered_gamma: bool = False, + lora_dropout: float = 0.0, + lora_dropout_position: str = "post", + lora_scale: float = 1.0, + ) -> None: + super().__init__() + + # Split adapter into two branches + # Main branch: norm, fork, linear + # LoRA branch: lora_a, lora_b, add + main_branch = [] + lora_branch = [] + + # Norm op + self.norm_main_branch_idx: Optional[int] = None + if norm_type is not None: + self.norm_main_branch_idx = len(main_branch) + norm_kwargs = { + "eps": norm_eps, + "device": device, + "dtype": dtype, + "zero_centered_gamma": norm_zero_centered_gamma, + } + if norm_type == "LayerNorm": + main_branch.append(te.ops.LayerNorm(in_features, **norm_kwargs)) + elif norm_type == "RMSNorm": + main_branch.append(te.ops.RMSNorm(in_features, **norm_kwargs)) + else: + raise ValueError(f"Unsupported normalization ({norm_type})") + main_branch.append(te.ops.Quantize(forward=True, backward=False)) + + # Fork to LoRA branch + main_branch.append(te.ops.MakeExtraOutput()) + + # Main branch linear op + self.linear_main_branch_idx: int = len(main_branch) + main_branch.append( + te.ops.Linear( + in_features, + out_features, + bias=bias, + device=device, + dtype=dtype, + tensor_parallel_mode=None, ### TODO Support TP + ) + ) + + # LoRA pre-processing + if lora_dropout > 0 and lora_dropout_position == "pre": + lora_branch.append(te.ops.Dropout(lora_dropout)) + + # LoRA linear ops + self.lora_a_lora_branch_idx: int = len(lora_branch) + lora_branch.append( + te.ops.Linear( + in_features, + lora_dim, + bias=False, + device=device, + dtype=dtype, + tensor_parallel_mode=None, ### TODO Support TP + ) + ) + self.lora_b_lora_branch_idx: int = len(lora_branch) + lora_branch.append( + te.ops.Linear( + lora_dim, + out_features, + bias=False, + device=device, + dtype=dtype, + tensor_parallel_mode=None, ### TODO Support TP + ) + ) + + # LoRA post-processing + if lora_scale != 1: + lora_branch.append(te.ops.ConstantScale(lora_scale)) + if lora_dropout > 0 and lora_dropout_position == "post": + lora_branch.append(te.ops.Dropout(lora_dropout)) + + # Add with main branch + lora_branch.append(te.ops.AddExtraInput()) + + # Fuse ops in each branch + self.main_branch = te.ops.Sequential(*main_branch) + self.lora_branch = te.ops.Sequential(*lora_branch) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + linear_output, linear_input = self.main_branch(x) + with te.fp8_autocast(enabled=False): + out = self.lora_branch(linear_input, linear_output) + return out, None + + @staticmethod + def make_from_lora_linear(lora_linear: LoRALinear) -> TEFusedLoRALinear: + """Construct a fused LoRA adapter with the same params as an unfused adapter""" + + # Check inputs + if not isinstance(lora_linear, LoRALinear): + raise ValueError(f"Expected LoRALinear, got {lora_linear.__class__}") + if not isinstance(lora_linear.to_wrap, (te.Linear, te.LayerNormLinear)): + raise ValueError( + f"Expected LoRALinear to wrap a Transformer Engine linear module, but found {lora_linear.to_wrap.__class__}" + ) + + # Args for TEFusedLoRALinear constructor + constructor_kwargs = {"device": "meta"} # Do not initialize params + + # Extract linear params from base linear module + linear = lora_linear.to_wrap + weight = orig_linear.weight + bias = orig_linear.bias + if isinstance(bias, torch.Tensor) and bias.numel() == 0: + bias = None + constructor_kwargs["in_features"] = weight.size(1) + constructor_kwargs["out_features"] = weight.size(0) + constructor_kwargs["bias"] = bias is not None + constructor_kwargs["dtype"] = weight.dtype + + # Extract norm params from base linear module + norm_type = None + norm_weight = None + norm_bias = None + if isinstance(orig_linear, te.LayerNormLinear): + norm_type = orig_linear.normalization + if norm_type == "LayerNorm": + norm_weight = orig_linear.layer_norm_weight + norm_bias = orig_linear.layer_norm_bias + elif norm_type == "RMSNorm": + norm_weight = orig_linear.layer_norm_weight + else: + raise RuntimeError("LayerNormLinear has unsupported norm type ({norm_type})") + constructor_kwargs["norm_type"] = norm_type + constructor_kwargs["norm_eps"] = orig_linear.eps + constructor_kwargs["norm_zero_centered_gamma"] = orig_linear.zero_centered_gamma + + # Extract params from LoRA adapter + adapter = lora_linear.adapter + lora_a_weight = None + lora_b_weight = None + if isinstance(adapter, (LinearAdapter, TELinearAdapter)): + lora_a_weight = adapter.lora_a.weight + lora_b_weight = adapter.lora_b.weight + constructor_kwargs["lora_dim"] = lora_a_weight.size(0) + constructor_kwargs["lora_dropout"] = adapter.dropout.p + constructor_kwargs["lora_dropout_position"] = adapter.dropout_position + constructor_kwargs["lora_scale"] = adapter.scale + elif isinstance(adapter, ParallelLinearAdapter): + lora_a_weight = adapter.linear_in.weight + lora_b_weight = adapter.linear_out.weight + constructor_kwargs["lora_dim"] = lora_a_weight.size(0) + constructor_kwargs["lora_dropout"] = adapter.dropout.p + constructor_kwargs["lora_dropout_position"] = adapter.dropout_position + constructor_kwargs["lora_scale"] = adapter.alpha / adapter.dim + + # Construct fused module + out = TEFusedLoRALinear(**constructor_kwargs) + + # Replace fused module params + if norm_type is not None: + norm_op = out.main_branch[out.norm_main_branch_idx] + norm_op.weight = norm_weight + if norm_bias: + norm_op.bias = norm_bias + linear_op = out.main_branch[out.linear_main_branch_idx] + linear_op.weight = weight + linear_op.bias = bias + out.lora_branch[out.lora_a_lora_branch_idx].weight = lora_a_weight + out.lora_branch[out.lora_b_lora_branch_idx].weight = lora_b_weight + + return out + if HAVE_TE: From 3afeedf6c8a859808c22ecf341cf69ab8a059374 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 4 Aug 2025 21:05:39 +0000 Subject: [PATCH 02/24] Get fused LoRA to run Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 124 ++++++++++++++++++++++-------- 1 file changed, 91 insertions(+), 33 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index d50d79cd0141..8835cfeec06a 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math from dataclasses import dataclass, field -from typing import List, Literal, Tuple +from typing import Literal, Optional import torch @@ -47,7 +49,7 @@ class LoRALinear(AdapterWrapper): class to provide a specific implementation of the forward method. """ - def forward(self, x: torch.Tensor, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, x: torch.Tensor, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: # pylint: disable=C0115,C0116 if getattr(self, "_enable_fused_impl", False): return self._fused_forward(x) @@ -56,7 +58,7 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> Tuple[torch.Tensor, torch adapter_output = adapter_output.reshape(linear_output.shape) return linear_output + adapter_output, bias - def _fused_forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Forward pass with Transformer Engine operation fuser The fused implementation is a PyTorch module that shares @@ -109,16 +111,48 @@ def __init__( ) -> None: super().__init__() - # Split adapter into two branches - # Main branch: norm, fork, linear - # LoRA branch: lora_a, lora_b, add - main_branch = [] - lora_branch = [] + self._make_main_branch( + in_features, + out_features, + bias=bias, + device=device, + dtype=dtype, + norm_type=norm_type, + norm_eps=norm_eps, + norm_zero_centered_gamma=norm_zero_centered_gamma, + ) + with te.fp8_model_init(enabled=False): + self._make_lora_branch( + in_features, + out_features, + lora_dim, + device=device, + dtype=dtype, + lora_dropout=lora_dropout, + lora_dropout_position=lora_dropout_position, + lora_scale=lora_scale, + ) + + def _make_main_branch( + self, + in_features: int, + out_features: int, + *, + bias: bool, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + norm_type: Optional[str], + norm_eps: float, + norm_zero_centered_gamma: bool, + ) -> None: + + # List of ops + ops = [] # Norm op self.norm_main_branch_idx: Optional[int] = None if norm_type is not None: - self.norm_main_branch_idx = len(main_branch) + self.norm_main_branch_idx = len(ops) norm_kwargs = { "eps": norm_eps, "device": device, @@ -126,19 +160,19 @@ def __init__( "zero_centered_gamma": norm_zero_centered_gamma, } if norm_type == "LayerNorm": - main_branch.append(te.ops.LayerNorm(in_features, **norm_kwargs)) + ops.append(te.ops.LayerNorm(in_features, **norm_kwargs)) elif norm_type == "RMSNorm": - main_branch.append(te.ops.RMSNorm(in_features, **norm_kwargs)) + ops.append(te.ops.RMSNorm(in_features, **norm_kwargs)) else: raise ValueError(f"Unsupported normalization ({norm_type})") - main_branch.append(te.ops.Quantize(forward=True, backward=False)) + ops.append(te.ops.Quantize(forward=True, backward=False)) # Fork to LoRA branch - main_branch.append(te.ops.MakeExtraOutput()) + ops.append(te.ops.MakeExtraOutput()) # Main branch linear op - self.linear_main_branch_idx: int = len(main_branch) - main_branch.append( + self.linear_main_branch_idx: int = len(ops) + ops.append( te.ops.Linear( in_features, out_features, @@ -149,13 +183,32 @@ def __init__( ) ) + # Fuse ops + self.main_branch = te.ops.Sequential(*ops) + + def _make_lora_branch( + self, + in_features: int, + out_features: int, + lora_dim: int, + *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + lora_dropout: float, + lora_dropout_position: str, + lora_scale: float, + ) -> None: + + # List of ops + ops = [] + # LoRA pre-processing if lora_dropout > 0 and lora_dropout_position == "pre": - lora_branch.append(te.ops.Dropout(lora_dropout)) + ops.append(te.ops.Dropout(lora_dropout)) # LoRA linear ops - self.lora_a_lora_branch_idx: int = len(lora_branch) - lora_branch.append( + self.lora_a_lora_branch_idx: int = len(ops) + ops.append( te.ops.Linear( in_features, lora_dim, @@ -165,8 +218,8 @@ def __init__( tensor_parallel_mode=None, ### TODO Support TP ) ) - self.lora_b_lora_branch_idx: int = len(lora_branch) - lora_branch.append( + self.lora_b_lora_branch_idx: int = len(ops) + ops.append( te.ops.Linear( lora_dim, out_features, @@ -179,18 +232,17 @@ def __init__( # LoRA post-processing if lora_scale != 1: - lora_branch.append(te.ops.ConstantScale(lora_scale)) + ops.append(te.ops.ConstantScale(lora_scale)) if lora_dropout > 0 and lora_dropout_position == "post": - lora_branch.append(te.ops.Dropout(lora_dropout)) + ops.append(te.ops.Dropout(lora_dropout)) # Add with main branch - lora_branch.append(te.ops.AddExtraInput()) + ops.append(te.ops.AddExtraInput()) - # Fuse ops in each branch - self.main_branch = te.ops.Sequential(*main_branch) - self.lora_branch = te.ops.Sequential(*lora_branch) + # Fuse ops + self.lora_branch = te.ops.Sequential(*ops) - def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: linear_output, linear_input = self.main_branch(x) with te.fp8_autocast(enabled=False): out = self.lora_branch(linear_input, linear_output) @@ -200,19 +252,25 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: def make_from_lora_linear(lora_linear: LoRALinear) -> TEFusedLoRALinear: """Construct a fused LoRA adapter with the same params as an unfused adapter""" + from nemo.collections.llm.peft.utils import ParallelLinearAdapter + # Check inputs if not isinstance(lora_linear, LoRALinear): raise ValueError(f"Expected LoRALinear, got {lora_linear.__class__}") - if not isinstance(lora_linear.to_wrap, (te.Linear, te.LayerNormLinear)): + if not isinstance(lora_linear.to_wrap, (te.Linear, te.LayerNormLinear, torch.nn.Linear)): + raise ValueError( + f"Unsupported class for LoRALinear wrapped linear ({lora_linear.to_wrap.__class__})" + ) + if not isinstance(lora_linear.adapter, (LinearAdapter, TELinearAdapter, ParallelLinearAdapter)): raise ValueError( - f"Expected LoRALinear to wrap a Transformer Engine linear module, but found {lora_linear.to_wrap.__class__}" + f"Unsupported class for LoRALinear adapter ({lora_linear.adapter.__class__})" ) # Args for TEFusedLoRALinear constructor constructor_kwargs = {"device": "meta"} # Do not initialize params # Extract linear params from base linear module - linear = lora_linear.to_wrap + orig_linear = lora_linear.to_wrap weight = orig_linear.weight bias = orig_linear.bias if isinstance(bias, torch.Tensor) and bias.numel() == 0: @@ -585,7 +643,7 @@ class LoRA(PEFT, ModuleMatcher): This class facilitates the application of LoRA to specific modules within the model architecture. Args: - target_modules (List[str], optional): A list of module names to apply LoRA to. + target_modules (list[str], optional): A list of module names to apply LoRA to. Defaults to all linear layers ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2']. - 'linear_qkv': Apply LoRA to the fused linear layer used for query, key, and value projections in self-attention. @@ -595,7 +653,7 @@ class LoRA(PEFT, ModuleMatcher): Target modules can also contain wildcards. For example, you can specify target_modules=['*.layers.0.*.linear_qkv', '*.layers.1.*.linear_qkv'] to add LoRA to only linear_qkv on the first two layers. - exclude_modules (List[str], optional): A list of module names not to apply LoRa to. It will + exclude_modules (list[str], optional): A list of module names not to apply LoRa to. It will match all nn.Linear & nn.Linear-adjacent modules whose name does not match any string in exclude_modules. If used, will require target_modules to be empty list or None. dim (int): Dimension of the low-rank projection space. Defaults to 32. @@ -626,7 +684,7 @@ class LoRA(PEFT, ModuleMatcher): ) """ - target_modules: List[str] = field( + target_modules: list[str] = field( default_factory=lambda: ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2'] ) dim: int = 32 From c651d90e9d7c9f65654bb63dc98427d27f741f1f Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 5 Aug 2025 22:26:05 +0000 Subject: [PATCH 03/24] Initial work toward tensor-parallel support Missing all-gather op Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 107 ++++++++++++++++++++++++++---- 1 file changed, 95 insertions(+), 12 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 8835cfeec06a..a85b9cc9fac1 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -19,6 +19,7 @@ from typing import Literal, Optional import torch +from megatron.core import parallel_state from nemo.utils.import_utils import safe_import @@ -49,16 +50,27 @@ class LoRALinear(AdapterWrapper): class to provide a specific implementation of the forward method. """ - def forward(self, x: torch.Tensor, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: + def forward( + self, + x: torch.Tensor, + *args, + **kwargs, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: # pylint: disable=C0115,C0116 - if getattr(self, "_enable_fused_impl", False): - return self._fused_forward(x) + + # Fused implementation + if ( + getattr(self, "_enable_fused_impl", False) + and parallel_state.get_tensor_model_parallel_world_size() == 1 + ): # TP is not yet supported + return self._fused_forward(x) + linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs) adapter_output = self.adapter(layernorm_output.contiguous()) adapter_output = adapter_output.reshape(linear_output.shape) return linear_output + adapter_output, bias - def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Forward pass with Transformer Engine operation fuser The fused implementation is a PyTorch module that shares @@ -102,6 +114,8 @@ def __init__( bias: bool = True, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, + tensor_parallel_mode: Optional[str] = None, + sequence_parallel: bool = False, norm_type: Optional[str] = None, norm_eps: float = 1e-5, norm_zero_centered_gamma: bool = False, @@ -111,12 +125,26 @@ def __init__( ) -> None: super().__init__() + # Tensor parallel config + tensor_parallel_group = None + if tensor_parallel_mode is not None: + if parallel_state.get_tensor_model_parallel_world_size() == 1: + tensor_parallel_mode = None + else: + tensor_parallel_group = parallel_state.get_tensor_model_parallel_group() + if tensor_parallel_group is not None: + raise NotImplementedError("Tensor parallelism is not yet supported") + + # Construct fused modules self._make_main_branch( in_features, out_features, bias=bias, device=device, dtype=dtype, + tensor_parallel_mode=tensor_parallel_mode, + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=sequence_parallel, norm_type=norm_type, norm_eps=norm_eps, norm_zero_centered_gamma=norm_zero_centered_gamma, @@ -128,6 +156,9 @@ def __init__( lora_dim, device=device, dtype=dtype, + tensor_parallel_mode=tensor_parallel_mode, + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=sequence_parallel, lora_dropout=lora_dropout, lora_dropout_position=lora_dropout_position, lora_scale=lora_scale, @@ -141,10 +172,14 @@ def _make_main_branch( bias: bool, device: Optional[torch.device], dtype: Optional[torch.dtype], + tensor_parallel_mode: Optional[str], + tensor_parallel_group: Optional[torch.distributed.ProcessGroup], + sequence_parallel: bool, norm_type: Optional[str], norm_eps: float, norm_zero_centered_gamma: bool, ) -> None: + """Construct fused module for main branch (norm + fork + linear)""" # List of ops ops = [] @@ -179,7 +214,9 @@ def _make_main_branch( bias=bias, device=device, dtype=dtype, - tensor_parallel_mode=None, ### TODO Support TP + tensor_parallel_mode=tensor_parallel_mode, + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=sequence_parallel, ) ) @@ -194,10 +231,14 @@ def _make_lora_branch( *, device: Optional[torch.device], dtype: Optional[torch.dtype], + tensor_parallel_mode: Optional[str], + tensor_parallel_group: Optional[torch.distributed.ProcessGroup], + sequence_parallel: bool, lora_dropout: float, lora_dropout_position: str, lora_scale: float, ) -> None: + """Construct fused module for LoRA branch (lora_a + lora_b + add)""" # List of ops ops = [] @@ -206,7 +247,7 @@ def _make_lora_branch( if lora_dropout > 0 and lora_dropout_position == "pre": ops.append(te.ops.Dropout(lora_dropout)) - # LoRA linear ops + # LoRA A linear op self.lora_a_lora_branch_idx: int = len(ops) ops.append( te.ops.Linear( @@ -215,9 +256,16 @@ def _make_lora_branch( bias=False, device=device, dtype=dtype, - tensor_parallel_mode=None, ### TODO Support TP + tensor_parallel_mode=tensor_parallel_mode, + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=sequence_parallel, ) ) + + # LoRA B linear op + if tensor_parallel_mode == "column": + # All-gather along dim -1 + raise NotImplementedError("Column tensor parallelism is not yet supported") self.lora_b_lora_branch_idx: int = len(ops) ops.append( te.ops.Linear( @@ -226,7 +274,9 @@ def _make_lora_branch( bias=False, device=device, dtype=dtype, - tensor_parallel_mode=None, ### TODO Support TP + tensor_parallel_mode=None if tensor_parallel_mode is None else "column", + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=False, ) ) @@ -235,6 +285,9 @@ def _make_lora_branch( ops.append(te.ops.ConstantScale(lora_scale)) if lora_dropout > 0 and lora_dropout_position == "post": ops.append(te.ops.Dropout(lora_dropout)) + if tensor_parallel_mode == "row": + # Note: All-gather along dim -1 + raise NotImplementedError("Row tensor parallelism is not yet supported") # Add with main branch ops.append(te.ops.AddExtraInput()) @@ -242,7 +295,7 @@ def _make_lora_branch( # Fuse ops self.lora_branch = te.ops.Sequential(*ops) - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: linear_output, linear_input = self.main_branch(x) with te.fp8_autocast(enabled=False): out = self.lora_branch(linear_input, linear_output) @@ -280,6 +333,21 @@ def make_from_lora_linear(lora_linear: LoRALinear) -> TEFusedLoRALinear: constructor_kwargs["bias"] = bias is not None constructor_kwargs["dtype"] = weight.dtype + # Extract tensor parallel config + tensor_parallel_size = parallel_state.get_tensor_model_parallel_world_size() + if tensor_parallel_size > 1: + tensor_parallel_mode = None + sequence_parallel = False + if isinstance(orig_linear, (te.Linear, te.LayerNormLinear)): + tensor_parallel_mode = orig_linear.parallel_mode + sequence_parallel = orig_linear.sequence_parallel + constructor_kwargs["tensor_parallel_mode"] = tensor_parallel_mode + constructor_kwargs["sequence_parallel"] = sequence_parallel + if tensor_parallel_mode == "row": + constructor_kwargs["in_features"] *= tensor_parallel_size + elif tensor_parallel_mode == "column": + constructor_kwargs["out_features"] *= tensor_parallel_size + # Extract norm params from base linear module norm_type = None norm_weight = None @@ -319,17 +387,32 @@ def make_from_lora_linear(lora_linear: LoRALinear) -> TEFusedLoRALinear: # Construct fused module out = TEFusedLoRALinear(**constructor_kwargs) - # Replace fused module params + # Replace norm params if norm_type is not None: norm_op = out.main_branch[out.norm_main_branch_idx] + assert norm_op.weight.size() == norm_weight.size() norm_op.weight = norm_weight if norm_bias: + assert norm_op.bias.size() == norm_bias.size() norm_op.bias = norm_bias + + # Replace base linear params linear_op = out.main_branch[out.linear_main_branch_idx] + assert linear_op.weight.size() == weight.size() + if bias is None: + assert linear_op.bias is None + else: + assert linear_op.bias.size() == bias.size() linear_op.weight = weight linear_op.bias = bias - out.lora_branch[out.lora_a_lora_branch_idx].weight = lora_a_weight - out.lora_branch[out.lora_b_lora_branch_idx].weight = lora_b_weight + + # Replace LoRA params + lora_a_op = out.lora_branch[out.lora_a_lora_branch_idx] + lora_b_op = out.lora_branch[out.lora_b_lora_branch_idx] + assert lora_a_op.weight.size() == lora_a_weight.size() + assert lora_b_op.weight.size() == lora_b_weight.size() + lora_a_op.weight = lora_a_weight + lora_b_op.weight = lora_b_weight return out From 67de5f4155ccbb08ada2c5af3697ce51658e1fcd Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 5 Aug 2025 23:35:08 +0000 Subject: [PATCH 04/24] Enable fused LoRA based on model config Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index a85b9cc9fac1..125a7a25fb8a 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -50,6 +50,20 @@ class LoRALinear(AdapterWrapper): class to provide a specific implementation of the forward method. """ + def __init__( + self, + to_wrap: nn.Module, + adapter: nn.Module, + enable_op_fuser: bool = False, + ): + super().__init__(to_wrap, adapter) + + # Whether to enable implementation with Transformer Engine operation fuser + self._op_fuser_enabled: bool = enable_op_fuser + if parallel_state.get_tensor_model_parallel_world_size() > 1: + # TP is not yet supported + self._op_fuser_enabled = False + def forward( self, x: torch.Tensor, @@ -59,11 +73,8 @@ def forward( # pylint: disable=C0115,C0116 # Fused implementation - if ( - getattr(self, "_enable_fused_impl", False) - and parallel_state.get_tensor_model_parallel_world_size() == 1 - ): # TP is not yet supported - return self._fused_forward(x) + if self._op_fuser_enabled: + return self._fused_forward(x) linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs) adapter_output = self.adapter(layernorm_output.contiguous()) @@ -80,12 +91,12 @@ def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch. """ # Construct fused impl if needed - fused_impl = getattr(self, "_fused_impl", (None,))[0] + fused_impl = getattr(self, "_op_fuser_impl", (None,))[0] if fused_impl is None: if not HAVE_TE: raise RuntimeError("Fused LoRALinear implementation requires Transformer Engine") fused_impl = TEFusedLoRALinear.make_from_lora_linear(self) - self._fused_impl = (fused_impl,) # Wrap in tuple to avoid registering submodule + self._op_fuser_impl = (fused_impl,) # Wrap in tuple to avoid registering submodule # Apply fused impl return fused_impl(x) @@ -815,6 +826,7 @@ def transform(self, m: nn.Module, name=None, prefix=None): else: lora_cls = LinearAdapter + # Construct LoRA module return lora_cls( m, dim=self.dim, @@ -827,6 +839,7 @@ def transform(self, m: nn.Module, name=None, prefix=None): input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = ( get_adapter_attributes_from_linear(m) ) + enable_op_fuser = hasattr(m, "config") and m.config.use_transformer_engine_op_fuser logging.info(f"Adding lora to: {full_name}") adapter = ParallelLinearAdapter( in_features, @@ -849,7 +862,7 @@ def transform(self, m: nn.Module, name=None, prefix=None): dropout_recompute=self.dropout_recompute, base_linear_is_parallel=base_linear_is_parallel, ) - return LoRALinear(m, adapter) + return LoRALinear(m, adapter, enable_op_fuser=enable_op_fuser) return m From 1711fd30d35dc736250545d359c3937cebc8c7dc Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 5 Aug 2025 23:43:43 +0000 Subject: [PATCH 05/24] Tweak comments Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 125a7a25fb8a..3d0066d32c6f 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -297,7 +297,7 @@ def _make_lora_branch( if lora_dropout > 0 and lora_dropout_position == "post": ops.append(te.ops.Dropout(lora_dropout)) if tensor_parallel_mode == "row": - # Note: All-gather along dim -1 + # All-gather along dim -1 raise NotImplementedError("Row tensor parallelism is not yet supported") # Add with main branch From b23d8e0eb2eeca924360384b604f8492fb9e032e Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Wed, 6 Aug 2025 00:04:16 +0000 Subject: [PATCH 06/24] Add TE version checks Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 3d0066d32c6f..946510d4ef83 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -40,6 +40,7 @@ from nemo.collections.llm.peft.utils import get_adapter_attributes_from_linear, is_expert_linear from nemo.lightning.pytorch.callbacks.peft import PEFT, AdapterWrapper from nemo.utils import logging +from nemo.utils.te_utils import te_version class LoRALinear(AdapterWrapper): @@ -59,8 +60,8 @@ def __init__( super().__init__(to_wrap, adapter) # Whether to enable implementation with Transformer Engine operation fuser - self._op_fuser_enabled: bool = enable_op_fuser - if parallel_state.get_tensor_model_parallel_world_size() > 1: + self._op_fuser_enabled: bool = HAVE_TE_FUSED_LORA and enable_op_fuser + if self._op_fuser_enabled and parallel_state.get_tensor_model_parallel_world_size() > 1: # TP is not yet supported self._op_fuser_enabled = False @@ -93,8 +94,8 @@ def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch. # Construct fused impl if needed fused_impl = getattr(self, "_op_fuser_impl", (None,))[0] if fused_impl is None: - if not HAVE_TE: - raise RuntimeError("Fused LoRALinear implementation requires Transformer Engine") + if not HAVE_TE_FUSED_LORA: + raise RuntimeError("Fused LoRALinear implementation requires Transformer Engine 2.7+") fused_impl = TEFusedLoRALinear.make_from_lora_linear(self) self._op_fuser_impl = (fused_impl,) # Wrap in tuple to avoid registering submodule @@ -102,7 +103,10 @@ def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch. return fused_impl(x) -if HAVE_TE: +# Fused LoRA requires Transformer Engine 2.7+ +HAVE_TE_FUSED_LORA: bool = HAVE_TE and te_version() >= (2, 7) + +if HAVE_TE_FUSED_LORA: class TEFusedLoRALinear(nn.Module): """A LoRA adapter wrapper using Transformer Engine operation fuser @@ -839,7 +843,10 @@ def transform(self, m: nn.Module, name=None, prefix=None): input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = ( get_adapter_attributes_from_linear(m) ) - enable_op_fuser = hasattr(m, "config") and m.config.use_transformer_engine_op_fuser + enable_op_fuser = ( + hasattr(m, "config") + and getattr(m.config, "use_transformer_engine_op_fuser", False) + ) logging.info(f"Adding lora to: {full_name}") adapter = ParallelLinearAdapter( in_features, From 155738ee0a5d27e0dd3d6685e977c64938333f84 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Wed, 6 Aug 2025 00:09:42 +0000 Subject: [PATCH 07/24] Fix linter warning Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 946510d4ef83..8bf7e0a2c684 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -311,6 +311,7 @@ def _make_lora_branch( self.lora_branch = te.ops.Sequential(*ops) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + # pylint: disable=C0115,C0116 linear_output, linear_input = self.main_branch(x) with te.fp8_autocast(enabled=False): out = self.lora_branch(linear_input, linear_output) From 86a937dd5e946099c2db36512c4f33ba35bd5df9 Mon Sep 17 00:00:00 2001 From: timmoon10 Date: Wed, 6 Aug 2025 00:12:50 +0000 Subject: [PATCH 08/24] Apply isort and black reformatting Signed-off-by: timmoon10 --- nemo/collections/llm/peft/lora.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 8bf7e0a2c684..6816b2fbf209 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -327,13 +327,9 @@ def make_from_lora_linear(lora_linear: LoRALinear) -> TEFusedLoRALinear: if not isinstance(lora_linear, LoRALinear): raise ValueError(f"Expected LoRALinear, got {lora_linear.__class__}") if not isinstance(lora_linear.to_wrap, (te.Linear, te.LayerNormLinear, torch.nn.Linear)): - raise ValueError( - f"Unsupported class for LoRALinear wrapped linear ({lora_linear.to_wrap.__class__})" - ) + raise ValueError(f"Unsupported class for LoRALinear wrapped linear ({lora_linear.to_wrap.__class__})") if not isinstance(lora_linear.adapter, (LinearAdapter, TELinearAdapter, ParallelLinearAdapter)): - raise ValueError( - f"Unsupported class for LoRALinear adapter ({lora_linear.adapter.__class__})" - ) + raise ValueError(f"Unsupported class for LoRALinear adapter ({lora_linear.adapter.__class__})") # Args for TEFusedLoRALinear constructor constructor_kwargs = {"device": "meta"} # Do not initialize params @@ -844,10 +840,7 @@ def transform(self, m: nn.Module, name=None, prefix=None): input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = ( get_adapter_attributes_from_linear(m) ) - enable_op_fuser = ( - hasattr(m, "config") - and getattr(m.config, "use_transformer_engine_op_fuser", False) - ) + enable_op_fuser = hasattr(m, "config") and getattr(m.config, "use_transformer_engine_op_fuser", False) logging.info(f"Adding lora to: {full_name}") adapter = ParallelLinearAdapter( in_features, From d9eccfa0b44126774e52395eaeda8937b347eb5c Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 12 Aug 2025 04:56:55 +0000 Subject: [PATCH 09/24] Use in-place fork/add ops to enable GEMMs with beta=1 Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 6816b2fbf209..6164345d6b0d 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -218,7 +218,8 @@ def _make_main_branch( ops.append(te.ops.Quantize(forward=True, backward=False)) # Fork to LoRA branch - ops.append(te.ops.MakeExtraOutput()) + # Note: GEMM with beta=1 in backward pass + ops.append(te.ops.MakeExtraOutput(in_place=True)) # Main branch linear op self.linear_main_branch_idx: int = len(ops) @@ -305,7 +306,8 @@ def _make_lora_branch( raise NotImplementedError("Row tensor parallelism is not yet supported") # Add with main branch - ops.append(te.ops.AddExtraInput()) + # Note: GEMM with beta=1 in forward pass + ops.append(te.ops.AddExtraInput(in_place=True)) # Fuse ops self.lora_branch = te.ops.Sequential(*ops) From 9b84bac56a086359e08cea68eb9f503946fb9047 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 12 Aug 2025 05:08:17 +0000 Subject: [PATCH 10/24] Add ops directly to te.op.Sequential Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 42 ++++++++++++++----------------- 1 file changed, 19 insertions(+), 23 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 6164345d6b0d..ca4a1c886b7b 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -197,33 +197,35 @@ def _make_main_branch( """Construct fused module for main branch (norm + fork + linear)""" # List of ops - ops = [] + self.main_branch = te.ops.Sequential() # Norm op self.norm_main_branch_idx: Optional[int] = None if norm_type is not None: - self.norm_main_branch_idx = len(ops) + self.norm_main_branch_idx = len(self.main_branch) norm_kwargs = { "eps": norm_eps, "device": device, "dtype": dtype, "zero_centered_gamma": norm_zero_centered_gamma, } + norm_op = None if norm_type == "LayerNorm": - ops.append(te.ops.LayerNorm(in_features, **norm_kwargs)) + norm_op = te.ops.LayerNorm(in_features, **norm_kwargs) elif norm_type == "RMSNorm": - ops.append(te.ops.RMSNorm(in_features, **norm_kwargs)) + norm_op = te.ops.RMSNorm(in_features, **norm_kwargs) else: raise ValueError(f"Unsupported normalization ({norm_type})") - ops.append(te.ops.Quantize(forward=True, backward=False)) + self.main_branch.append(norm_op) + self.main_branch.append(te.ops.Quantize(forward=True, backward=False)) # Fork to LoRA branch # Note: GEMM with beta=1 in backward pass - ops.append(te.ops.MakeExtraOutput(in_place=True)) + self.main_branch.append(te.ops.MakeExtraOutput(in_place=True)) # Main branch linear op - self.linear_main_branch_idx: int = len(ops) - ops.append( + self.linear_main_branch_idx: int = len(self.main_branch) + self.main_branch.append( te.ops.Linear( in_features, out_features, @@ -236,9 +238,6 @@ def _make_main_branch( ) ) - # Fuse ops - self.main_branch = te.ops.Sequential(*ops) - def _make_lora_branch( self, in_features: int, @@ -257,15 +256,15 @@ def _make_lora_branch( """Construct fused module for LoRA branch (lora_a + lora_b + add)""" # List of ops - ops = [] + self.lora_branch = te.ops.Sequential() # LoRA pre-processing if lora_dropout > 0 and lora_dropout_position == "pre": - ops.append(te.ops.Dropout(lora_dropout)) + self.lora_branch.append(te.ops.Dropout(lora_dropout)) # LoRA A linear op - self.lora_a_lora_branch_idx: int = len(ops) - ops.append( + self.lora_a_lora_branch_idx: int = len(self.lora_branch) + self.lora_branch.append( te.ops.Linear( in_features, lora_dim, @@ -282,8 +281,8 @@ def _make_lora_branch( if tensor_parallel_mode == "column": # All-gather along dim -1 raise NotImplementedError("Column tensor parallelism is not yet supported") - self.lora_b_lora_branch_idx: int = len(ops) - ops.append( + self.lora_b_lora_branch_idx: int = len(self.lora_branch) + self.lora_branch.append( te.ops.Linear( lora_dim, out_features, @@ -298,19 +297,16 @@ def _make_lora_branch( # LoRA post-processing if lora_scale != 1: - ops.append(te.ops.ConstantScale(lora_scale)) + self.lora_branch.append(te.ops.ConstantScale(lora_scale)) if lora_dropout > 0 and lora_dropout_position == "post": - ops.append(te.ops.Dropout(lora_dropout)) + self.lora_branch.append(te.ops.Dropout(lora_dropout)) if tensor_parallel_mode == "row": # All-gather along dim -1 raise NotImplementedError("Row tensor parallelism is not yet supported") # Add with main branch # Note: GEMM with beta=1 in forward pass - ops.append(te.ops.AddExtraInput(in_place=True)) - - # Fuse ops - self.lora_branch = te.ops.Sequential(*ops) + self.lora_branch.append(te.ops.AddExtraInput(in_place=True)) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: # pylint: disable=C0115,C0116 From d6136c94e02f4e31b9ec3f604bda5518a9ba2d3b Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 14 Aug 2025 01:52:47 +0000 Subject: [PATCH 11/24] Move fused LoRA impl into LoRALinear subclass Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 463 +++++++++++------------------- 1 file changed, 162 insertions(+), 301 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index ca4a1c886b7b..bd4fa1d5272e 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -51,381 +51,234 @@ class LoRALinear(AdapterWrapper): class to provide a specific implementation of the forward method. """ - def __init__( - self, - to_wrap: nn.Module, - adapter: nn.Module, - enable_op_fuser: bool = False, - ): - super().__init__(to_wrap, adapter) - - # Whether to enable implementation with Transformer Engine operation fuser - self._op_fuser_enabled: bool = HAVE_TE_FUSED_LORA and enable_op_fuser - if self._op_fuser_enabled and parallel_state.get_tensor_model_parallel_world_size() > 1: - # TP is not yet supported - self._op_fuser_enabled = False - - def forward( - self, + def forward(self, x: torch.Tensor, *args, **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: # pylint: disable=C0115,C0116 - - # Fused implementation - if self._op_fuser_enabled: - return self._fused_forward(x) - linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs) adapter_output = self.adapter(layernorm_output.contiguous()) adapter_output = adapter_output.reshape(linear_output.shape) return linear_output + adapter_output, bias - def _fused_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Forward pass with Transformer Engine operation fuser - - The fused implementation is a PyTorch module that shares - params with this module. Since it owns no state, there is no - need for extra checkpointing logic. - - """ - - # Construct fused impl if needed - fused_impl = getattr(self, "_op_fuser_impl", (None,))[0] - if fused_impl is None: - if not HAVE_TE_FUSED_LORA: - raise RuntimeError("Fused LoRALinear implementation requires Transformer Engine 2.7+") - fused_impl = TEFusedLoRALinear.make_from_lora_linear(self) - self._op_fuser_impl = (fused_impl,) # Wrap in tuple to avoid registering submodule - - # Apply fused impl - return fused_impl(x) - # Fused LoRA requires Transformer Engine 2.7+ HAVE_TE_FUSED_LORA: bool = HAVE_TE and te_version() >= (2, 7) if HAVE_TE_FUSED_LORA: - class TEFusedLoRALinear(nn.Module): - """A LoRA adapter wrapper using Transformer Engine operation fuser + class TEFusedLoRALinear(LoRALinear): + """LoRA adapter wrapper using Transformer Engine operation fuser""" - Its compute is equivalent to LoRALinear. + def __init__(self, to_wrap: nn.Module, adapter: nn.Module): + super().__init__(to_wrap, adapter) + self._fused_branches: Optional[tuple[te.ops.Sequential, te.ops.Sequential]] = None - There are no guarantees on the checkpoint structure, either - for compatibility with other modules or for backward - compatibility. LoRALinear works around this by treating - TEFusedLoRALinear as stateless. + def _make_fused_branches(self) -> tuple[te.ops.Sequential, te.ops.Sequential]: + """Construct fused modules for main and LoRA branches""" - """ + # Extract layer size and tensor parallel config + kwargs = { + "in_features": self.to_wrap.weight.size(1), + "out_features": self.to_wrap.weight.size(0), + "tensor_parallel_mode": None, + "tensor_parallel_group": None, + "sequence_parallel": False, + } + tensor_parallel_size = parallel_state.get_tensor_model_parallel_world_size() + if tensor_parallel_size > 1: + kwargs["tensor_parallel_group"] = parallel_state.get_tensor_model_parallel_group() + if isinstance(self.to_wrap, (te.Linear, te.LayerNormLinear)): + kwargs["tensor_parallel_mode"] = self.to_wrap.parallel_mode + kwargs["sequence_parallel"] = self.to_wrap.sequence_parallel + if kwargs["tensor_parallel_mode"] == "row": + kwargs["in_features"] *= tensor_parallel_size + elif kwargs["tensor_parallel_mode"] == "column": + kwargs["out_features"] *= tensor_parallel_size - def __init__( - self, - in_features: int, - out_features: int, - lora_dim: int, - *, - bias: bool = True, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - tensor_parallel_mode: Optional[str] = None, - sequence_parallel: bool = False, - norm_type: Optional[str] = None, - norm_eps: float = 1e-5, - norm_zero_centered_gamma: bool = False, - lora_dropout: float = 0.0, - lora_dropout_position: str = "post", - lora_scale: float = 1.0, - ) -> None: - super().__init__() - - # Tensor parallel config - tensor_parallel_group = None - if tensor_parallel_mode is not None: - if parallel_state.get_tensor_model_parallel_world_size() == 1: - tensor_parallel_mode = None - else: - tensor_parallel_group = parallel_state.get_tensor_model_parallel_group() - if tensor_parallel_group is not None: - raise NotImplementedError("Tensor parallelism is not yet supported") + # Construct fused branches + main_branch = self._make_main_branch(**kwargs) + lora_branch = self._make_lora_branch(**kwargs) - # Construct fused modules - self._make_main_branch( - in_features, - out_features, - bias=bias, - device=device, - dtype=dtype, - tensor_parallel_mode=tensor_parallel_mode, - tensor_parallel_group=tensor_parallel_group, - sequence_parallel=sequence_parallel, - norm_type=norm_type, - norm_eps=norm_eps, - norm_zero_centered_gamma=norm_zero_centered_gamma, - ) - with te.fp8_model_init(enabled=False): - self._make_lora_branch( - in_features, - out_features, - lora_dim, - device=device, - dtype=dtype, - tensor_parallel_mode=tensor_parallel_mode, - tensor_parallel_group=tensor_parallel_group, - sequence_parallel=sequence_parallel, - lora_dropout=lora_dropout, - lora_dropout_position=lora_dropout_position, - lora_scale=lora_scale, - ) + return main_branch, lora_branch def _make_main_branch( self, + *, in_features: int, out_features: int, - *, - bias: bool, - device: Optional[torch.device], - dtype: Optional[torch.dtype], tensor_parallel_mode: Optional[str], tensor_parallel_group: Optional[torch.distributed.ProcessGroup], sequence_parallel: bool, - norm_type: Optional[str], - norm_eps: float, - norm_zero_centered_gamma: bool, - ) -> None: + ) -> te.ops.Sequential: """Construct fused module for main branch (norm + fork + linear)""" - # List of ops - self.main_branch = te.ops.Sequential() + # Check wrapped linear class + if not isinstance(self.to_wrap, (te.Linear, te.LayerNormLinear, torch.nn.Linear)): + raise ValueError( + f"Unsupported class for wrapped linear ({self.to_wrap.__class__.__name__})" + ) + + # Ops in main branch + main_branch = te.ops.Sequential() # Norm op - self.norm_main_branch_idx: Optional[int] = None - if norm_type is not None: - self.norm_main_branch_idx = len(self.main_branch) - norm_kwargs = { - "eps": norm_eps, - "device": device, - "dtype": dtype, - "zero_centered_gamma": norm_zero_centered_gamma, + if isinstance(self.to_wrap, te.LayerNormLinear): + norm_type = self.to_wrap.normalization + kwargs = { + "eps": self.to_wrap.eps, + "device": "meta", + "dtype": self.to_wrap.layer_norm_weight.dtype, + "zero_centered_gamma": self.to_wrap.zero_centered_gamma, } - norm_op = None + op = None if norm_type == "LayerNorm": - norm_op = te.ops.LayerNorm(in_features, **norm_kwargs) + op = te.ops.LayerNorm(in_features, **kwargs) + op.weight = self.to_wrap.layer_norm_weight + op.bias = self.to_wrap.layer_norm_bias elif norm_type == "RMSNorm": - norm_op = te.ops.RMSNorm(in_features, **norm_kwargs) + op = te.ops.RMSNorm(in_features, **kwargs) + op.weight = self.to_wrap.layer_norm_weight else: raise ValueError(f"Unsupported normalization ({norm_type})") - self.main_branch.append(norm_op) - self.main_branch.append(te.ops.Quantize(forward=True, backward=False)) + main_branch.append(op) # Fork to LoRA branch # Note: GEMM with beta=1 in backward pass - self.main_branch.append(te.ops.MakeExtraOutput(in_place=True)) - - # Main branch linear op - self.linear_main_branch_idx: int = len(self.main_branch) - self.main_branch.append( - te.ops.Linear( - in_features, - out_features, - bias=bias, - device=device, - dtype=dtype, - tensor_parallel_mode=tensor_parallel_mode, - tensor_parallel_group=tensor_parallel_group, - sequence_parallel=sequence_parallel, - ) + main_branch.append(te.ops.MakeExtraOutput(in_place=True)) + + # Linear op + weight = self.to_wrap.weight + bias = self.to_wrap.bias + if isinstance(bias, torch.Tensor) and bias.numel() == 0: + bias = None + op = te.ops.Linear( + in_features, + out_features, + bias=bias is not None, + device="meta", + dtype=weight.dtype, + tensor_parallel_mode=tensor_parallel_mode, + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=sequence_parallel, ) + op.weight = weight + op.bias = bias + main_branch.append(op) + + return main_branch def _make_lora_branch( self, + *, in_features: int, out_features: int, - lora_dim: int, - *, - device: Optional[torch.device], - dtype: Optional[torch.dtype], tensor_parallel_mode: Optional[str], tensor_parallel_group: Optional[torch.distributed.ProcessGroup], sequence_parallel: bool, - lora_dropout: float, - lora_dropout_position: str, - lora_scale: float, - ) -> None: + ) -> te.ops.Sequential: """Construct fused module for LoRA branch (lora_a + lora_b + add)""" - # List of ops - self.lora_branch = te.ops.Sequential() + from nemo.collections.llm.peft.utils import ParallelLinearAdapter + + # Extract params from LoRA adapter + lora_a_weight = None + lora_b_weight = None + lora_dim = None + dropout = None + dropout_position = None + scale = None + if isinstance(self.adapter, (LinearAdapter, TELinearAdapter)): + lora_a_weight = self.adapter.lora_a.weight + lora_b_weight = self.adapter.lora_b.weight + lora_dim = lora_b_weight.size(1) + dropout = self.adapter.dropout.p + dropout_position = self.adapter.dropout_position + scale = self.adapter.scale + elif isinstance(self.adapter, ParallelLinearAdapter): + lora_a_weight = self.adapter.linear_in.weight + lora_b_weight = self.adapter.linear_out.weight + lora_dim = lora_b_weight.size(1) + dropout = self.adapter.dropout.p + dropout_position = self.adapter.dropout_position + scale = self.adapter.alpha / self.adapter.dim + else: + raise ValueError( + f"Unsupported class for LoRA adapter ({self.adapter.__class__.__name__})" + ) + + # Ops in LoRA branch + lora_branch = te.ops.Sequential() # LoRA pre-processing - if lora_dropout > 0 and lora_dropout_position == "pre": - self.lora_branch.append(te.ops.Dropout(lora_dropout)) + if dropout > 0 and dropout_position == "pre": + lora_branch.append(te.ops.Dropout(dropout)) # LoRA A linear op - self.lora_a_lora_branch_idx: int = len(self.lora_branch) - self.lora_branch.append( - te.ops.Linear( - in_features, - lora_dim, - bias=False, - device=device, - dtype=dtype, - tensor_parallel_mode=tensor_parallel_mode, - tensor_parallel_group=tensor_parallel_group, - sequence_parallel=sequence_parallel, - ) + op = te.ops.Linear( + in_features, + lora_dim, + bias=False, + device="meta", + dtype=lora_a_weight.dtype, + tensor_parallel_mode=tensor_parallel_mode, + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=sequence_parallel, ) + op.weight = lora_a_weight + lora_branch.append(op) # LoRA B linear op if tensor_parallel_mode == "column": # All-gather along dim -1 raise NotImplementedError("Column tensor parallelism is not yet supported") - self.lora_b_lora_branch_idx: int = len(self.lora_branch) - self.lora_branch.append( - te.ops.Linear( - lora_dim, - out_features, - bias=False, - device=device, - dtype=dtype, - tensor_parallel_mode=None if tensor_parallel_mode is None else "column", - tensor_parallel_group=tensor_parallel_group, - sequence_parallel=False, - ) + op = te.ops.Linear( + lora_dim, + out_features, + bias=False, + device="meta", + dtype=lora_b_weight.dtype, + tensor_parallel_mode=None if tensor_parallel_mode is None else "column", + tensor_parallel_group=tensor_parallel_group, + sequence_parallel=False, ) + op.weight = lora_b_weight + lora_branch.append(op) # LoRA post-processing - if lora_scale != 1: - self.lora_branch.append(te.ops.ConstantScale(lora_scale)) - if lora_dropout > 0 and lora_dropout_position == "post": - self.lora_branch.append(te.ops.Dropout(lora_dropout)) + if scale != 1: + lora_branch.append(te.ops.ConstantScale(scale)) + if dropout > 0 and dropout_position == "post": + lora_branch.append(te.ops.Dropout(dropout)) if tensor_parallel_mode == "row": # All-gather along dim -1 raise NotImplementedError("Row tensor parallelism is not yet supported") # Add with main branch # Note: GEMM with beta=1 in forward pass - self.lora_branch.append(te.ops.AddExtraInput(in_place=True)) + lora_branch.append(te.ops.AddExtraInput(in_place=True)) + + return lora_branch - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: # pylint: disable=C0115,C0116 - linear_output, linear_input = self.main_branch(x) + + # Construct fused impl if needed + # Note: We initialize during the first forward pass in + # case the params are modified after the constructor. + # Note: The fused impl is stored in a tuple to avoid + # registering submodules. + if self._fused_branches is None: + self._fused_branches = self._make_fused_branches() + + # Apply fused impl + main_branch, lora_branch = self._fused_branches + linear_output, linear_input = main_branch(x) with te.fp8_autocast(enabled=False): - out = self.lora_branch(linear_input, linear_output) + out = lora_branch(linear_input, linear_output) return out, None - @staticmethod - def make_from_lora_linear(lora_linear: LoRALinear) -> TEFusedLoRALinear: - """Construct a fused LoRA adapter with the same params as an unfused adapter""" - - from nemo.collections.llm.peft.utils import ParallelLinearAdapter - - # Check inputs - if not isinstance(lora_linear, LoRALinear): - raise ValueError(f"Expected LoRALinear, got {lora_linear.__class__}") - if not isinstance(lora_linear.to_wrap, (te.Linear, te.LayerNormLinear, torch.nn.Linear)): - raise ValueError(f"Unsupported class for LoRALinear wrapped linear ({lora_linear.to_wrap.__class__})") - if not isinstance(lora_linear.adapter, (LinearAdapter, TELinearAdapter, ParallelLinearAdapter)): - raise ValueError(f"Unsupported class for LoRALinear adapter ({lora_linear.adapter.__class__})") - - # Args for TEFusedLoRALinear constructor - constructor_kwargs = {"device": "meta"} # Do not initialize params - - # Extract linear params from base linear module - orig_linear = lora_linear.to_wrap - weight = orig_linear.weight - bias = orig_linear.bias - if isinstance(bias, torch.Tensor) and bias.numel() == 0: - bias = None - constructor_kwargs["in_features"] = weight.size(1) - constructor_kwargs["out_features"] = weight.size(0) - constructor_kwargs["bias"] = bias is not None - constructor_kwargs["dtype"] = weight.dtype - - # Extract tensor parallel config - tensor_parallel_size = parallel_state.get_tensor_model_parallel_world_size() - if tensor_parallel_size > 1: - tensor_parallel_mode = None - sequence_parallel = False - if isinstance(orig_linear, (te.Linear, te.LayerNormLinear)): - tensor_parallel_mode = orig_linear.parallel_mode - sequence_parallel = orig_linear.sequence_parallel - constructor_kwargs["tensor_parallel_mode"] = tensor_parallel_mode - constructor_kwargs["sequence_parallel"] = sequence_parallel - if tensor_parallel_mode == "row": - constructor_kwargs["in_features"] *= tensor_parallel_size - elif tensor_parallel_mode == "column": - constructor_kwargs["out_features"] *= tensor_parallel_size - - # Extract norm params from base linear module - norm_type = None - norm_weight = None - norm_bias = None - if isinstance(orig_linear, te.LayerNormLinear): - norm_type = orig_linear.normalization - if norm_type == "LayerNorm": - norm_weight = orig_linear.layer_norm_weight - norm_bias = orig_linear.layer_norm_bias - elif norm_type == "RMSNorm": - norm_weight = orig_linear.layer_norm_weight - else: - raise RuntimeError("LayerNormLinear has unsupported norm type ({norm_type})") - constructor_kwargs["norm_type"] = norm_type - constructor_kwargs["norm_eps"] = orig_linear.eps - constructor_kwargs["norm_zero_centered_gamma"] = orig_linear.zero_centered_gamma - - # Extract params from LoRA adapter - adapter = lora_linear.adapter - lora_a_weight = None - lora_b_weight = None - if isinstance(adapter, (LinearAdapter, TELinearAdapter)): - lora_a_weight = adapter.lora_a.weight - lora_b_weight = adapter.lora_b.weight - constructor_kwargs["lora_dim"] = lora_a_weight.size(0) - constructor_kwargs["lora_dropout"] = adapter.dropout.p - constructor_kwargs["lora_dropout_position"] = adapter.dropout_position - constructor_kwargs["lora_scale"] = adapter.scale - elif isinstance(adapter, ParallelLinearAdapter): - lora_a_weight = adapter.linear_in.weight - lora_b_weight = adapter.linear_out.weight - constructor_kwargs["lora_dim"] = lora_a_weight.size(0) - constructor_kwargs["lora_dropout"] = adapter.dropout.p - constructor_kwargs["lora_dropout_position"] = adapter.dropout_position - constructor_kwargs["lora_scale"] = adapter.alpha / adapter.dim - - # Construct fused module - out = TEFusedLoRALinear(**constructor_kwargs) - - # Replace norm params - if norm_type is not None: - norm_op = out.main_branch[out.norm_main_branch_idx] - assert norm_op.weight.size() == norm_weight.size() - norm_op.weight = norm_weight - if norm_bias: - assert norm_op.bias.size() == norm_bias.size() - norm_op.bias = norm_bias - - # Replace base linear params - linear_op = out.main_branch[out.linear_main_branch_idx] - assert linear_op.weight.size() == weight.size() - if bias is None: - assert linear_op.bias is None - else: - assert linear_op.bias.size() == bias.size() - linear_op.weight = weight - linear_op.bias = bias - - # Replace LoRA params - lora_a_op = out.lora_branch[out.lora_a_lora_branch_idx] - lora_b_op = out.lora_branch[out.lora_b_lora_branch_idx] - assert lora_a_op.weight.size() == lora_a_weight.size() - assert lora_b_op.weight.size() == lora_b_weight.size() - lora_a_op.weight = lora_a_weight - lora_b_op.weight = lora_b_weight - - return out - if HAVE_TE: @@ -838,7 +691,12 @@ def transform(self, m: nn.Module, name=None, prefix=None): input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = ( get_adapter_attributes_from_linear(m) ) - enable_op_fuser = hasattr(m, "config") and getattr(m.config, "use_transformer_engine_op_fuser", False) + enable_op_fuser = ( + HAVE_TE_FUSED_LORA + and hasattr(m, "config") + and getattr(m.config, "use_transformer_engine_op_fuser", False) + and not base_linear_is_parallel # TP not yet supported + ) logging.info(f"Adding lora to: {full_name}") adapter = ParallelLinearAdapter( in_features, @@ -861,7 +719,10 @@ def transform(self, m: nn.Module, name=None, prefix=None): dropout_recompute=self.dropout_recompute, base_linear_is_parallel=base_linear_is_parallel, ) - return LoRALinear(m, adapter, enable_op_fuser=enable_op_fuser) + if enable_op_fuser: + return TEFusedLoRALinear(m, adapter) + else: + return LoRALinear(m, adapter) return m From 9f7df8d43cf293e7c20cb151028b172f4df119bb Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 14 Aug 2025 02:42:32 +0000 Subject: [PATCH 12/24] Fix bug where fused impl was always disabled Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index bd4fa1d5272e..08dd06251d1a 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -143,6 +143,7 @@ def _make_main_branch( else: raise ValueError(f"Unsupported normalization ({norm_type})") main_branch.append(op) + main_branch.append(te.ops.Quantize(forward=True, backward=False)) # Fork to LoRA branch # Note: GEMM with beta=1 in backward pass @@ -695,7 +696,8 @@ def transform(self, m: nn.Module, name=None, prefix=None): HAVE_TE_FUSED_LORA and hasattr(m, "config") and getattr(m.config, "use_transformer_engine_op_fuser", False) - and not base_linear_is_parallel # TP not yet supported + # TP not yet supported + and parallel_state.get_tensor_model_parallel_world_size() == 1 ) logging.info(f"Adding lora to: {full_name}") adapter = ParallelLinearAdapter( From 0044d7785e6f1b3a91328000896ad59b7110f0bc Mon Sep 17 00:00:00 2001 From: timmoon10 Date: Thu, 14 Aug 2025 02:43:20 +0000 Subject: [PATCH 13/24] Apply isort and black reformatting Signed-off-by: timmoon10 --- nemo/collections/llm/peft/lora.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 08dd06251d1a..1d8a0fa7a665 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -51,7 +51,8 @@ class LoRALinear(AdapterWrapper): class to provide a specific implementation of the forward method. """ - def forward(self, + def forward( + self, x: torch.Tensor, *args, **kwargs, @@ -116,9 +117,7 @@ def _make_main_branch( # Check wrapped linear class if not isinstance(self.to_wrap, (te.Linear, te.LayerNormLinear, torch.nn.Linear)): - raise ValueError( - f"Unsupported class for wrapped linear ({self.to_wrap.__class__.__name__})" - ) + raise ValueError(f"Unsupported class for wrapped linear ({self.to_wrap.__class__.__name__})") # Ops in main branch main_branch = te.ops.Sequential() @@ -205,9 +204,7 @@ def _make_lora_branch( dropout_position = self.adapter.dropout_position scale = self.adapter.alpha / self.adapter.dim else: - raise ValueError( - f"Unsupported class for LoRA adapter ({self.adapter.__class__.__name__})" - ) + raise ValueError(f"Unsupported class for LoRA adapter ({self.adapter.__class__.__name__})") # Ops in LoRA branch lora_branch = te.ops.Sequential() From c4b9aced46021cc161478b0e278e9b51adf54edf Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 22 Aug 2025 00:33:16 +0000 Subject: [PATCH 14/24] Support wgrad accumulation fusion Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index c6ed6f2e1d1b..87b3cf6cc41a 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -98,6 +98,12 @@ def _make_fused_branches(self) -> tuple[te.ops.Sequential, te.ops.Sequential]: elif kwargs["tensor_parallel_mode"] == "column": kwargs["out_features"] *= tensor_parallel_size + # wgrad accumulation fusion + accumulate_into_main_grad = False + if isinstance(self.to_wrap, (te.Linear, te.LayerNormLinear)): + accumulate_into_main_grad = self.to_wrap.fuse_wgrad_accumulation + kwargs["accumulate_into_main_grad"] = accumulate_into_main_grad + # Construct fused branches main_branch = self._make_main_branch(**kwargs) lora_branch = self._make_lora_branch(**kwargs) @@ -112,6 +118,7 @@ def _make_main_branch( tensor_parallel_mode: Optional[str], tensor_parallel_group: Optional[torch.distributed.ProcessGroup], sequence_parallel: bool, + accumulate_into_main_grad: bool, ) -> te.ops.Sequential: """Construct fused module for main branch (norm + fork + linear)""" @@ -162,6 +169,7 @@ def _make_main_branch( tensor_parallel_mode=tensor_parallel_mode, tensor_parallel_group=tensor_parallel_group, sequence_parallel=sequence_parallel, + accumulate_into_main_grad=accumulate_into_main_grad, ) op.weight = weight op.bias = bias @@ -177,6 +185,7 @@ def _make_lora_branch( tensor_parallel_mode: Optional[str], tensor_parallel_group: Optional[torch.distributed.ProcessGroup], sequence_parallel: bool, + accumulate_into_main_grad: bool, ) -> te.ops.Sequential: """Construct fused module for LoRA branch (lora_a + lora_b + add)""" @@ -223,6 +232,7 @@ def _make_lora_branch( tensor_parallel_mode=tensor_parallel_mode, tensor_parallel_group=tensor_parallel_group, sequence_parallel=sequence_parallel, + accumulate_into_main_grad=accumulate_into_main_grad, ) op.weight = lora_a_weight lora_branch.append(op) @@ -240,6 +250,7 @@ def _make_lora_branch( tensor_parallel_mode=None if tensor_parallel_mode is None else "column", tensor_parallel_group=tensor_parallel_group, sequence_parallel=False, + accumulate_into_main_grad=accumulate_into_main_grad, ) op.weight = lora_b_weight lora_branch.append(op) From 79dd433be3ee591c59cd81bf46f551515002cc4c Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 22 Aug 2025 03:41:43 +0000 Subject: [PATCH 15/24] Add integration test for TE op fuser Signed-off-by: Tim Moon --- .github/workflows/cicd-main-nemo2.yml | 2 ++ nemo/collections/llm/peft/lora.py | 5 +-- tests/collections/llm/common.py | 1 + tests/collections/llm/gpt_finetuning.py | 7 +++- ...NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser.sh | 34 +++++++++++++++++++ 5 files changed, 46 insertions(+), 3 deletions(-) create mode 100644 tests/functional_tests/L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser.sh diff --git a/.github/workflows/cicd-main-nemo2.yml b/.github/workflows/cicd-main-nemo2.yml index 092ac96c655f..46b345c39701 100644 --- a/.github/workflows/cicd-main-nemo2.yml +++ b/.github/workflows/cicd-main-nemo2.yml @@ -143,6 +143,8 @@ jobs: runner: self-hosted-azure - script: L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_Chat runner: self-hosted-azure + - script: L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser + runner: self-hosted-azure - script: L2_NeMo_2_Mixtral_LoRA_EP2PP1_MBS2_exclude runner: self-hosted-azure - script: L2_NeMo_2_Mixtral_LoRA_EP2PP1_MBS2 diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 87b3cf6cc41a..3787150c9f0e 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -195,7 +195,7 @@ def _make_lora_branch( lora_a_weight = None lora_b_weight = None lora_dim = None - dropout = None + dropout = 0 dropout_position = None scale = None if isinstance(self.adapter, (LinearAdapter, TELinearAdapter)): @@ -209,7 +209,8 @@ def _make_lora_branch( lora_a_weight = self.adapter.linear_in.weight lora_b_weight = self.adapter.linear_out.weight lora_dim = lora_b_weight.size(1) - dropout = self.adapter.dropout.p + if self.adapter.dropout is not None: + dropout = self.adapter.dropout.p dropout_position = self.adapter.dropout_position scale = self.adapter.alpha / self.adapter.dim else: diff --git a/tests/collections/llm/common.py b/tests/collections/llm/common.py index 6e1f75d501f0..b34dcb7b38f1 100644 --- a/tests/collections/llm/common.py +++ b/tests/collections/llm/common.py @@ -224,3 +224,4 @@ class Llama3ConfigCI(llm.Llama3Config8B): ffn_hidden_size: int = 3072 num_attention_heads: int = 8 vocab_size: int = 50304 + use_transformer_engine_op_fuser: bool = False diff --git a/tests/collections/llm/gpt_finetuning.py b/tests/collections/llm/gpt_finetuning.py index 150080263b5f..ac5a1e5e0cbe 100644 --- a/tests/collections/llm/gpt_finetuning.py +++ b/tests/collections/llm/gpt_finetuning.py @@ -42,6 +42,7 @@ def get_args(): parser.add_argument('--pp_size', type=int, default=1, help="pipeline parallel size") parser.add_argument('--packed', action='store_true', help="use packed sequence dataset") parser.add_argument('--dataset', type=str, default="dolly", choices=['dolly', 'chat'], help="Dataset to use") + parser.add_argument('--te_op_fuser', action='store_true', help="Enable Transformer Engine operation fuser") return parser.parse_args() @@ -98,6 +99,9 @@ def get_args(): peft = llm.peft.PEFT_STR2CLS[args.peft]() else: peft = None + if args.te_op_fuser: + # TE op fuser replaces MLP, so only finetune qkv and proj layers + peft.target_modules=["linear_proj", "linear_qkv"] packed_sequence_specs = ( PackedSequenceSpecs(packed_sequence_size=2048, tokenizer_model_name="dummy_tokenizer") if args.packed else None @@ -126,7 +130,8 @@ def get_args(): assert str(data.dataset_root).startswith(os.environ.get("NEMO_HOME")) tokenizer = get_nmt_tokenizer(tokenizer_model=os.path.join(args.restore_path, "dummy_tokenizer.model")) - llama3_8b = llm.LlamaModel(Llama3ConfigCI(), tokenizer=tokenizer) + model_config = Llama3ConfigCI(use_transformer_engine_op_fuser=args.te_op_fuser) + llama3_8b = llm.LlamaModel(model_config, tokenizer=tokenizer) resume = nl.AutoResume( restore_config=nl.RestoreConfig(path=args.restore_path), diff --git a/tests/functional_tests/L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser.sh b/tests/functional_tests/L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser.sh new file mode 100644 index 000000000000..b0f40a3d681f --- /dev/null +++ b/tests/functional_tests/L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser.sh @@ -0,0 +1,34 @@ +# Copyright (c) 2020-2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +coverage run -a --data-file=/workspace/.coverage --source=/workspace/nemo tests/collections/llm/gpt_finetuning.py \ + --restore_path /home/TestData/nemo2_ckpt/llama_68M_v4 \ + --devices 1 \ + --max_steps 3 \ + --experiment_dir /tmp/nemo2_gpt_finetune/$RUN_ID \ + --peft lora \ + --tp_size 1 \ + --pp_size 1 \ + --mbs 1 \ + --te_op_fuser + +coverage run -a --data-file=/workspace/.coverage --source=/workspace/nemo tests/collections/llm/gpt_finetuning.py \ + --restore_path /home/TestData/nemo2_ckpt/llama_68M_v4 \ + --devices 1 \ + --max_steps 6 \ + --experiment_dir /tmp/nemo2_gpt_finetune/$RUN_ID \ + --peft lora \ + --tp_size 1 \ + --pp_size 1 \ + --mbs 1 \ + --te_op_fuser From 793cc28b6ddec7fbe3c3eca4287eb35d4e10a6f2 Mon Sep 17 00:00:00 2001 From: timmoon10 Date: Fri, 22 Aug 2025 03:45:25 +0000 Subject: [PATCH 16/24] Apply isort and black reformatting Signed-off-by: timmoon10 --- tests/collections/llm/gpt_finetuning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/collections/llm/gpt_finetuning.py b/tests/collections/llm/gpt_finetuning.py index ac5a1e5e0cbe..56c9320d51fa 100644 --- a/tests/collections/llm/gpt_finetuning.py +++ b/tests/collections/llm/gpt_finetuning.py @@ -101,7 +101,7 @@ def get_args(): peft = None if args.te_op_fuser: # TE op fuser replaces MLP, so only finetune qkv and proj layers - peft.target_modules=["linear_proj", "linear_qkv"] + peft.target_modules = ["linear_proj", "linear_qkv"] packed_sequence_specs = ( PackedSequenceSpecs(packed_sequence_size=2048, tokenizer_model_name="dummy_tokenizer") if args.packed else None From e6ce9b86a22ef6c2f2a9c3f492e54a9f4c3f284d Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 22 Aug 2025 22:25:52 +0000 Subject: [PATCH 17/24] Explicitly list module containers that are compatible with list or dict APIs Mcore subclasses of te.ops.Sequential are iterable, but are not compatible with list API. Signed-off-by: Tim Moon --- nemo/collections/llm/fn/base.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/nemo/collections/llm/fn/base.py b/nemo/collections/llm/fn/base.py index 7fdf455fbfb8..335c289111ce 100644 --- a/nemo/collections/llm/fn/base.py +++ b/nemo/collections/llm/fn/base.py @@ -72,13 +72,21 @@ def double_weights(m): if not kwargs.pop("_skip_map", False) and hasattr(module, "map"): return module.map(func, leaf_only=leaf_only, **kwargs) + if ( + isinstance(module, nn.Module) + and not isinstance(module, (nn.Sequential, nn.ModuleList, nn.ModuleDict)) + ): + return _map_module(module, func, leaf_only=leaf_only, **kwargs) elif isinstance(module, Iterable): + # Assume iterable is API-compatible with dict or list if all(hasattr(module, key) for key in ["items", "values", "keys"]): return _map_module_dict(module, func, leaf_only=leaf_only, **kwargs) - return _map_module_list(module, func, leaf_only=leaf_only, **kwargs) else: - return _map_module(module, func, leaf_only=leaf_only, **kwargs) + raise ValueError( + "Expected `module` to be a PyTorch module or a collection of modules, " + f"but got {module.__class__.__name__}." + ) def walk( From a8055ed1c2eb30f1bfc82de8c4cc230a1b2b122e Mon Sep 17 00:00:00 2001 From: timmoon10 Date: Fri, 22 Aug 2025 22:26:40 +0000 Subject: [PATCH 18/24] Apply isort and black reformatting Signed-off-by: timmoon10 --- nemo/collections/llm/fn/base.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/nemo/collections/llm/fn/base.py b/nemo/collections/llm/fn/base.py index 335c289111ce..084f5d4e5773 100644 --- a/nemo/collections/llm/fn/base.py +++ b/nemo/collections/llm/fn/base.py @@ -72,10 +72,7 @@ def double_weights(m): if not kwargs.pop("_skip_map", False) and hasattr(module, "map"): return module.map(func, leaf_only=leaf_only, **kwargs) - if ( - isinstance(module, nn.Module) - and not isinstance(module, (nn.Sequential, nn.ModuleList, nn.ModuleDict)) - ): + if isinstance(module, nn.Module) and not isinstance(module, (nn.Sequential, nn.ModuleList, nn.ModuleDict)): return _map_module(module, func, leaf_only=leaf_only, **kwargs) elif isinstance(module, Iterable): # Assume iterable is API-compatible with dict or list From e9ceb353c34f291715fcda3915e7ebc6dd60a66c Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 22 Aug 2025 22:37:41 +0000 Subject: [PATCH 19/24] Add missing docstring Signed-off-by: Tim Moon --- nemo/collections/llm/fn/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nemo/collections/llm/fn/base.py b/nemo/collections/llm/fn/base.py index 084f5d4e5773..e84cc4b96b57 100644 --- a/nemo/collections/llm/fn/base.py +++ b/nemo/collections/llm/fn/base.py @@ -20,6 +20,7 @@ @runtime_checkable class HasBool(Protocol): + """Protocol for objects with bool operation""" def __bool__(self) -> bool: ... From 709656c612060df578dbc75e5f7d07086e0d6d28 Mon Sep 17 00:00:00 2001 From: timmoon10 Date: Fri, 22 Aug 2025 22:38:30 +0000 Subject: [PATCH 20/24] Apply isort and black reformatting Signed-off-by: timmoon10 --- nemo/collections/llm/fn/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nemo/collections/llm/fn/base.py b/nemo/collections/llm/fn/base.py index e84cc4b96b57..b2e1162ef20e 100644 --- a/nemo/collections/llm/fn/base.py +++ b/nemo/collections/llm/fn/base.py @@ -21,6 +21,7 @@ @runtime_checkable class HasBool(Protocol): """Protocol for objects with bool operation""" + def __bool__(self) -> bool: ... From d521c5ab204fde09bfd964452843ae5cc3ff3c94 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 28 Aug 2025 17:53:32 +0000 Subject: [PATCH 21/24] Update Mcore version Signed-off-by: Tim Moon --- requirements/manifest.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/manifest.json b/requirements/manifest.json index 0e48e043a557..631697804434 100644 --- a/requirements/manifest.json +++ b/requirements/manifest.json @@ -11,7 +11,7 @@ }, "megatron-lm": { "repo": "https://github.com/NVIDIA/Megatron-LM", - "ref": "7f7439f543288f50f134e44832069192a3e1d98e" + "ref": "d1a87770595a060f577d47f7796724400a917c52" }, "trt-llm": { "repo": "https://github.com/NVIDIA/TensorRT-LLM.git", From 95d50b9aa1b4ba3a92afa9150576ad5cf7a595e6 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Wed, 3 Sep 2025 18:46:59 +0000 Subject: [PATCH 22/24] Update Megatron-LM commit Signed-off-by: Tim Moon --- requirements/manifest.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/manifest.json b/requirements/manifest.json index 8fcf3cf943c9..f48e9c9e48dc 100644 --- a/requirements/manifest.json +++ b/requirements/manifest.json @@ -11,7 +11,7 @@ }, "megatron-lm": { "repo": "https://github.com/NVIDIA/Megatron-LM", - "ref": "69b65e00fe43bc77f361a3244f67e646c58b2627" + "ref": "122324c52d0ba0abe30ed5b8356044c557750897" }, "trt-llm": { "repo": "https://github.com/NVIDIA/TensorRT-LLM.git", From 942dc35f338a617018a7e5ebc74d51fc28862cde Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sun, 7 Sep 2025 22:45:34 +0000 Subject: [PATCH 23/24] Attempt to support forward hooks in fused LoRA Signed-off-by: Tim Moon --- nemo/collections/llm/peft/lora.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index 3787150c9f0e..f13411c299d4 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -108,6 +108,33 @@ def _make_fused_branches(self) -> tuple[te.ops.Sequential, te.ops.Sequential]: main_branch = self._make_main_branch(**kwargs) lora_branch = self._make_lora_branch(**kwargs) + # Get submodule forward hooks + forward_pre_hooks = [] + forward_post_hooks = [] + for submodule in self.modules(): + for hook in submodule._forward_pre_hooks.values(): + forward_pre_hooks.append((submodule, hook)) + for hook in submodule._forward_hooks.values(): + forward_post_hooks.append((submodule, hook)) + + # Attempt to emulate submodule forward hooks if needed + # Note: Assume hooks do not interact with submodule inputs + # or outputs since they are internal to the op fuser. + if forward_pre_hooks: + def forward_pre_hook(module, *_) -> None: + for submodule, hook in forward_pre_hooks: + # Assume that hook does not interact with + # input + hook(submodule, None) + main_branch.register_forward_pre_hook(forward_pre_hook) + if forward_post_hooks: + def forward_post_hook(module, *_) -> None: + for submodule, hook in forward_post_hooks: + # Assume that hook does not interact with + # input or output + hook(submodule, None, None) + lora_branch.register_forward_hook(forward_post_hook) + return main_branch, lora_branch def _make_main_branch( From 264e6a1c28fe48c05e1df5c5bbc4c02a2cc57b08 Mon Sep 17 00:00:00 2001 From: timmoon10 Date: Mon, 8 Sep 2025 02:19:46 +0000 Subject: [PATCH 24/24] Apply isort and black reformatting Signed-off-by: timmoon10 --- nemo/collections/llm/peft/lora.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index f13411c299d4..a8ea118e847f 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -121,18 +121,22 @@ def _make_fused_branches(self) -> tuple[te.ops.Sequential, te.ops.Sequential]: # Note: Assume hooks do not interact with submodule inputs # or outputs since they are internal to the op fuser. if forward_pre_hooks: + def forward_pre_hook(module, *_) -> None: for submodule, hook in forward_pre_hooks: # Assume that hook does not interact with # input hook(submodule, None) + main_branch.register_forward_pre_hook(forward_pre_hook) if forward_post_hooks: + def forward_post_hook(module, *_) -> None: for submodule, hook in forward_post_hooks: # Assume that hook does not interact with # input or output hook(submodule, None, None) + lora_branch.register_forward_hook(forward_post_hook) return main_branch, lora_branch