inspect_config.py — Model parameter accounting

scripts/inspect_config.py answers a static architecture-budget question: how many unique stored parameters does this configured model have, and how many are on the selected-expert path for one token? It builds the real module graph but performs no forward pass, optimization or data loading. By default, construction occurs on PyTorch’s meta device so parameter tensors have shapes without allocating their numerical storage. The counted module graph belongs to model.py and stable_latent_moe.py.

Code excerpts below are dedented from their enclosing functions for readability; the displayed statements and ordering are preserved.

Invocation and configuration

The CLI requires the positional model YAML config. --device {meta,cpu} defaults to meta; --by-layer adds one row for each decoder layer; --show-config prints every normalized KimiK3Config dataclass field before the report. There is no CUDA option and no checkpoint input.

A read-only example from the repo root uses the existing model configuration:

python scripts/inspect_config.py configs/model/kimi_k3_300m.yaml --by-layer

The YAML is parsed using load_config() and the normal model constructor. The script reports what this implementation would instantiate, rather than independently estimating dimensions from a paper table. --device cpu constructs real initialized parameters in RAM and can be expensive for very large configs; meta avoids weight storage but still builds Python module objects and shapes.

Count unique parameter objects, not repeated references

unique_numel() counts the number of elements of each distinct Parameter object once:

for p in parameters:
    pid = id(p)
    if pid not in seen:
        seen.add(pid)
        total += p.numel()

This matters for tied token embeddings and LM head: the same parameter can be reachable through two module paths but should occupy one parameter budget. The routine deduplicates by Python identity, not by numerical value or by shared storage between distinct objects. It uses numel(), independent of dtype and device. module_params() applies it to any module’s parameter iterator.

Count routed experts as active only when selected

The MoE helper expects the project’s current StableLatentMoE representation: routed experts are a nn.ModuleList named experts, and top_k is defined on the MoE module. Nonmatching modules contribute zero routed-expert parameters. Equal expert sizes are an explicit invariant; unequal sizes raise rather than produce a misleading total:

sizes = [module_params(expert) for expert in experts]
if len(set(sizes)) != 1:
    raise RuntimeError(
        "Routed experts have unequal parameter counts; "
        "the active-parameter estimate needs updating."
    )

With E routed experts, k = top_k, and P parameters per routed expert, the helper returns (E*P, k*P). The total parameter count includes all experts; the per-token active estimate subtracts (E-k)*P for each MoE layer. Both router weights and shared experts remain counted as active. Dense FFNs, attention and embeddings are unchanged. This is a count of structurally active model parameters for a token, not an exact count of floating-point operations, unique tensors touched over a batch, or memory allocated during inference.

The model-level count subtracts inactive routed experts directly from the unique model total. The per-layer helper instead walks each layer’s child modules and subtracts inactive routed-expert elements, so the same rule underlies both summaries.

Categorize parameters without double-counting tying

classify_parameter() assigns names to buckets using module paths and the loaded configuration. Its layers.<index>.attn entries are labeled KDA or Gated MLA according to layer_types[index], while layers.<index>.ffn is Dense FFN for layers before num_dense_layers and Stable LatentMoE thereafter. Other decoder-layer parameters are classed as decoder norms/AttnRes; final norm and output AttnRes have their own bucket. Embedding/LM head tying is captured by the embed_tokens path. A defensive lm_head bucket exists if the implementation ever changes tying.

The breakdown iterates model.named_parameters(remove_duplicate=True) for unique names, and subtracts only inactive MoE routed experts from the Stable LatentMoE active bucket. Bucket labels are a view over parameter names; they do not represent separately instantiated submodels. In particular, the output is sensitive to module naming and num_dense_layers. If model structure changes, the classification rules must be reviewed.

Formatting and printed sections

print_report() prints the architecture summary (vocabulary, width, layer mix, expert counts, max context and KDA settings), the parameter budget and component breakdown. It reports four different counts:

The Active / total percentage uses the exact integer counts internally. Counts are displayed both as rounded K/M/B suffixes and exact comma-separated totals for main budget rows. The component % total column uses total parameters in that bucket, not active parameters. --by-layer prints index, attention type, Dense/MoE marker, total and active elements for each layer, but does not allocate or score a token.

The “Parameter-only memory” block multiplies total parameter elements by four bytes for FP32 or two for BF16/FP16 and presents MiB/GiB using powers of 1024. It is a storage-only estimate: no gradients, optimizer slots, activations, allocator overhead, temporary buffers or shared memory. It is not a claim that an entire model will fit in that amount of device RAM.

The last entry point optionally prints normalized dataclass fields, constructs with torch.device("meta") or on CPU, prints the config path and invokes the report. There are no writes or generated artifacts; stdout is the only output. Invalid YAML/configuration or a model construction failure surfaces as an exception.

Verification and boundaries

The source’s correctness depends principally on KimiK3 and StableLatentMoE. The inspected test_recipe.py asserts that optimizer groups cover tied weights once, a related deduplication invariant. The test_train_cli.py real-model case checks the selected architecture configuration in a saved run. No direct inspect_config.py snapshot/count test was identified in the inspected suite; this page does not claim one or that any parameter report was executed.