dew.lora

Adapt a variables tree with low-rank deltas (LoRA, arXiv 2106.09685).

An adapter is a set of rank-r deltas on the kernels of Dense modules. Beside a target kernel W, [in..., out...], the tree holds lora_A, [in..., r], and lora_B, [r, out...], and the module computes x W + scale * (x A) B with scale = alpha / r (alpha / sqrt(r) for rsLoRA, arXiv 2312.03732). Merged, scale * A B is added into the kernel and the factors are gone. A target is the path of its module’s leaves in the tree, ("params", "layers_0", "self_attn", "q_proj"), so one description serves every model and every collection.

One object carries all of it. LoRA.fresh draws an adapter on the projections a model binds, LoRA.load reads one off disk onto them, and the adapter that comes back adapts the model (adapt), says what trains (trainable), folds itself in (merge) and writes itself out (save), since it holds the bindings it was built over.

LoRA.adapt makes a model compute the branch without a module of its own for every projection: it wraps apply and init in a Flax method interceptor that adds the branch to each target Dense’s output and reads or creates the factors as that module’s own parameters. What it hides is the per-module bookkeeping PEFT does with wrapper layers: the factor shapes of a DenseGeneral with several contracted axes, the compute dtype the kernel path uses, dropout on the branch input, and the parameter naming the merge, the export and the trainable filter agree on.

Two files are read and written, both what the references produce natively. PEFT’s directory (adapter_config.json, adapter_model.safetensors, keys base_model.model.<module>.lora_A.weight) is what Transformers loads; the Diffusers file (pytorch_lora_weights.safetensors, keys <component>.<module>.lora_A.weight, the PEFT config per component in the header’s lora_adapter_metadata) is what a pipeline’s load_lora_weights reads. Kohya/sgm keys (lora_unet_..., .alpha tensors) are not accepted. Source module names resolve to tree paths through Pretrained.layouts, the bindings a source’s export runs backwards, so an adapter is placed exactly where the base tensor it modifies went. An adapter attaches to a model and its variables, not to a loader: a model built from the registry passes no layouts and bound_layouts reads the names and shapes off its own kernels, so a run that never touched a published checkpoint adapts the same way. RunConfig.lora is that path from a config: the run adapts the module its objective trains and freezes everything but the factors.

Name	Summary
PEFT_CONFIG
PEFT_WEIGHTS
PEFT_PREFIX
DIFFUSERS_WEIGHTS
DIFFUSERS_METADATA
FACTORS
INIT_A
INIT_B
Target	Holds one adapted kernel’s rank and alpha.
LoRA	Says which kernels carry a low-rank delta, and how the delta scales.
bound_layouts	Return the projections an adapter can bind on model, by the name a file uses.
PeftConfig	Describes adapter_config.json as PEFT writes and reads it.
Adaptable	Declares an objective an adapter attaches to.
attach	Adapt the module objective trains and freeze all but the factors.

PEFT_CONFIG

attribute source

PEFT_CONFIG = 'adapter_config.json'

PEFT_WEIGHTS

attribute source

PEFT_WEIGHTS = 'adapter_model.safetensors'

PEFT_PREFIX

attribute source

PEFT_PREFIX = 'base_model.model.'

DIFFUSERS_WEIGHTS

attribute source

DIFFUSERS_WEIGHTS = 'pytorch_lora_weights.safetensors'

DIFFUSERS_METADATA

attribute source

DIFFUSERS_METADATA = 'lora_adapter_metadata'

FACTORS

attribute source

FACTORS = ('lora_A', 'lora_B')

INIT_A

attribute source

INIT_A = nn.initializers.variance_scaling(1 / 3, 'fan_in', 'uniform')

INIT_B

attribute source

INIT_B = nn.initializers.zeros

Target

dataclass source

class Target(rank: int, alpha: float)

Holds one adapted kernel’s rank and alpha.

LoRA

dataclass source

class LoRA(
    targets: Mapping[Path, Target],
    rslora: bool = False,
    dropout: float = 0.0,
    layouts: Mapping[str, WeightLayout] = dict(),
)

Says which kernels carry a low-rank delta, and how the delta scales.

layouts: Mapping[str, WeightLayout]

The bindings the targets were bound over, by the name a file writes them under, which is what save writes the factors back through. fresh and load fill it with the source’s own projections; an adapter a run config declares carries its targets alone, so it compares by them and cannot save itself. A run record leaves it out: the source the record names is where the bindings come from.

LoRA.scale

def scale(target: Target) -> float

LoRA.trainable

def trainable(path: Path) -> bool

Return the adapter’s own leaves: the PathFilter a partial run trains.

LoRA.target_at

def target_at(path: Path) -> Target | None

Return the target a module path names, seen through the stack that runs it: a scanned run’s module layers_3_7 stands for layers 3 through 7, whose targets must agree, since the run’s kernels share one stacked factor pair.

LoRA.adapt

def adapt(model: nn.Module, root: Path = ('params',)) -> nn.Module

Return model computing the adapter branch in every target module.

The result is an instance of a subclass of model’s class with the same fields and methods, so it builds, checks and generates as the model does; only apply and init change, running under the interceptor (Flax’s init dispatches through init_with_output). Adapt the module that is applied: a target inside a submodule is reached through its parent’s apply. root is where the model’s own params sit in the tree the targets are named in: the whole tree’s params for a model applied on it, deeper for a tower applied on a subtree.

LoRA.merge

def merge(variables: Variables) -> Variables

Return variables with every delta added into its kernel and the factors removed.

The sum runs in at least fp32 at full precision and lands in the kernel’s dtype, PEFT’s merge_and_unload.

LoRA.fresh

def fresh(
    model: nn.Module,
    variables: Variables,
    layouts: Mapping[str, WeightLayout],
    *,
    rank: int,
    modules: Sequence[str],
    key: jax.Array,
    alpha: float | None = None,
    rslora: bool = False,
    dropout: float = 0.0,
) -> tuple[LoRA, Variables]

Build a new adapter on the projections modules name, and add its factors.

modules are PEFT’s target_modules: a projection matches when its name relative to the model (model.layers.0.self_attn.q_proj, or to_q under a pipeline component) is the entry or ends in . and the entry. An entry that matches no projection is refused. alpha unset is PEFT’s own default for a rank, twice it. A is drawn from key the way PEFT draws it and B is zero, so the fresh adapter is the identity.

LoRA.load

def load(
    model: nn.Module,
    variables: Variables,
    layouts: Mapping[str, WeightLayout],
    path: str | FilePath,
) -> tuple[LoRA, Variables]

Load an adapter for model and variables with the factors in place.

layouts are the bindings a file’s module names resolve through: Pretrained.layouts for a loaded source, an empty mapping for a model built from the registry, whose own module paths are its names.

path is a PEFT adapter directory or a Diffusers file (or the directory holding one). Targets the model does not bind, tensors whose shapes do not fit the bound weight, ranks that disagree with the config, and PEFT features this loader does not carry are refused by name.

LoRA.save

def save(variables: Variables, path: str | FilePath) -> None

Write the adapter’s factors from variables under its own module names.

The names are the ones this adapter bound at construction, so a run saves what it trained with the tree it trained it in. One unnamed component writes PEFT’s directory, which is a decoder source and a registry-built model; a pipeline source, whose weights are named under several components, writes the Diffusers file with each component’s PEFT config in its header.

bound_layouts

function source

def bound_layouts(
    model: nn.Module,
    variables: Variables,
    layouts: Mapping[str, WeightLayout],
) -> Mapping[str, WeightLayout]

Return the projections an adapter can bind on model, by the name a file uses.

A loaded source publishes Pretrained.layouts, the bindings its export runs backwards, and those names are the ones PEFT and Diffusers write. A model built from the registry has no published file, so its own module path under params is the name and its kernel in variables is the shape; that is the mapping this derives when none is given.

A derived layout covers a two-axis kernel, which stores [in, out] and exports as PEFT’s [out, in]. A kernel of more axes splits into contracted and feature axes that the module decides and the tree does not record, so it carries no derived name and a target that asks for it is refused as unbound.

PeftConfig

class source

class PeftConfig(TypedDict)

Describes adapter_config.json as PEFT writes and reads it.

It carries the defaults every target takes and the per-module exceptions. fan_in_fan_out, bias, init_lora_weights and inference_mode are fixed because dew builds its adapters one way. Nothing here reads those four back; they are written so a PEFT reader finds the keys it expects.

Adaptable

class source

class Adaptable(Protocol)

Declares an objective an adapter attaches to.

It trains one module, which is its model, and it takes the filter that says which of that module’s leaves the optimizer moves. LMObjective and BlockDiffusionObjective are the two; an objective that keeps no model or selects what trains some other way is refused by name.

attach

function source

def attach(objective: object, adapter: LoRA) -> None

Adapt the module objective trains and freeze all but the factors.

RunConfig.train calls this once, after a recipe has built the objective and before anything initialises it, so the adapted module is what the run traces and the adapter’s own leaves are the only ones the optimizer moves. The adapted module is a subclass of the same class with the same fields, so what the objective read off the model at construction still holds.