mstar.engine.resources.lora.checkpoint

Reading LoRA checkpoints: shared by the static merge and the per-request adapter pool.

Accepted checkpoint layouts (.safetensors):

• diffusers / PEFT: [transformer.]<module>.lora_A.weight + .lora_B.weight (also the base_model.model. prefix), optionally <module>.alpha; without an alpha the scaling comes from an adapter_config.json / config.json next to the weights when there is one (lora_alpha / r, per-module alpha_pattern / rank_pattern, and use_rslora for lora_alpha / sqrt(r), as PEFT applies them), else it is 1 (lora_alpha = r).

• Kohya-style naming lora_down / lora_up for A / B.

• model-native layouts (e.g. the BFL double_blocks.N.img_attn.qkv keys of FLUX.2) through a model-supplied convert_keys hook that rewrites them to the reference module paths.

<module> is a reference-library module path; the model’s checkpoint remap (the same one the weight loader uses) and its stacked-shard rules turn it into the native parameter and, for fused projections, the shard it lands on (the model side’s ModuleLoraMap, and the static merge’s resolve_lora_module).

Functions

load_lora_file(path[, device])
normalize_lora_key(key)	A raw checkpoint key in reference-library naming: known prefixes stripped, Kohya's lora_down / lora_up renamed to lora_A / lora_B.
normalize_lora_state_dict(state_dict[, ...])	Group raw LoRA tensors into per-module (A, B, alpha) triples in reference-library naming.
parse_lora_key(key)	(module, "A" | "B" | "alpha") of a normalized key; raises on anything else.
read_lora_file(path[, convert_keys])	normalize_lora_state_dict over a safetensors file, with its sidecar's scaling.
sidecar_default_alpha(path)	(module, rank) -> alpha reproducing PEFT's scaling for an adapter whose safetensors has no alpha keys.

Classes

LoraAdapter([layers])	module path -> LoraLayer in reference-library naming.
LoraFile(path[, scale])	One adapter checkpoint: a safetensors path and its strength.
LoraLayer(a, b, alpha)	One module's adapter: W += scale * alpha / r * b @ a.
SidecarScaling(lora_alpha, rslora, ...)	(module, rank) -> alpha from a PEFT adapter config: lora_alpha (None: alpha = r), unless an alpha_pattern names the module; rslora folds PEFT's alpha / sqrt(r) into the alpha.

class mstar.engine.resources.lora.checkpoint.LoraAdapter(layers=<factory>)

Bases: object

module path -> LoraLayer in reference-library naming.

Parameters:

layers (dict[str, LoraLayer])

layers: dict[str, LoraLayer]

class mstar.engine.resources.lora.checkpoint.LoraFile(path, scale=1.0)

Bases: object

One adapter checkpoint: a safetensors path and its strength.

Parameters:

• path (str)

• scale (float)

classmethod parse(raw)

Return type:

LoraFile

path: str

scale: float = 1.0

class mstar.engine.resources.lora.checkpoint.LoraLayer(a, b, alpha)

Bases: NamedTuple

One module’s adapter: W += scale * alpha / r * b @ a.

Parameters:

• a (Tensor)

• b (Tensor)

• alpha (float)

a: Tensor

Alias for field number 0

alpha: float

Alias for field number 2

b: Tensor

Alias for field number 1

class mstar.engine.resources.lora.checkpoint.SidecarScaling(lora_alpha, rslora, alpha_pattern, rank_pattern)

Bases: object

(module, rank) -> alpha from a PEFT adapter config: lora_alpha (None: alpha = r), unless an alpha_pattern names the module; rslora folds PEFT’s alpha / sqrt(r) into the alpha. A module a rank_pattern names must be that rank.

Parameters:

• lora_alpha (float | None)

• rslora (bool)

• alpha_pattern (dict)

• rank_pattern (dict)

alpha_pattern: dict

lora_alpha: float | None

rank_pattern: dict

rslora: bool

mstar.engine.resources.lora.checkpoint.load_lora_file(path, device='cpu')

Parameters:

path (str)

Return type:

dict[str, Tensor]

mstar.engine.resources.lora.checkpoint.normalize_lora_key(key)

A raw checkpoint key in reference-library naming: known prefixes stripped, Kohya’s lora_down / lora_up renamed to lora_A / lora_B.

Parameters:

key (str)

Return type:

str

mstar.engine.resources.lora.checkpoint.normalize_lora_state_dict(state_dict, convert_keys=None, default_alpha=None)

Group raw LoRA tensors into per-module (A, B, alpha) triples in reference-library naming.

Parameters:

• state_dict (dict[str, Tensor])

• convert_keys (Callable[[dict[str, Tensor]], dict[str, Tensor]] | None)

• default_alpha (Callable[[str, int], float] | None)

Return type:

LoraAdapter

mstar.engine.resources.lora.checkpoint.parse_lora_key(key)

(module, "A" | "B" | "alpha") of a normalized key; raises on anything else.

Parameters:

key (str)

Return type:

tuple[str, str]

mstar.engine.resources.lora.checkpoint.read_lora_file(path, convert_keys=None)

normalize_lora_state_dict over a safetensors file, with its sidecar’s scaling.

Parameters:

• path (str)

• convert_keys (Callable[[dict[str, Tensor]], dict[str, Tensor]] | None)

Return type:

LoraAdapter

mstar.engine.resources.lora.checkpoint.sidecar_default_alpha(path)

(module, rank) -> alpha reproducing PEFT’s scaling for an adapter whose safetensors has no alpha keys.

PEFT reads lora_alpha (and use_rslora) from the adapter’s config, so load_lora_weights applies lora_alpha / r — or lora_alpha / sqrt(r) for rank-stabilized adapters — where the bare weights would suggest 1; a module an alpha_pattern names takes that alpha instead, and one a rank_pattern names has that rank (which its weights must have). Looks for adapter_config.json then config.json next to path and accepts the training-config spellings too (alpha, lora_rank / rank). None when there is nothing to read.

Parameters:

path (str)

Return type:

Callable[[str, int], float] | None