diff --git a/src/megatron/bridge/peft/lora.py b/src/megatron/bridge/peft/lora.py index a69b635a6d..f1944b6d14 100644 --- a/src/megatron/bridge/peft/lora.py +++ b/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 diff --git a/src/megatron/bridge/peft/lora_layers.py b/src/megatron/bridge/peft/lora_layers.py index 7c61b857dc..74d78e7653 100644 --- a/src/megatron/bridge/peft/lora_layers.py +++ b/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) diff --git a/tests/unit_tests/peft/test_lora_layers.py b/tests/unit_tests/peft/test_lora_layers.py index 297651db01..e8aa81bf26 100644 --- a/tests/unit_tests/peft/test_lora_layers.py +++ b/tests/unit_tests/peft/test_lora_layers.py @@ -19,15 +19,18 @@ functionality for Parameter-Efficient Fine-Tuning. """ +import os from copy import deepcopy +import megatron.core.parallel_state as parallel_state import pytest import torch +import torch.distributed as dist import torch.nn as nn import transformer_engine.pytorch as te from megatron.bridge.peft.lora import TELinearAdapter -from megatron.bridge.peft.lora_layers import LinearAdapter, LoRALinear, patch_linear_module +from megatron.bridge.peft.lora_layers import LinearAdapter, LoRALinear, TEFusedLoRALinear, patch_linear_module class MockLinearWithTupleReturn(nn.Module): @@ -329,6 +332,239 @@ def test_patch_linear_module_parameters(self, dim, alpha): assert patched_linear.lora_b.in_features == dim +class TestTEFusedLoRALinear: + """Test the TEFusedLoRALinear adapter wrapper with fused operations.""" + + @pytest.fixture(autouse=True) + def setup_and_teardown_parallel_state(self): + """Setup and teardown parallel state for Megatron tests.""" + + if not dist.is_initialized(): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29500" + os.environ["RANK"] = "0" + os.environ["LOCAL_RANK"] = "0" + os.environ["WORLD_SIZE"] = "1" + + device_count = torch.cuda.device_count() + if device_count > 0: + torch.cuda.set_device(0) + + init_process_group_kwargs = { + "backend": "nccl" if device_count > 0 else "gloo", + "world_size": 1, + "rank": 0, + } + + dist.init_process_group(**init_process_group_kwargs) + + assert dist.is_initialized(), "Distributed backend not initialized" + if not parallel_state.model_parallel_is_initialized(): + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=1, + ) + + assert parallel_state.model_parallel_is_initialized(), "Model parallel not initialized" + + from megatron.bridge.training.initialize import _set_random_seed + + _set_random_seed( + seed_=1234, + data_parallel_random_init=False, + te_rng_tracker=True, + inference_rng_tracker=False, + ) + + yield + + try: + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + # Clean up environment variables + for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "LOCAL_RANK", "WORLD_SIZE"]: + os.environ.pop(key, None) + except (NameError, AttributeError, RuntimeError): + pass + + @pytest.fixture + def te_layer_norm_linear_layernorm(self): + """Create a TELayerNormLinear with LayerNorm.""" + return te.LayerNormLinear(in_features=10, out_features=5, bias=True, normalization="LayerNorm", device="cuda") + + @pytest.fixture + def te_layer_norm_linear_rmsnorm(self): + """Create a TELayerNormLinear with RMSNorm.""" + return te.LayerNormLinear(in_features=10, out_features=5, bias=True, normalization="RMSNorm", device="cuda") + + @pytest.fixture + def te_linear(self): + """Create a basic TE linear layer.""" + return te.Linear(10, 5, device="cuda") + + @pytest.fixture + def linear_adapter(self): + """Create a LinearAdapter for LoRA.""" + linear = nn.Linear(10, 5, device="cuda") + return LinearAdapter(linear, dim=4, alpha=8) + + @pytest.fixture + def parallel_linear_adapter(self): + """Create a ParallelLinearAdapter for LoRA.""" + from megatron.bridge.peft.utils import ParallelLinearAdapter + + return ParallelLinearAdapter( + in_features=10, + out_features=5, + dim=4, + base_linear_name="test_linear", + alpha=8, + dropout=0.0, + ).cuda() + + def test_fused_lora_linear_with_layernorm(self, te_layer_norm_linear_layernorm, linear_adapter): + """Test TEFusedLoRALinear with LayerNormLinear (LayerNorm variant).""" + fused_lora = TEFusedLoRALinear(te_layer_norm_linear_layernorm, linear_adapter) + x = torch.randn(3, 10, device="cuda") + + output, bias = fused_lora(x) + + assert output.shape == (3, 5) + assert bias is None + + def test_fused_lora_linear_with_rmsnorm(self, te_layer_norm_linear_rmsnorm, linear_adapter): + """Test TEFusedLoRALinear with LayerNormLinear (RMSNorm variant).""" + fused_lora = TEFusedLoRALinear(te_layer_norm_linear_rmsnorm, linear_adapter) + x = torch.randn(3, 10, device="cuda") + + output, bias = fused_lora(x) + + assert output.shape == (3, 5) + assert bias is None + + def test_fused_lora_linear_with_te_linear(self, te_linear, linear_adapter): + """Test TEFusedLoRALinear with basic TELinear.""" + fused_lora = TEFusedLoRALinear(te_linear, linear_adapter) + x = torch.randn(3, 10, device="cuda") + + output, bias = fused_lora(x) + + assert output.shape == (3, 5) + assert bias is None + + def test_fused_lora_linear_with_parallel_adapter(self, te_linear, parallel_linear_adapter): + """Test TEFusedLoRALinear with ParallelLinearAdapter.""" + fused_lora = TEFusedLoRALinear(te_linear, parallel_linear_adapter) + x = torch.randn(3, 10, device="cuda") + + output, bias = fused_lora(x) + + assert output.shape == (3, 5) + assert bias is None + + def test_fused_lora_linear_with_te_linear_adapter(self, te_linear): + """Test TEFusedLoRALinear with TELinearAdapter.""" + te_adapter = TELinearAdapter(te_linear, dim=4, alpha=8) + fused_lora = TEFusedLoRALinear(te_linear, te_adapter) + x = torch.randn(3, 10, device="cuda") + + output, bias = fused_lora(x) + + assert output.shape == (3, 5) + assert bias is None + + def test_fused_lora_linear_unsupported_normalization(self, te_linear, linear_adapter): + """Test TEFusedLoRALinear with unsupported normalization type.""" + # Manually create a LayerNormLinear with an unsupported normalization + te_layer_norm = te.LayerNormLinear(10, 5, device="cuda", normalization="LayerNorm") + # Hack the normalization type to trigger the error + te_layer_norm.normalization = "UnsupportedNorm" + + fused_lora = TEFusedLoRALinear(te_layer_norm, linear_adapter) + x = torch.randn(3, 10, device="cuda") + + with pytest.raises(ValueError, match="Unsupported normalization"): + fused_lora(x) + + def test_fused_lora_linear_unsupported_adapter(self, te_linear): + """Test TEFusedLoRALinear with unsupported adapter type.""" + # Create an unsupported adapter type + unsupported_adapter = nn.Linear(10, 5, device="cuda") + + fused_lora = TEFusedLoRALinear(te_linear, unsupported_adapter) + x = torch.randn(3, 10, device="cuda") + + with pytest.raises(ValueError, match="Unsupported class for LoRA adapter"): + fused_lora(x) + + def test_fused_lora_linear_unsupported_wrapped_module(self, linear_adapter): + """Test TEFusedLoRALinear with unsupported wrapped module type.""" + # Create an unsupported wrapped module + conv = nn.Conv2d(3, 3, 3).cuda() + + fused_lora = TEFusedLoRALinear(conv, linear_adapter) + x = torch.randn(1, 3, 5, 5, device="cuda") + + with pytest.raises(ValueError, match="Unsupported class for wrapped linear"): + fused_lora(x) + + def test_fused_lora_linear_multiple_forward_passes(self, te_linear, linear_adapter): + """Test that fused branches are reused across forward passes.""" + fused_lora = TEFusedLoRALinear(te_linear, linear_adapter) + x = torch.randn(3, 10, device="cuda") + + # First forward pass initializes fused branches + output1, _ = fused_lora(x) + assert fused_lora._fused_branches is not None + + # Store reference to fused branches + fused_branches_ref = fused_lora._fused_branches + + # Second forward pass should reuse the same fused branches + output2, _ = fused_lora(x) + assert fused_lora._fused_branches is fused_branches_ref + + def test_fused_lora_linear_with_dropout(self, te_linear): + """Test TEFusedLoRALinear with dropout in adapter.""" + adapter = LinearAdapter(nn.Linear(10, 5, device="cuda"), dim=4, dropout=0.5) + fused_lora = TEFusedLoRALinear(te_linear, adapter) + x = torch.randn(3, 10, device="cuda") + + # Train mode + fused_lora.train() + output_train, _ = fused_lora(x) + + # Eval mode + fused_lora.eval() + output_eval, _ = fused_lora(x) + + assert output_train.shape == output_eval.shape == (3, 5) + + def test_fused_lora_linear_with_dropout_pre(self, te_linear): + """Test TEFusedLoRALinear with pre-dropout position.""" + adapter = LinearAdapter(nn.Linear(10, 5, device="cuda"), dim=4, dropout=0.3, dropout_position="pre") + fused_lora = TEFusedLoRALinear(te_linear, adapter) + x = torch.randn(3, 10, device="cuda") + + output, _ = fused_lora(x) + assert output.shape == (3, 5) + + def test_fused_lora_linear_with_scale(self, te_linear): + """Test TEFusedLoRALinear with different scale values.""" + adapter = LinearAdapter(nn.Linear(10, 5, device="cuda"), dim=4, alpha=16) + fused_lora = TEFusedLoRALinear(te_linear, adapter) + x = torch.randn(3, 10, device="cuda") + + output, _ = fused_lora(x) + assert output.shape == (3, 5) + # Verify scale is correctly set (alpha/dim = 16/4 = 4) + assert adapter.scale == 4.0 + + class TestTELinearAdapter: """Test the TELinearAdapter class."""