Add option for LoRA with Transformer Engine op fuser

src/megatron/bridge/peft/lora.py

@@ -19,10 +19,17 @@
 import torch
 import torch.nn as nn
 import transformer_engine.pytorch as te
+from megatron.core import parallel_state
 from megatron.core.utils import unwrap_model
 
 from megatron.bridge.peft.base import PEFT
-from megatron.bridge.peft.lora_layers import LinearAdapter, LoRALinear, TELinearAdapter, patch_linear_module
+from megatron.bridge.peft.lora_layers import (
+    LinearAdapter,
+    LoRALinear,
+    TEFusedLoRALinear,
+    TELinearAdapter,
+    patch_linear_module,
+)
 from megatron.bridge.peft.module_matcher import ModuleMatcher
 from megatron.bridge.peft.utils import ParallelLinearAdapter, get_adapter_attributes_from_linear, is_expert_linear
 
@@ -130,6 +137,14 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio
             input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = (
                 get_adapter_attributes_from_linear(module)
             )
+
+            enable_op_fuser = (
+                hasattr(module, "config")
+                and getattr(module.config, "use_transformer_engine_op_fuser", False)
+                # TP not yet supported
+                and parallel_state.get_tensor_model_parallel_world_size() == 1
+            )
+
             logging.info(f"Adding lora to: {full_name}")
             adapter = ParallelLinearAdapter(
                 in_features,
@@ -151,7 +166,10 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio
                 disable_sequence_parallel_comm=disable_sp_comm,
                 base_linear_is_parallel=base_linear_is_parallel,
             )
-            return LoRALinear(module, adapter)
+            if enable_op_fuser:
+                return TEFusedLoRALinear(module, adapter)
+            else:
+                return LoRALinear(module, adapter)
         return module
 
 

src/megatron/bridge/peft/lora_layers.py

@@ -191,6 +191,259 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return res + lora_res
 
 
+class TEFusedLoRALinear(LoRALinear):
+    """LoRA adapter wrapper using Transformer Engine operation fuser"""
+
+    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
+
+    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,
+        }
+        # TODO: Restore once TP is supported
+        # 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
+
+        # 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)
+
+        # 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(
+        self,
+        *,
+        in_features: int,
+        out_features: int,
+        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)"""
+
+        # 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
+        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,
+            }
+            op = None
+            if norm_type == "LayerNorm":
+                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":
+                op = te.ops.RMSNorm(in_features, **kwargs)
+                op.weight = self.to_wrap.layer_norm_weight
+            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
+        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,
+            accumulate_into_main_grad=accumulate_into_main_grad,
+        )
+        op.weight = weight
+        op.bias = bias
+        main_branch.append(op)
+
+        return main_branch
+
+    def _make_lora_branch(
+        self,
+        *,
+        in_features: int,
+        out_features: int,
+        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)"""
+
+        from megatron.bridge.peft.utils import ParallelLinearAdapter
+
+        # Extract params from LoRA adapter
+        lora_a_weight = None
+        lora_b_weight = None
+        lora_dim = None
+        dropout = 0
+        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)
+            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:
+            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 dropout > 0 and dropout_position == "pre":
+            lora_branch.append(te.ops.Dropout(dropout))
+
+        # LoRA A linear op
+        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,
+            accumulate_into_main_grad=accumulate_into_main_grad,
+        )
+        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")
+        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,
+            accumulate_into_main_grad=accumulate_into_main_grad,
+        )
+        op.weight = lora_b_weight
+        lora_branch.append(op)
+
+        # LoRA post-processing
+        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
+        lora_branch.append(te.ops.AddExtraInput(in_place=True))
+
+        return lora_branch
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
+        # pylint: disable=C0115,C0116
+
+        # 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 = lora_branch(linear_input, linear_output)
+        return out, None
+
+
 class LinearAdapter(nn.Linear):
     """
     Linear + LoRA, maintains ckpts structure (i.e. Linear's weight/bias remain at the same FQN)