mstar.model.glm52.weight_loader

HF GLM-5.2 checkpoint loading.

Functions

build_glm52_read_plan(checkpoint_keys, ...)	Keys-to-read + per-key slice specs for the TP fast read path.
build_glm52_stacked_params(n_routed_experts)
glm52_name_remapper(name)
load_glm52_hf_weights(module, weights, ...)
load_weights(module, source[, device])
restore_fp32_params(module)	Re-widen params the checkpoint stores fp32 before loading into them.
skip_phase_b_keys(weights, num_hidden_layers)	Drop MTP-layer keys — and, with load_indexer=False, indexer keys — before any dequant buffering.

mstar.model.glm52.weight_loader.build_glm52_read_plan(checkpoint_keys, config, tp_rank, tp_size, load_indexer=None)

Keys-to-read + per-key slice specs for the TP fast read path.

load_indexer defaults to config.dsa_long_context: flag-off the model builds no indexer to fill, so its keys are never read.

Parameters:

• checkpoint_keys (Iterable[str])

• tp_rank (int)

• tp_size (int)

• load_indexer (bool | None)

Return type:

tuple[set[str], dict[str, tuple[int, int, int]]]

mstar.model.glm52.weight_loader.build_glm52_stacked_params(n_routed_experts, fp8_experts=False)

Parameters:

• n_routed_experts (int)

• fp8_experts (bool)

Return type:

list[StackedParamRule]

mstar.model.glm52.weight_loader.glm52_name_remapper(name)

Parameters:

name (str)

Return type:

str | None

mstar.model.glm52.weight_loader.load_glm52_hf_weights(module, weights, n_routed_experts, quant_config=None, fp8_experts=False, num_hidden_layers=78, load_indexer=True)

Parameters:

• module (nn.Module)

• weights (Iterable[tuple[str, torch.Tensor]])

• n_routed_experts (int)

• quant_config (Fp8BlockQuantConfig | None)

• fp8_experts (bool)

• num_hidden_layers (int)

• load_indexer (bool)

Return type:

set[str]

mstar.model.glm52.weight_loader.load_weights(module, source, device='cpu')

Parameters:

• module (Module)

• source (str | Path)

• device (device | str)

Return type:

set[str]

mstar.model.glm52.weight_loader.restore_fp32_params(module)

Re-widen params the checkpoint stores fp32 before loading into them.

Parameters:

module (Module)

Return type:

None

mstar.model.glm52.weight_loader.skip_phase_b_keys(weights, num_hidden_layers, load_indexer=True)

Drop MTP-layer keys — and, with load_indexer=False, indexer keys — before any dequant buffering.

Parameters:

• weights (Iterable[tuple[str, Tensor]])

• num_hidden_layers (int)

• load_indexer (bool)

Return type:

Iterator[tuple[str, Tensor]]