LoRA Fine-Tuning#

Low-Rank Adaptation (LoRA) fine-tunes a large pretrained model by freezing its weights and training only a small set of low-rank adapter matrices injected beside selected layers. Compared with full fine-tuning, this approach:

  • Adapts a model to a new dataset at a fraction of the cost

  • Ships a tiny adapter checkpoint (hundreds of KB instead of tens of MB)

  • Reduces overfitting and catastrophic forgetting in the small-data regime typical of scientific machine learning

Note

Parameter-Efficient Fine-Tuning (PEFT) is an experimental feature (physicsnemo.experimental.peft). APIs and functionality may change in future releases without backward compatibility guarantees. Contributions are welcome.

Location: physicsnemo.experimental.peft

The subpackage is a native, self-contained LoRA implementation. It operates on any torch.nn.Module. It does not require physicsnemo.Module or the .mdlus checkpoint format, but it works with them: load a physicsnemo.Module from a .mdlus checkpoint before applying LoRA. It supports both torch.nn.Linear and NVIDIA Transformer Engine layers (transformer_engine.pytorch.Linear and the fused LayerNormMLP), which is why it uses a native implementation rather than a third-party one.

Overview#

The subpackage exposes a small set of entry points:

Entry Point

Purpose

LoRAConfig

Declare the adapter (rank, alpha, which layers to adapt)

apply_lora

Inject adapters into a model in place and freeze the base

split_params_for_optimizer

Separate adapter parameters for optimizer construction

save_adapter and load_adapter

Save or load the adapter (a small archive)

merge_lora

Fold adapters into the base weights for zero-overhead inference

print_trainable_parameters

Report the trainable-parameter fraction

Tip

When to use LoRA. Prefer LoRA over full fine-tuning when adapting a pretrained foundation model to a new chemistry, geometry class, or operating regime. It is especially valuable with a small target dataset, a tight memory budget (frozen layers drop their saved activations), or a need to distribute small adapters separately from their base models. Full fine-tuning can still win when a large, high-quality dataset is available. For a complete workflow, refer to the GeoTransolver LoRA fine-tuning recipe.

How LoRA Works#

A linear layer computes y = W x. LoRA leaves W frozen and learns a low-rank update ΔW = B A, where A and B have an inner dimension rank much smaller than the layer width. The adapted layer computes:

y = W x + (alpha / rank) * (B A) x

Only A and B are trained. rank controls adapter capacity and alpha controls its strength (scaling = alpha / rank). B is initialized to zero, so at the start of training the adapter is exactly zero and the model reproduces the base model precisely. Setting the adapter to zero at any time recovers the original model. A uses Kaiming-uniform initialization by default. To choose another scheme, pass init to LoRAConfig. It accepts the string "default" or a callable that initializes the lora_A tensor in place, for example init=lambda t: torch.nn.init.normal_(t, std=0.01). B stays zero-initialized regardless, so the adapter is always identity at the start.

Quickstart#

The workflow applies, trains, and saves the adapter, then loads or merges it for deployment:

import torch
from physicsnemo.experimental.peft import (
    LoRAConfig, apply_lora, split_params_for_optimizer,
    print_trainable_parameters, save_adapter,
)

model = build_model()                 # any torch.nn.Module
model.load_state_dict(pretrained)     # start from a pretrained base

# 1) Inject adapters into the attention projections; freeze everything else.
config = LoRAConfig(rank=16, alpha=16,
                    target_pattern=r"blocks\.\d+\.attn\.(q|k|v|out)_proj")
apply_lora(model, config)
print_trainable_parameters(model)     # "trainable params: N (X% of M total)"

# 2) Train only the adapter parameters with AdamW.
groups = split_params_for_optimizer(model)
optimizer = torch.optim.AdamW(groups["lora"] + groups["extras"], lr=5e-4)
# ... standard training loop ...

# 3) Save the (small) adapter.
save_adapter(model, "adapter.lora")

Choosing Which Layers to Adapt#

A LoRAConfig must set exactly one of three selectors. All three match against fully qualified module names (for example, blocks.3.Attn.qkv_project), not bare leaf names. The same short name often appears in many submodules, so matching the full path avoids adapting the wrong layers.

The three selectors, and when to use each:

Selector

You Provide

Use When

target_modules

a list of exact fully qualified names

you want a hand-picked set of layers

target_pattern

a regular expression (re.search)

“all attention projections in every block” — the common case

target_filter

a predicate (name, module) -> bool

the choice depends on something a regex cannot see (shape, type, index)

 
# Exact names:
LoRAConfig(target_modules=["blocks.0.Attn.qkv_project", "head"])

# Regex (note the anchor on the block path, so only attention layers match):
LoRAConfig(target_pattern=r"blocks\.\d+\.Attn\.(qkv_project|out_linear|cross_[qkv])")

# Predicate (Python-only; cannot be expressed in a YAML config):
LoRAConfig(target_filter=lambda name, m: "Attn" in name and m.weight.shape[0] > 256)

Only registered layer types are eligible to be wrapped (refer to Transformer Engine Support and Implementing a Custom Wrapper). A name match on, say, a LayerNorm is skipped. If a selector matches zero wrappable layers, apply_lora raises. A mistyped pattern fails loudly rather than silently training nothing.

Two optional modifiers refine the selection:

  • wrap_mlp=True — additively adapts the transformer feed-forward (FFN) sub-block: the position-wise Linear → activation → Linear that follows attention in each block (in PhysicsNeMo transformer blocks this is the ln_mlp1 module — the fused te.LayerNormMLP under Transformer Engine, otherwise a Sequential(LayerNorm, Mlp)). It targets that sub-block specifically, not arbitrary or standalone MLPs, and not the whole model. It is also a no-op on architectures without that feed-forward structure. For non-transformer models, select layers explicitly with the selectors above.

  • extras_trainable=[...] — names modules to train fully (not low-rank), for example a final normalization or output head.

LoRAConfig is Hydra-instantiable, for example in a training config:

peft:
  _target_: physicsnemo.experimental.peft.LoRAConfig
  rank: 16
  alpha: 16
  target_pattern: 'blocks\.\d+\.Attn\.(qkv_project|out_linear|cross_[qkv])'
  wrap_mlp: false

Optimizer Setup#

apply_lora freezes the base and leaves only the adapter (and any extras_trainable) parameters trainable. Use split_params_for_optimizer to collect them and route them to AdamW:

groups = split_params_for_optimizer(model)        # {'lora', 'extras', 'frozen'}
optimizer = torch.optim.AdamW(groups["lora"] + groups["extras"], lr=5e-4)

Note

Do not route adapter parameters to optimizers such as Muon whose Newton-Schulz orthogonalization is degenerate on low-rank factors. This is the purpose of split_params_for_optimizer. It keeps the adapter parameters separate so they go to AdamW even when the base training script uses a different optimizer for the (now frozen) full-rank weights.

Under Distributed Data Parallel, wrap the model with find_unused_parameters=True because the frozen base parameters receive no gradient:

ddp_model = torch.nn.parallel.DistributedDataParallel(
    model, find_unused_parameters=True,
)

Saving and Loading Adapters#

save_adapter writes a small archive containing only the trainable adapter tensors. It does not store the frozen base, which is why adapters are tiny. The archive is a plain ZIP holding three entries:

  • adapter_config.json — the adapter configuration, including rank, alpha, dropout, init, and the explicit list of the actually wrapped layer names, so that the adapter reloads identically regardless of how the layers were originally selected.

  • adapter_model.pt — the adapter tensors and any extras_trainable params.

  • metadata.json — kind="lora_adapter", versions, and a base fingerprint (a hash of the base model’s structure, not its weights).

Important

Only load_adapter reads the adapter archive. It is neither a torch.save file nor a physicsnemo.Module (.mdlus) checkpoint, and torch.load and Module.load cannot read it. load_adapter accepts any file extension, but use a dedicated one such as .lora so that the name does not imply the wrong loader.

load_adapter reconstructs the adapters on a fresh base model. It verifies that the file is a LoRA adapter and checks the base fingerprint, rejecting a mismatched architecture or warning when you pass strict=False. It then re-applies LoRA and loads the adapter weights:

base = build_model()
base.load_state_dict(pretrained)
load_adapter(base, "adapter.lora")    # ready for inference (adapter active)

Note

load_adapter reads the adapter tensors with weights_only=True, so an adapter from an untrusted source cannot execute arbitrary code on load. This matters because adapters may be distributed independently of the base model.

The current API installs one adapter on a fresh base model. To use a different adapter, construct another fresh instance of the same base model and call load_adapter on that instance. load_adapter does not replace an adapter that is already installed.

Merging for Inference#

For zero inference overhead, fold the adapters into the base weights with merge_lora. It replaces the wrappers with ordinary layers, and the model is a plain model again:

from physicsnemo.experimental.peft import merge_lora

merge_lora(base)      # in place; base is now a standard model with no adapters

To compare base and adapter behavior without merging or reloading, toggle the adapters with set_adapter_enabled.

Transformer Engine Support#

LoRA wraps NVIDIA Transformer Engine layers in addition to torch.nn.Linear:

  • transformer_engine.pytorch.Linear is adapted exactly like nn.Linear, with Transformer Engine keyword arguments passed through to the base layer.

  • The fused transformer_engine.pytorch.LayerNormMLP (LayerNorm → fc1 → activation → fc2 in a single op) has no per-matrix linear to wrap, so LoRA adapts it with a single low-rank residual across the whole feed-forward sub-block. This preserves the fused (and FP8) kernel but, unlike per-matrix adapters, cannot be merged into the fused weights. merge_lora leaves it in place and logs a warning, and you deploy it using load_adapter rather than merging.

When wrap_mlp=True, LoRA adapts the feed-forward layers automatically through the appropriate mechanism for the backend (per-matrix for the non-TE MLP, the residual adapter for the fused TE MLP).

Implementing a Custom Wrapper#

A registry drives adapter support, mapping a base-layer type to a wrapper that knows how to adapt it. Register one pair, and everything else works unchanged, including targeting, freezing, the optimizer split, saving, and merging. This works because the rest of the package identifies adapters structurally (through isinstance(module, LoRALayer)) rather than by type.

The Wrapper Contract#

A wrapper is an nn.Module that subclasses LoRALayer and provides:

  • Constructor. A wrapper’s constructor:

    • __init__(self, base_layer, *, rank, alpha, dropout=0.0, init="default").

    • apply_lora instantiates wrappers as wrapper(base_layer, rank=, alpha=, dropout=, init=). It accepts **kwargs to stay forward-compatible with options added later.

  • Attributes. A wrapper exposes these attributes:

    • base_layer is the wrapped, frozen module.

    • lora_A and lora_B are the trainable nn.Parameter factors. They are the only parameters left with requires_grad=True, which is how save_adapter slices the adapter. They may be @property objects if the factors live in submodules.

    • enabled is a bool that toggles the delta.

    • mergeable is a bool that defaults to False. Set it True only if you also implement merge_into_base.

  • Methods. A wrapper implements:

    • forward adds the low-rank delta to the base output when enabled.

    • merge_into_base folds the delta into the base weight. It is required only when mergeable is True.

    • is_compatible(base_layer) is an optional classmethod that returns False to skip instances this wrapper cannot adapt. It defaults to True.

The base class gives you two optional conveniences for the common tensor case:

  • _make_lora_params(...) creates the lora_A and lora_B factors (lora_B zero-initialized, lora_A per the init argument)

  • lora_delta(x) computes ((dropout(x) @ A) @ B) * scaling.

A wrapper whose adapter is not a plain matmul (Matrix Multiplication) builds its own factors and delta instead.

Case 1: Reuse a Built-In Wrapper for a Linear-Like Layer#

If your custom layer behaves like nn.Linear, meaning it exposes in_features and out_features or a 2-D weight, you don’t need to write a wrapper at all. Register the built-in LoRALinear for it:

from physicsnemo.experimental.peft import register_lora_wrapper, LoRALinear

register_lora_wrapper(MyLinearLikeLayer, LoRALinear)

Case 2: A New Wrapper for a Different Adapter Algebra#

When the adaptation is genuinely different from a matmul, write a wrapper. A good example is nn.Embedding, whose forward is an index lookup. The adapter looks up rows of A for the input indices and projects them through B. This is why an embedding cannot go through the Linear path, because its 2-D weight is (num_embeddings, embedding_dim), not (out, in), and its forward is not a matmul.

import torch
import torch.nn as nn
import torch.nn.functional as F
from physicsnemo.experimental.peft import LoRALayer, register_lora_wrapper

class LoRAEmbedding(nn.Module, LoRALayer):
    """LoRA for nn.Embedding: delta = (lookup rows of A) @ B * scaling."""

    mergeable = True  # delta is (num_embeddings, dim) — foldable into the table

    def __init__(self, base_layer, rank, alpha, dropout=0.0, init="default"):
        nn.Module.__init__(self)
        self.base_layer = base_layer
        for p in self.base_layer.parameters():
            p.requires_grad = False
        # Reuse the convenience: lora_A is (num_embeddings, rank),
        # lora_B is (rank, embedding_dim); lora_B is zero so the delta is
        # zero at init. base_layer.weight supplies device/dtype.
        self._make_lora_params(
            base_layer.num_embeddings, base_layer.embedding_dim,
            base_layer.weight, rank, alpha, dropout, init,
        )

    def forward(self, idx):
        out = self.base_layer(idx)
        if self.enabled:
            out = out + (F.embedding(idx, self.lora_A) @ self.lora_B) * self.scaling
        return out

    @torch.no_grad()
    def merge_into_base(self):
        delta = (self.lora_A @ self.lora_B) * self.scaling
        self.base_layer.weight.add_(delta.to(self.base_layer.weight.dtype))

register_lora_wrapper(nn.Embedding, LoRAEmbedding)

After registration, the new type is targeted exactly like any other layer — by name in a LoRAConfig selector — and participates in freezing, the optimizer split, save_adapter and load_adapter, and (because it set mergeable = True) merge_lora.

Instance-Level Eligibility#

Some wrappers can handle a type but not every instance of it. Override is_compatible to veto those. apply_lora then skips the layer (with a warning if your selector matched it) instead of failing when it tries to wrap it. For example, an equivariant e3nn.o3.Linear adapter requires at least one shared input and output irrep:

class LoRA_o3_Linear(nn.Module, LoRALayer):
    @classmethod
    def is_compatible(cls, base_layer):
        shared = set(base_layer.irreps_in) & set(base_layer.irreps_out)
        return len(shared) > 0
    # ... __init__ / forward / merge_into_base ...

Some deltas cannot be folded into the base weights, such as the fused te.LayerNormMLP residual adapter. For these, keep the default mergeable = False. merge_lora then skips the wrapper, and the adapter stays active at inference.

End-to-End Example#

A complete recipe for fine-tuning a pretrained GeoTransolver on a custom external-aerodynamics dataset ships with the transformer models example, in its own src/finetune/ folder (separate from the main src/ training and inference scripts):

  • finetune.py — load a pretrained base, apply LoRA, train the adapters, save the adapter.

  • deploy.py — load the adapter onto a base, or merge for deployment.

  • finetune_lora.yaml — the Hydra configuration, including the peft block.

Refer to the LoRA fine-tuning README for the full walkthrough.

API Reference#

class physicsnemo.experimental.peft.LoRAConfig(rank: int = 16, alpha: float | None = None, target_modules: list[str] | None = None, target_pattern: str | None = None, target_filter: Callable[[str, 'nn.Module'], bool] | None = None, lora_dropout: float = 0.0, extras_trainable: list[str] = <factory>, wrap_mlp: bool = False, init: LoRAInit = 'default')[source]#

Bases: object

Configuration for applying LoRA to a model.

Exactly one of target_modules, target_pattern or target_filter must be provided. They select layers by fully-qualified module name (e.g. blocks.3.Attn.qkv_project), NOT bare leaf names — leaf names are not unique (the same short name can appear in many submodules).

Parameters:

  • rank (int) – Low-rank dimension r. Must be positive.

  • alpha (float | None) – LoRA scaling numerator; scaling = alpha / rank. None defaults alpha to rank (scaling 1.0).

  • target_modules (list[str] | None) – Exact fully-qualified module names to wrap.

  • target_pattern (str | None) – Regex (re.search) matched against fully-qualified module names.

  • target_filter (Callable[[str, nn.Module], bool] | None) – Predicate (name, module) -> bool (most flexible selector).

  • lora_dropout (float) – Dropout on the LoRA input path; 0.0 disables it. In [0.0, 1.0).

  • extras_trainable (list[str]) – Additional fully-qualified module names to leave fully trainable (not low-rank), e.g. a final head or norm.

  • wrap_mlp (bool) – Convenience flag to also adapt the transformer feed-forward (FFN) sub-block — the position-wise Linear -> activation -> Linear that follows attention in a transformer block (NOT arbitrary or standalone MLPs, and not the model as a whole). In PhysicsNeMo transformer blocks this is the ln_mlp1 module: under Transformer Engine the fused te.LayerNormMLP, otherwise a Sequential(LayerNorm, Mlp). Matched by the known feed-forward naming of those blocks, so it is a no-op on models without that structure. Additive to the selector above.

  • init ({"default"} or callable) – Initialization for the lora_A factor (lora_B is always zero, so the adapter is identity at init). "default" uses kaiming_uniform_(a=√5) — matching nn.Linear and the common PEFT default. Pass a callable (tensor) -> None to initialize lora_A in place with a custom scheme (e.g. lambda t: nn.init.normal_(t, std=0.01) for a Gaussian with a scale you control). Honored by wrappers built on _make_lora_params; wrappers with their own parameterization (e.g. equivariant layers) initialize themselves and ignore this.

property effective_alpha: float#

alpha if set, else equal to rank (→ scaling 1.0).

property scaling: float#

The LoRA scaling factor alpha / rank.

physicsnemo.experimental.peft.apply_lora(
model: Module,
config: LoRAConfig,
compute_fingerprint: bool = True,
) → ApplyResult[source]#

In-place: wrap matched Linear / te.Linear layers with LoRA and freeze the base (except extras_trainable).

Parameters:

  • model (nn.Module) – Model to mutate in place.

  • config (LoRAConfig) – LoRA configuration controlling target selection, rank, scaling, initialization, dropout, and any extra trainable modules.

  • compute_fingerprint (bool, optional) – If True, compute and store a fingerprint of the pristine base model before wrapping layers. Set to False to skip this cost.

Raises:

ValueError – If the model already contains LoRA layers (not re-entrant), or if zero layers match the selector (silent-miss prevention).

class physicsnemo.experimental.peft.ApplyResult(
n_wrapped: int,
n_trainable: int,
n_frozen: int,
trainable_names: list[str] = <factory>,
base_fingerprint: str = '',
)[source]#

Summary report from apply_lora() (the model is mutated in place).

physicsnemo.experimental.peft.split_params_for_optimizer(
model: Module,
) → dict[str, list][source]#

Split parameters into {'lora', 'extras', 'frozen'}.

Route lora + extras to AdamW — NOT to optimizers like Muon whose Newton-Schulz orthogonalization is degenerate on low-rank factors. frozen is returned for reporting only.

physicsnemo.experimental.peft.save_adapter(
model: Module,
path: str | Path,
) → None[source]#

Save adapter-only state for a LoRA-wrapped model to path.

The archive is a plain multi-file ZIP (contents below) — load it with load_adapter(), never torch.load or physicsnemo.Module.load. Any file extension is accepted, but a dedicated one such as .lora is recommended: .pt implies torch.load and .mdlus implies Module.load, and neither can read this archive. The model must have been processed by apply_lora.

Archive contents:

  • adapter_config.json — the adapter config (rank, alpha, dropout, init, and an explicit target_modules list of the actually-wrapped names, so it reloads identically regardless of the original selector, including a non-serializable target_filter).

  • adapter_model.pt — the trainable tensors only: lora_A/lora_B and any extras_trainable params (the frozen base is NOT stored).

  • metadata.json — kind="lora_adapter", format/library versions, the base fingerprint, and a summary (n_wrapped, rank, alpha, timestamp).

physicsnemo.experimental.peft.load_adapter(
model: Module,
path: str | Path,
strict: bool = True,
) → None[source]#

Load an adapter into a compatible base model (mutated in place): verify it is a LoRA adapter, check the base fingerprint, re-apply LoRA to the same modules, then load the adapter weights.

Parameters:

strict (bool) – If True (default), a base-fingerprint mismatch raises. If False, it only logs a warning (you assert the base is compatible).

physicsnemo.experimental.peft.merge_lora(
model: Module,
) → Module[source]#

In-place: for each mergeable LoRA-wrapped module, fold its delta into the base weight and replace the wrapper with the (now-updated) underlying nn.Linear / te.Linear. Returns model for chaining.

Non-mergeable adapters (the fused te.LayerNormMLP residual, whose update can’t be folded into the fused weights) are left in place and a warning is logged. After merging, mergeable wrappers are gone, so a model with only those can be saved as a normal .mdlus and served with zero adapter overhead. Idempotent — a second call is a no-op for merged layers.

physicsnemo.experimental.peft.set_adapter_enabled(
model: Module,
enabled: bool,
) → None[source]#

Enable/disable all LoRA deltas in model.

Lets you run a base-only forward (e.g. for an adapter-vs-base comparison) without merging or reloading.

physicsnemo.experimental.peft.print_trainable_parameters(
model: Module,
use_logger: bool = False,
) → str[source]#

Emit a one-line trainable params: N (X% of M total) summary and return the message string.

Extension API#

physicsnemo.experimental.peft.register_lora_wrapper(
layer_type: type,
wrapper_factory: Callable[[...], Module],
) → None[source]#

Register a LoRA wrapper for layer_type.

This is how new architectures (e.g. equivariant, tensor, or MoE layers) plug in without touching the targeting / apply / merge core.

Parameters:

  • layer_type (type) – The base layer class to wrap (e.g. a custom nn.Module subclass). Matched against each module’s MRO, so subclasses are handled too.

  • wrapper_factory (Callable[..., nn.Module]) – Called as wrapper_factory(base_layer, rank=, alpha=, dropout=, init=) and must return an nn.Module that subclasses LoRALayer (see its docstring for the full attribute/method contract). The subclass requirement is enforced by apply_lora: freeze/save/merge identify LoRA layers via isinstance(module, LoRALayer), so a wrapper that does not subclass it would otherwise be silently skipped.

class physicsnemo.experimental.peft.LoRALayer[source]#

Generic LoRA mixin: holds lora_A/lora_B, scaling, dropout and the enable flag, and computes the low-rank delta. Makes no assumption about the base layer’s parameter shapes — combined with a base layer type by the wrapper subclasses.

Math: with lora_A: (in, r) and lora_B: (r, out) the forward adds ((dropout(x) @ A) @ B) * scaling. B is zero at init so the delta is exactly zero — the wrapped forward equals the base forward until trained.

Wrapper contract#

Every LoRA wrapper — the built-ins below and any registered via register_lora_wrapper() — is an nn.Module that subclasses LoRALayer and exposes the surface the apply / freeze / save / merge / enable utilities depend on:

  • Constructor: __init__(self, base_layer, *, rank, alpha, dropout=0.0, init="default"). apply_lora instantiates wrappers as wrapper(base_layer, rank=, alpha=, dropout=, init=) — accept **kwargs if you want to be forward-compatible with options added later.

  • Attributes: base_layer (the wrapped, frozen module); lora_A / lora_B (trainable nn.Parameters — the only params left with requires_grad=True, which is how save_adapter slices the adapter; these may be plain attributes or @propertys that resolve to the Parameters, e.g. for layers whose factors are submodules); enabled (bool toggling the delta); mergeable (bool; False by default — opt in only if you also implement merge_into_base).

  • Methods: forward (adds the low-rank delta to the base output when enabled); merge_into_base (folds the delta into the base weight) — required only when mergeable is True. Optionally override the classmethod is_compatible(base_layer) to veto instances of a registered type this wrapper can’t adapt (defaults to accepting all).

_make_lora_params(...) and lora_delta(...) are optional conveniences for the standard tensor case (2-D lora_A/lora_B with the ((x @ A) @ B) * scaling delta): a wrapper may instead create its own parameters, init, and delta (e.g. an equivariant wrapper whose factors are themselves equivariant layers) as long as it ends up satisfying the contract above. Wrappers for Linear-like bases (.weight shaped (out, in), or exposing in_features/out_features) should subclass _LinearLoRALayer instead — it adds in/out inference at init and a weight-folding merge_into_base. Only generic, non-Linear wrappers (e.g. the fused te.LayerNormMLP residual) inherit LoRALayer directly.

classmethod is_compatible(
base_layer: Module,
) → bool[source]#

Whether this wrapper can adapt base_layer beyond simple type match.

resolve_targets calls this on a selected, registered-type layer before wrapping it and skips the layer if it returns False — letting a wrapper veto instances it can’t actually handle (e.g. an equivariant adapter that needs at least one shared input/output irrep). Defaults to True (every instance of a registered type is adaptable); override in subclasses that need an instance-level check.

lora_delta(x: Tensor) → Tensor[source]#

The low-rank update added to the base output: ((dropout(x) @ lora_A) @ lora_B) * scaling. Zero at init since lora_B starts at zero.

class physicsnemo.experimental.peft.LoRALinear(
base_layer: Linear,
rank: int,
alpha: float,
dropout: float = 0.0,
init: Literal['default'] | Callable[[Tensor], None] = 'default',
)[source]#

Bases: Module, _LinearLoRALayer

LoRA wrapper for torch.nn.Linear.

Wraps a frozen nn.Linear and adds a trainable low-rank update to its output (base(x) + lora_delta(x)). Only lora_A/lora_B train; the base layer’s weight and bias are frozen in place.

Parameters:

  • base_layer (nn.Linear) – The linear layer to wrap; its parameters are frozen.

  • rank (int) – Low-rank dimension r of the adapter.

  • alpha (float) – LoRA scaling numerator; the delta is scaled by alpha / rank.

  • dropout (float, optional) – Dropout applied to the adapter input path. Defaults to 0.0.

  • init (str or callable, optional) – How lora_A is initialized (lora_B is always zero). See LoRAConfig. Defaults to "default" (kaiming_uniform_).

forward(x: Tensor) → Tensor[source]#

Frozen base output plus the LoRA delta (when enabled).

physicsnemo.experimental.peft.is_lora_layer(module: Module) → bool[source]#

Return whether module is a LoRA wrapper.

Parameters:

module (nn.Module) – The module to test.

Returns:

True if module is a LoRALayer instance.

Return type:

bool

physicsnemo.experimental.peft.wrappable_types() → tuple[type, ...][source]#

Return the layer types currently registered as wrappable.

Returns:

The registered base layer types (e.g. nn.Linear and, when available, the Transformer Engine types).

Return type:

tuple[type, …]