mstar.engine.resources.lora.store

mstar.engine.resources.lora.store#

The host tier of the adapter pool: adapters read, sharded for this rank, and pinned.

An adapter’s tensors are laid out exactly like one device slot (rank padded to the pool’s max_rank with zeros), so moving it onto the device is one copy per target.

Functions

adapter_targets(spec, path)

The targets of spec the checkpoint at path holds LoRA weights for.

as_linear(shape)

A LoRA factor's shape as a linear's: a conv adapter's A [r, c, *kernel] and B [out, r, 1, ...] flatten to [r, c * prod(kernel)] and [out, r], which is the delta on a conv whose kernel is its stride, over the flattened patch.

check_adapter_file(spec, name, path[, others])

The adapter at path checked against spec and the model's others (its LoRA specs for other nodes, a vision tower's say, which serve the same adapters), each from the checkpoint's header; returns its digest.

check_adapters(spec[, others])

Every adapter spec serves, through check_adapter_file (against others too); returns each one's digest.

inspect_adapter(spec, path)

Rank, unmapped and misfitting modules of the checkpoint at path.

inspect_keys(spec, items)

inspect_adapter over (checkpoint key, shape) pairs, as a header lists them.

misfit(spec, module, a_shape, b_shape)

Why a mapped module's A / B shapes do not fit its target, or None.

Classes

LoadedAdapter(name, digest, rank, scaling, a, b)

One adapter, ready to copy into a slot.

LoraAdapterStore(spec[, rank, world_size, ...])

Adapters held in pinned host memory.
class mstar.engine.resources.lora.store.LoadedAdapter(name, digest, rank, scaling, a, b)[source]#

Bases: object

One adapter, ready to copy into a slot.

a[target] is [target layers, slices * max_rank, in] and b[target] is [target layers, out, max_rank], both already this rank’s shard. A target the adapter does not touch is absent. digest is adapter_digest of the unsharded weights, so every rank agrees on it, computed as the API server computes the one it roots the adapter’s prefix-cache keys in.

Parameters:
a: dict[str, Tensor]#
b: dict[str, Tensor]#
digest: bytes#
name: str#
rank: int#
scaling: float#
class mstar.engine.resources.lora.store.LoraAdapterStore(spec, rank=0, world_size=1, pin=None, dtype=None)[source]#

Bases: object

Adapters held in pinned host memory.

Which adapters it holds follows from the hold calls alone, in their order: each makes an adapter the most recent, dropping the least recent past capacity but none the call’s keep names. A held adapter’s copy arrives later, by put, once read. So when a read lands or fails never changes what is held, and ranks that make the same calls hold the same set.

Parameters:
drop(name)[source]#
Parameters:

name (str)

Return type:

None

get(name)[source]#
Parameters:

name (str)

Return type:

LoadedAdapter | None

hold(name, keep=())[source]#

Make name the most recent adapter held, dropping the least recent past capacity, but none named in keep (which may hold the tier past capacity for a while). True when it was not held, so its copy is the caller’s to read in.

Parameters:
Return type:

bool

load(name, path, digest=None)[source]#

Hold name and read it from path unless its copy is in.

Parameters:
Return type:

LoadedAdapter

put(adapter)[source]#

adapter’s copy, kept if it is still held.

Parameters:

adapter (LoadedAdapter)

Return type:

bool

read(name, path, digest=None)[source]#

path read, checked against the spec and sharded for this rank, without holding it: a caller can read outside whatever lock guards put. With digest, weights that digest otherwise are refused: the file changed since the adapter was checked and keyed, and serving it would put other weights under the same prefix-cache root, or different weights on different ranks.

Parameters:
Return type:

LoadedAdapter

mstar.engine.resources.lora.store.adapter_targets(spec, path)[source]#

The targets of spec the checkpoint at path holds LoRA weights for.

Parameters:
Return type:

set[str]

mstar.engine.resources.lora.store.as_linear(shape)[source]#

A LoRA factor’s shape as a linear’s: a conv adapter’s A [r, c, *kernel] and B [out, r, 1, ...] flatten to [r, c * prod(kernel)] and [out, r], which is the delta on a conv whose kernel is its stride, over the flattened patch.

Return type:

tuple[int, …]

mstar.engine.resources.lora.store.check_adapter_file(spec, name, path, others=())[source]#

The adapter at path checked against spec and the model’s others (its LoRA specs for other nodes, a vision tower’s say, which serve the same adapters), each from the checkpoint’s header; returns its digest. The digest is read once per way the specs convert the checkpoint (convert_keys), so once for specs that read it alike. Raises ValueError for an adapter the deployment cannot serve.

Parameters:
Return type:

bytes

mstar.engine.resources.lora.store.check_adapters(spec, others=())[source]#

Every adapter spec serves, through check_adapter_file (against others too); returns each one’s digest.

Raises naming every adapter that fails, so a bad adapter in the serving config stops the deployment at startup, the same whether or not a worker would have preloaded it.

Parameters:
Return type:

dict[str, bytes]

mstar.engine.resources.lora.store.inspect_adapter(spec, path)[source]#

Rank, unmapped and misfitting modules of the checkpoint at path.

Reads only the safetensors header unless the model converts a native layout, which needs the tensors.

Parameters:
Return type:

AdapterInfo

mstar.engine.resources.lora.store.inspect_keys(spec, items)[source]#

inspect_adapter over (checkpoint key, shape) pairs, as a header lists them.

Parameters:

spec (LoraSpec)

Return type:

AdapterInfo

mstar.engine.resources.lora.store.misfit(spec, module, a_shape, b_shape)[source]#

Why a mapped module’s A / B shapes do not fit its target, or None.

Parameters:
Return type:

str | None