Wrap nn.Linear LoRA path with LoRALinear (to_wrap/adapter) (#5310)

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

src/megatron/bridge/models/conversion/peft_bridge.py

@@ -39,6 +39,7 @@
     persistent_buffers,
 )
 from megatron.bridge.peft.canonical_lora import ModuleDict
+from megatron.bridge.peft.lora_layers import LinearAdapter
 from megatron.bridge.peft.lora_merge import LoRAMerge
 from megatron.bridge.peft.utils import (
     get_adapter_attributes_from_linear,
@@ -595,6 +596,12 @@ def _megatron_global_adapters_info_all_pp_ranks(
                     input_is_parallel = adapter.input_is_parallel
                     base_linear_is_parallel = True
                     requires_expert_splits = adapter.linear_in.weight.ndim > 2
+                elif isinstance(adapter, LinearAdapter):
+                    # Adapter wrapping a plain nn.Linear: no parallelism layout
+                    # and the wrapped module has no Megatron ``config`` to derive one from.
+                    input_is_parallel = False
+                    base_linear_is_parallel = False
+                    requires_expert_splits = False
                 else:
                     attrs = get_adapter_attributes_from_linear(to_wrap)
                     input_is_parallel = attrs.input_is_parallel

src/megatron/bridge/peft/adapter_wrapper.py

@@ -16,6 +16,7 @@
 
 import torch
 import torch.nn as nn
+from megatron.core.utils import make_sharded_tensor_for_checkpoint
 
 from megatron.bridge.peft.utils import ParallelLinearAdapter
 
@@ -24,6 +25,33 @@
     from megatron.core.dist_checkpointing.mapping import ShardedStateDict
 
 
+def _plain_module_sharded_state_dict(
+    module: nn.Module,
+    prefix: str,
+    sharded_offsets: Tuple[Tuple[int, int, int], ...] = (),
+) -> "ShardedStateDict":
+    """Build a sharded state dict for a module that has no ``sharded_state_dict``.
+
+    Used as a fallback for plain ``nn.Linear`` / delta-only ``LinearAdapter`` leaves on
+    the simple (non-parallel) path: each parameter is wrapped as a replicated
+    :class:`ShardedTensor` so distributed checkpointing preserves the trained weights
+    instead of silently dropping them.
+
+    Args:
+        module: The plain module whose parameters to shard.
+        prefix: Key prefix applied to every parameter.
+        sharded_offsets: Offsets prepended to each sharded tensor.
+
+    Returns:
+        A sharded state dict mapping ``f"{prefix}{param_name}"`` to ShardedTensor.
+    """
+    sharded_state_dict: Dict[str, Any] = {}
+    for name, param in module.state_dict(prefix="", keep_vars=True).items():
+        key = f"{prefix}{name}"
+        sharded_state_dict[key] = make_sharded_tensor_for_checkpoint(param, key, prepend_offsets=sharded_offsets)
+    return sharded_state_dict
+
+
 def _compute_mamba_dim_info(wrapped_module: nn.Module) -> Dict[str, int]:
     """Compute Mamba dimension information from a wrapped module's config.
 
@@ -101,6 +129,7 @@ def __init__(self, to_wrap: nn.Module, adapter: nn.Module) -> None:
         self.to_wrap = to_wrap
         self.adapter = adapter
         self._adapter_enabled = True
+        self._base_returns_tuple = True
 
     def enable_adapter_layers(self) -> None:
         """Enable the adapter layers, allowing them to contribute to the forward pass output."""
@@ -140,13 +169,24 @@ def base_linear_forward(
             4. both: (out, bias, ln_out)
         """
         linear_output = self.to_wrap(x, *args, **kwargs)
-        assert isinstance(linear_output, tuple), (
-            f"{self.to_wrap} should return a tuple but instead returns {linear_output}"
-        )
 
         bias = None
         layernorm_output = x
 
+        # Plain nn.Linear (and similar modules) return a bare tensor rather than the
+        # Megatron-style tuple. Handle that case without asserting, so the wrapper stays
+        # a drop-in for nn.Linear in simple (non-parallel) models. Record the return
+        # shape so ``forward`` implementations can mirror it (tensor vs. tuple).
+        if isinstance(linear_output, torch.Tensor):
+            self._base_returns_tuple = False
+            bias = getattr(self.to_wrap, "bias", None)
+            return linear_output, bias, layernorm_output
+
+        assert isinstance(linear_output, tuple), (
+            f"{self.to_wrap} should return a tensor or a tuple but instead returns {type(linear_output).__name__}"
+        )
+        self._base_returns_tuple = True
+
         if len(linear_output) == 2:
             linear_output, bias = linear_output
             if isinstance(linear_output, tuple) and len(linear_output) == 2:
@@ -213,10 +253,24 @@ def sharded_state_dict(
             adapter_sharded_state_dict_kwargs["mamba_dim_info"] = _compute_mamba_dim_info(self.to_wrap)
 
         sharded_state_dict = {}
-        sharded_state_dict.update(self.to_wrap.sharded_state_dict(prefix, sharded_offsets, metadata))
-        sharded_state_dict.update(
-            self.adapter.sharded_state_dict(
-                f"{prefix}adapter.", sharded_offsets, metadata, **adapter_sharded_state_dict_kwargs
+        # The wrapped module may be a plain nn.Linear (simple, non-parallel path) that
+        # does not implement distributed sharding. Fall back to wrapping each parameter
+        # as a replicated ShardedTensor so the trained weights are preserved.
+        if hasattr(self.to_wrap, "sharded_state_dict"):
+            sharded_state_dict.update(self.to_wrap.sharded_state_dict(prefix, sharded_offsets, metadata))
+        else:
+            sharded_state_dict.update(_plain_module_sharded_state_dict(self.to_wrap, prefix, sharded_offsets))
+        # Likewise, a delta-only LinearAdapter has no sharded layout; only parallel adapters
+        # carry distributed-sharding information. Fall back to replicated ShardedTensors for
+        # the plain-adapter case so the trained LoRA matrices are not dropped.
+        if hasattr(self.adapter, "sharded_state_dict"):
+            sharded_state_dict.update(
+                self.adapter.sharded_state_dict(
+                    f"{prefix}adapter.", sharded_offsets, metadata, **adapter_sharded_state_dict_kwargs
+                )
+            )
+        else:
+            sharded_state_dict.update(
+                _plain_module_sharded_state_dict(self.adapter, f"{prefix}adapter.", sharded_offsets)
             )
-        )
         return sharded_state_dict

src/megatron/bridge/peft/canonical_lora.py

@@ -319,15 +319,16 @@ def transform(self, m: nn.Module, name: Optional[str] = None, prefix: Optional[s
         """
 
         # Skip already transformed modules
-        if isinstance(m, (LinearAdapter, LoRALinear, LoRALinearSplitQKV, LoRALinearSplitFC1UpGate, LoRATopKRouter)):
+        if isinstance(m, (LoRALinear, LoRALinearSplitQKV, LoRALinearSplitFC1UpGate, LoRATopKRouter)):
             return m
 
         if (ans := self.match(m, name, prefix)) is not None:
             (match, full_name) = ans
             if isinstance(m, nn.Linear) and not is_modelopt_linear(m):
-                return LinearAdapter(
+                adapter = LinearAdapter(
                     m, dim=self.dim, alpha=self.alpha, dropout=self.dropout, lora_A_init_method=self.lora_A_init_method
                 )
+                return LoRALinear(m, adapter)
 
             is_expert = is_expert_linear(full_name)
             attrs = get_adapter_attributes_from_linear(m, is_expert=is_expert)

src/megatron/bridge/peft/lora.py

@@ -127,21 +127,22 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio
             nn.Module: The modified module with LoRA applied, or the original module if not a target.
         """
         # Skip already transformed modules
-        adapter_types = (LinearAdapter, LoRALinear, LoRATopKRouter)
+        adapter_types = (LoRALinear, LoRATopKRouter)
         if isinstance(module, adapter_types):
             return module
 
         if (ans := self.match(module, name, prefix)) is not None:
             _, full_name = ans
             if isinstance(module, nn.Linear) and not is_modelopt_linear(module):
-                return LinearAdapter(
+                adapter = LinearAdapter(
                     module,
                     dim=self.dim,
                     alpha=self.alpha,
                     dropout=self.dropout,
                     lora_A_init_method=self.lora_A_init_method,
                     lora_dtype=self.lora_dtype,
                 )
+                return LoRALinear(module, adapter)
 
             is_expert = is_expert_linear(full_name)
             attrs = get_adapter_attributes_from_linear(

src/megatron/bridge/peft/lora_layers.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import math
-from typing import Any, Literal, Optional, Tuple
+from typing import Any, Literal, Optional
 
 import torch
 import torch.nn as nn
@@ -62,7 +62,7 @@ def bias(self) -> torch.Tensor | None:
         """Return the wrapped linear bias."""
         return getattr(self.to_wrap, "bias", None)
 
-    def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
         """Forward pass that combines the wrapped module output with the adapter output.
 
         Args:
@@ -71,18 +71,21 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten
             **kwargs: Additional keyword arguments for the wrapped module.
 
         Returns:
-            A tuple containing:
-                - Combined output (linear_output + adapter_output) if adapter is enabled,
-                  otherwise just the linear_output
-                - Bias term (if present, otherwise None)
+            When the wrapped module returns a Megatron-style tuple, a
+            ``(combined_output, bias)`` tuple; when it returns a bare tensor
+            (e.g. a plain ``nn.Linear``), a bare tensor so the wrapper stays a
+            drop-in replacement in simple (non-parallel) models.
         """
         # pylint: disable=C0115,C0116
         linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs)
         if not self._adapter_enabled:
-            return linear_output, bias
+            return linear_output if not self._base_returns_tuple else (linear_output, bias)
         adapter_output = self.adapter_forward(self.adapter, layernorm_output.contiguous(), *args, **kwargs)
         adapter_output = adapter_output.reshape(linear_output.shape)
-        return linear_output + adapter_output, bias
+        combined = linear_output + adapter_output
+        if not self._base_returns_tuple:
+            return combined
+        return combined, bias
 
 
 class LoRATopKRouter(AdapterWrapper):
@@ -343,11 +346,18 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
         return out, None
 
 
-class LinearAdapter(nn.Linear):
-    """Linear with LoRA that preserves the base weight and bias checkpoint keys.
+class LinearAdapter(nn.Module):
+    """Delta-only LoRA adapter for a plain ``nn.Linear``, mirroring :class:`ParallelLinearAdapter`'s role.
+
+    This adapter holds *only* the low-rank LoRA delta (``linear_in`` -> ``linear_out``) and
+    produces just the scaled adaptation term. It is intended to be wrapped together with the
+    original linear by :class:`LoRALinear` (the ``to_wrap`` / ``adapter`` pattern), so that
+    base weights and adapter weights live in distinct submodules and adapter state is
+    checkpointed under the ``adapter.`` prefix.
 
     Args:
-        orig_linear: The linear module to augment.
+        orig_linear: The linear module to augment (only its shape/dtype/device are used; its
+            weights are *not* copied).
         dim: LoRA's dimension (in_features -> dim -> out_features).
         alpha: LoRA's scaling alpha.
         dropout: Dropout probability (default: 0.0).
@@ -366,47 +376,31 @@ def __init__(
         lora_A_init_method: Literal["xavier", "uniform"] = "xavier",
         lora_dtype: Optional[torch.dtype] = None,
     ) -> None:
-        """Initialize LinearAdapter by copying from original Linear and adding LoRA components.
+        """Initialize the LoRA delta weights from the original Linear's shape and dtype.
 
         Args:
-            orig_linear: The original Linear module to adapt.
+            orig_linear: The original Linear module to adapt (weights are not copied).
             dim: LoRA rank dimension.
             alpha: LoRA scaling factor.
             dropout: Dropout probability.
             dropout_position: When to apply dropout ('pre' or 'post' LoRA computation).
             lora_A_init_method: Initialization method for LoRA matrix A.
             lora_dtype: Data type for LoRA weights.
         """
+        super().__init__()
         assert isinstance(orig_linear, nn.Linear)
-        super(LinearAdapter, self).__init__(
-            in_features=orig_linear.in_features,
-            out_features=orig_linear.out_features,
-            bias=orig_linear.bias is not None,
-            device=orig_linear.weight.device,
-            dtype=orig_linear.weight.dtype,
-        )
-        # copy weights
-        self.weight.data.copy_(orig_linear.weight.data)
-        if orig_linear.bias is not None:
-            self.bias.data.copy_(orig_linear.bias.data)
-        # initialize the adapter
+        self.in_features = orig_linear.in_features
+        self.out_features = orig_linear.out_features
         self._init_adapter(
             dim=dim,
             alpha=alpha,
             dropout=dropout,
             dropout_position=dropout_position,
             lora_A_init_method=lora_A_init_method,
             lora_dtype=lora_dtype,
+            device=orig_linear.weight.device,
+            base_dtype=orig_linear.weight.dtype,
         )
-        self._adapter_enabled = True
-
-    def enable_adapter_layers(self) -> None:
-        """Enable the adapter layers, allowing them to contribute to the forward pass output."""
-        self._adapter_enabled = True
-
-    def disable_adapter_layers(self) -> None:
-        """Disable the adapter layers, making the forward pass return only the base module output."""
-        self._adapter_enabled = False
 
     @torch.no_grad
     def _init_adapter(
@@ -417,8 +411,10 @@ def _init_adapter(
         dropout_position: Literal["pre", "post"] = "pre",
         lora_A_init_method: Literal["xavier", "uniform"] = "xavier",
         lora_dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        base_dtype: Optional[torch.dtype] = None,
     ) -> None:
-        """Initialize the LoRA weights and freeze the base parameters.
+        """Initialize the LoRA delta weights.
 
         Args:
             dim: LoRA's dimension (in_features -> dim -> out_features).
@@ -427,20 +423,16 @@ def _init_adapter(
             dropout_position: Where to apply dropout relative to LoRA (choices: ['pre', 'post'], default='pre').
             lora_A_init_method: Initialization method for lora_A (choices: ['xavier', 'uniform']).
             lora_dtype: Adapter weight dtype. Defaults to the base weight dtype.
+            device: Device for the LoRA weights.
+            base_dtype: Base weight dtype, used when ``lora_dtype`` is not provided.
         """
         self.dim = dim
         self.alpha = alpha
         self.scale = alpha / dim
 
-        # Freeze original weights
-        device = self.weight.device
-        self.weight.requires_grad = False
-        if self.bias is not None:
-            self.bias.requires_grad = False
-
         in_features = self.in_features
         out_features = self.out_features
-        dtype = lora_dtype or self.weight.dtype
+        dtype = lora_dtype or base_dtype
 
         self.linear_in = nn.Linear(in_features, dim, bias=False, dtype=dtype, device=device)
         self.linear_out = nn.Linear(dim, out_features, bias=False, dtype=dtype, device=device)
@@ -457,24 +449,20 @@ def _init_adapter(
         self.dropout_position = dropout_position
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass combining Linear output with LoRA adaptation.
+        """Compute the scaled LoRA delta only (no base-weight term).
 
         Args:
             x: Input tensor.
 
         Returns:
-            Combined output from original linear layer and LoRA adaptation.
+            The scaled low-rank adaptation ``scale * linear_out(linear_in(x))`` with dropout
+            applied per ``dropout_position``.
         """
         # pylint: disable=C0115,C0116
-        res = torch.nn.functional.linear(x, self.weight, self.bias)
-
-        if not self._adapter_enabled:
-            return res
-
         if self.dropout_position == "pre":
             x = self.dropout(x)
         lora_res = self.linear_out(self.linear_in(x))
         lora_res = lora_res * self.scale
         if self.dropout_position == "post":
             lora_res = self.dropout(lora_res)
-        return res + lora_res
+        return lora_res