mstar.model.cosmos3.loader

Weight loading for the Cosmos3 generator backbone and the Edge reasoner’s vision tower.

The published checkpoint is the diffusers transformer/ layout: flat layers.N.* keys with unfused attention projections (to_q/to_k/to_v for the understanding pathway, add_q_proj/add_k_proj/add_v_proj for the generation pathway) and _moe_gen-suffixed GEN MLP/norms. Our backbone module mirrors that layout one-to-one, so loading needs no key remapping and no stacked-parameter fusion. The text lm_head is dropped for generators that never decode text and loaded for checkpoints served as a reasoner.

Edge checkpoints add vision_encoder/model.safetensors (the SigLIP2-style tower under model.visual.* and the patch-merger projector under model.projector.*); load_vision_encoder_weights streams it into the reasoner’s vision module, whose submodules are named visual and projector so only the model. prefix is stripped.

Functions

cosmos3_name_remapper(name[, with_lm_head])	Map a checkpoint key to a backbone parameter path, or None to drop.
load_transformer_weights(model, checkpoint_dir)	Stream the transformer/ shards into model and return loaded keys.
load_vision_encoder_weights(model, ...[, device])	Stream vision_encoder/model.safetensors into the reasoner's vision module (tower + projector) and return the loaded keys, raising if any parameter is left unfilled.
read_transformer_weight_keys(checkpoint_dir)	Return every tensor key declared by the transformer/ shard index.
read_transformer_weight_shapes(checkpoint_dir)	Return {key: shape} for every transformer/ tensor by reading only the safetensors headers — no tensor data is materialized.
read_vision_encoder_weight_shapes(checkpoint_dir)	{remapped key: shape} of vision_encoder/model.safetensors, from the safetensors header only (CPU-side shape verification of the meta-built vision module).
vision_encoder_name_remapper(name)	model.visual.* / model.projector.* -> visual.* / projector.*.

mstar.model.cosmos3.loader.cosmos3_name_remapper(name, with_lm_head=False)

Map a checkpoint key to a backbone parameter path, or None to drop.

Identity for every key the backbone owns; None for the intentional drop-list (lm_head.weight unless the backbone serves the reasoner). Kept explicit so an unexpected checkpoint key surfaces as a coverage failure rather than being silently ignored.

Parameters:

• name (str)

• with_lm_head (bool)

Return type:

str | None

mstar.model.cosmos3.loader.load_transformer_weights(model, checkpoint_dir, device='cpu')

Stream the transformer/ shards into model and return loaded keys.

Mirrors the meta-device + load_hf_weights path the other model packages use, but resolves the shard list from the diffusers diffusion_pytorch_model index (the generic iterator only knows the model.safetensors name). No stacked-parameter rules: the checkpoint’s projections are unfused and match the backbone parameter names directly. Raises if any backbone parameter is left unfilled — the completeness guarantee bagel’s loader also enforces.

Parameters:

• model (Module)

• checkpoint_dir (str | Path)

• device (str)

Return type:

set[str]

mstar.model.cosmos3.loader.load_vision_encoder_weights(model, checkpoint_dir, device='cpu')

Stream vision_encoder/model.safetensors into the reasoner’s vision module (tower + projector) and return the loaded keys, raising if any parameter is left unfilled.

Parameters:

• model (Module)

• checkpoint_dir (str | Path)

• device (str)

Return type:

set[str]

mstar.model.cosmos3.loader.read_transformer_weight_keys(checkpoint_dir)

Return every tensor key declared by the transformer/ shard index.

Parameters:

checkpoint_dir (str | Path)

Return type:

set[str]

mstar.model.cosmos3.loader.read_transformer_weight_shapes(checkpoint_dir)

Return {key: shape} for every transformer/ tensor by reading only the safetensors headers — no tensor data is materialized. Enables CPU-side shape verification of the meta-built backbone against the checkpoint.

Parameters:

checkpoint_dir (str | Path)

Return type:

dict[str, tuple[int, …]]

mstar.model.cosmos3.loader.read_vision_encoder_weight_shapes(checkpoint_dir)

{remapped key: shape} of vision_encoder/model.safetensors, from the safetensors header only (CPU-side shape verification of the meta-built vision module).

Parameters:

checkpoint_dir (str | Path)

Return type:

dict[str, tuple[int, …]]

mstar.model.cosmos3.loader.vision_encoder_name_remapper(name)

model.visual.* / model.projector.* -> visual.* / projector.*.

Parameters:

name (str)

Return type:

str | None