sgl-project/sglang · critical · RuntimeError

Pi05 weight load failed: {len(missing)} missing weights, {mi

Error message

Pi05 weight load failed: {len(missing)} missing weights, {mismatched} mismatched weights. Running a robot policy with uninitialized or partially loaded weights is unsafe.

What it means

Pi05Policy._load_weights copies a checkpoint state_dict into the model and then verifies that every expected parameter was loaded with a matching shape; it reports counts of missing and shape-mismatched weights. Because running a robot policy with random or partially initialized weights produces dangerous physical actions, this check is a hard RuntimeError, not a warning. It runs during __init__/from_pretrained.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py:771

                    target_params,
                )
                if target_weight is None:
                    unexpected += 1
                    continue
                target_key, shard_id = target_weight
                target = self._target_tensor_for_key(
                    target_key,
                    target_state,
                    target_params,
                )
                if not self._load_tensor_to_target(target, tensor, shard_id):
                    mismatched += 1
                    continue
                loaded_keys.add(target_key)

        missing = [key for key in target_state if key not in loaded_keys]
        if missing or mismatched:
            raise RuntimeError(
                f"Pi05 weight load failed: {len(missing)} missing weights, "
                f"{mismatched} mismatched weights. Running a robot policy with "
                "uninitialized or partially loaded weights is unsafe."
            )
        if unexpected:
            logger.warning(
                "Pi05 weight load: %d loaded, %d unexpected",
                len(loaded_keys),
                unexpected,
            )
        else:
            logger.info("Pi05 weights loaded successfully")

    def build_prefix_cache_key(
        self,
        observation: VLAObservationBatch,
    ) -> str:
        camera_order = tuple(observation.metadata.get("camera_order", ()))

View on GitHub (pinned to 0132848349)

Solutions

1. Verify all checkpoint shards are present in the model path (incomplete downloads are the #1 cause)
2. Confirm the model config (variant, action_dim) matches the checkpoint — re-download the original config.json
3. If you finetuned and intentionally dropped modules, export the full state_dict instead
4. Log missing/mismatched key lists before the raise (or in a debugger) to see exactly which modules fail

Example fix

# before
# incomplete download: only model-00001-of-00002.safetensors present
policy = Pi05Policy.from_pretrained("./pi05_ckpt")

# after
# complete both shards, then:
policy = Pi05Policy.from_pretrained("./pi05_ckpt")

import os, glob, torch, json
shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
assert shards, "no safetensors found"
from safetensors import safe_open
ckpt_keys = set()
for s in shards:
    with safe_open(s, framework="pt") as f:
        ckk_keys |= set(f.keys())
# quick sanity: every ckpt key starts with a known Pi05 module prefix
bad = [k for k in ck_keys if not k.startswith(("vision_tower.", "language_model.", "noise_proj.", "action_in_proj.", "action_out_proj."))]
assert not bad, f"foreign keys: {bad[:5]}"

try:
    policy = Pi05Policy.from_pretrained(path)
except RuntimeError as e:
    if "Pi05 weight load failed" in str(e):
        # do NOT fall back to partial weights for a robot policy
        raise SystemExit(
            f"Checkpoint incomplete/mismatched at {path}. "
            "Re-download or fix config. Details: {e}"
        ) from e
    raise

When it happens

Trigger: Loading a checkpoint that doesn't match the model config: missing keys (partial finetune exports, renamed modules), or shape mismatches (e.g. different action dimension, different Gemma variant depth, or a checkpoint from another Pi0/Pi05 architecture).

Common situations: Checkpoint/model version skew after upgrading sglang; finetuned checkpoints that omit vision-tower weights; editing model config (action_dim, num_steps) so shapes no longer line up; mixed safetensors shards not all present in the directory.

Related errors

• shard_offset and shard_size must be provided
• Unsupported Pi05 dtype: {dtype_name}
• qkv_proj weight {name}: unexpected shape {tuple(loaded_weigh
• scale_shift_table must have shape [9, D]
• expected a tensor with at least one dimension

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/801db6b60513ef2a. Report an issue: GitHub.