Skip to content
mini-sglang from Scratch Stage 1 · Computing it right: from skeleton to real structure

Model and weight loading

With the layers in place, this chapter assembles them into a complete model: pull the fields the engine cares about out of Hugging Face's config.json, create the model by architecture name, then use a streaming loader that reads the weights out of safetensors while merging and sharding them on the way in. Four architectures are supported: Llama, Qwen2, Qwen3 and Qwen3-MoE.

Self-test: if you can answer these, skip the chapter

  1. How do the decoder layers of Llama, Qwen2 and Qwen3 differ?
  2. Why does the weight loader stream? What is wrong with one torch.load of the whole checkpoint?
  3. q_proj, k_proj and v_proj are stored separately in the checkpoint. How are they merged into qkv_proj while loading? Does the order matter?
  4. Where does RoPE's rope_theta live in a transformers 4.x config, and where in a 5.x one?
Answers (try it yourself first, then expand)
  1. The structure is the same and only three switches differ: Qwen2's q, k and v projections have biases; Qwen3 applies QK-Norm to q and k before RoPE (and has no qkv bias); Qwen3-MoE replaces the MLP with an MoE.
  2. One torch.load reads the whole checkpoint into memory, and the peak can be several times the model (the raw weights plus the sharded and merged copies); streaming shards and merges as it reads, so the peak stays small.
  3. Once q, k and v have each been read, they are sharded (under tensor parallelism) and then concatenated along the output dimension into one qkv_proj weight. The order matters: it must match the order in which the model splits qkv_proj's output back into q, k and v, or they end up swapped.
  4. transformers 4.x puts it at the top level of the config (rope_theta); 5.x moves it into a rope_parameters dictionary. ModelConfig has to accept both.

Files you will write: models/config.py, models/base.py, models/utils.py, models/decoder.py, models/register.py, models/weight.py.

Files so far: 28 (3 new, 0 modified, 0 brought in early)
minisgl/
├── __init__.py
├── core.py
├── distributed/  (3 files, unchanged)
├── kernel/  (2 files, unchanged)
├── layers/  (8 files, unchanged)
├── models/
│   ├── config.py
│   ├── base.py
│   ├── utils.py
│   ├── decoder.py
│   ├── __init__.py  ← new in this chapter
│   ├── register.py  ← new in this chapter
│   ├── weight.py  ← new in this chapter
├── utils/  (6 files, unchanged)
This step's main: examples/ch03_models.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

ModelConfig: only the fields we need

A Hugging Face config class has dozens of fields, and different models use different names and defaults. ModelConfig collapses the fields the engine cares about into one frozen dataclass:

minisgl/models/config.py · ModelConfig
@dataclass(frozen=True)
class ModelConfig:
    """Pull the fields the engine cares about out of Hugging Face's config.json."""

    num_layers: int
    num_qo_heads: int
    num_kv_heads: int
    head_dim: int
    hidden_size: int
    vocab_size: int
    intermediate_size: int
    rms_norm_eps: float
    rotary_config: RotaryConfig
    hidden_act: str
    tie_word_embeddings: bool
    attention_bias: bool
    model_type: str
    architectures: List[str]
    # MoE only (0 for a dense model)
    num_experts: int = 0
    num_experts_per_tok: int = 0
    moe_intermediate_size: int = 0
    norm_topk_prob: bool = False
    mlp_only_layers: List[int] = field(default_factory=list)

    @property
    def is_moe(self) -> bool:
        return self.num_experts > 0

    @classmethod
    def from_hf(cls, config: PretrainedConfig) -> ModelConfig:
        num_heads = config.num_attention_heads
        head_dim = getattr(config, "head_dim", None) or config.hidden_size // num_heads
        # transformers 5 merged rope_theta and rope_scaling into rope_parameters; older versions keep them apart
        rope = dict(getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {})
        base = getattr(config, "rope_theta", None) or rope.get("rope_theta", 10000.0)
        scaling = rope if rope.get("rope_type", "default") != "default" else None
        model_type = getattr(config, "model_type", "llama")
        return cls(
            num_layers=config.num_hidden_layers,
            num_qo_heads=num_heads,
            num_kv_heads=getattr(config, "num_key_value_heads", None) or num_heads,
            head_dim=head_dim,
            hidden_size=config.hidden_size,
            vocab_size=config.vocab_size,
            intermediate_size=config.intermediate_size,
            rms_norm_eps=config.rms_norm_eps,
            rotary_config=RotaryConfig(
                head_dim=head_dim,
                rotary_dim=head_dim,
                max_position=config.max_position_embeddings,
                base=float(base),
                scaling=scaling,
            ),
            hidden_act=config.hidden_act,
            tie_word_embeddings=bool(getattr(config, "tie_word_embeddings", False)),
            # Qwen2's q/k/v projections carry a bias, but its config has no attention_bias field
            attention_bias=bool(getattr(config, "attention_bias", model_type == "qwen2")),
            model_type=model_type,
            architectures=list(getattr(config, "architectures", None) or ["LlamaForCausalLM"]),
            num_experts=getattr(config, "num_experts", 0) or 0,
            num_experts_per_tok=getattr(config, "num_experts_per_tok", 0) or 0,
            moe_intermediate_size=getattr(config, "moe_intermediate_size", 0) or 0,
            norm_topk_prob=bool(getattr(config, "norm_topk_prob", False)),
            mlp_only_layers=list(getattr(config, "mlp_only_layers", None) or []),
        )
A few easy traps:

  • head_dim: Qwen3's head_dim (128) is not hidden_size / num_attention_heads (1024 / 16 = 64), so the explicit field has to win;
  • RoPE parameters: transformers 5 merged rope_theta and rope_scaling into a rope_parameters dictionary ({"rope_theta": 1000000, "rope_type": "default"}), while 4.x keeps them as two separate fields. from_hf accepts both;
  • Qwen2's biases: Qwen2's q/k/v projections have biases, but its config has no attention_bias field, so it has to be decided from model_type.

examples/ch03_models.py
"""Chapter 3: from config.json to a model that can compute."""

import time

import torch
from minisgl.distributed import set_tp_info
from minisgl.layers import set_rope_device
from minisgl.models import ModelConfig, create_model, load_weight
from minisgl.utils import cached_load_hf_config, torch_dtype

set_tp_info(0, 1)
set_rope_device(torch.device("cpu"))
config = ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B"))
print(f"{config.num_layers} layers, {config.num_qo_heads} q heads, {config.num_kv_heads} kv heads, "
      f"head_dim {config.head_dim}, hidden {config.hidden_size}, vocab {config.vocab_size}")
print(f"RoPE base {config.rotary_config.base:.0f}, tied embeddings {config.tie_word_embeddings}, "
      f"architecture {config.architectures[0]}")

t = time.perf_counter()
with torch.device("meta"), torch_dtype(torch.bfloat16):
    model = create_model(config)
print(f"building the model on the meta device: {(time.perf_counter() - t) * 1e3:.0f} ms, "
      f"{len(model.state_dict())} weight tensors")

t = time.perf_counter()
weights = {}
for i, (name, tensor) in enumerate(load_weight("models/Qwen3-0.6B", torch.device("cpu"))):
    if i < 6:
        print(f"  yield {name:<48} {tuple(tensor.shape)}")
    weights[name] = tensor
print(f"streaming {len(weights)} tensors: {time.perf_counter() - t:.2f} s")
model.load_state_dict(weights)  # lm_head.weight is dropped by the output layer with tied embeddings
w = model.model.layers.op_list[0].self_attn.qkv_proj.weight
print("qkv_proj after loading:", tuple(w.shape), w.dtype, w.device)
What it prints
28 layers, 16 q heads, 8 kv heads, head_dim 128, hidden 1024, vocab 151936
RoPE base 1000000, tied embeddings True, architecture Qwen3ForCausalLM
building the model on the meta device: 22 ms, 226 weight tensors
  yield lm_head.weight                                   (151936, 1024)
  yield model.embed_tokens.weight                        (151936, 1024)
  yield model.layers.0.input_layernorm.weight            (1024,)
  yield model.layers.0.mlp.down_proj.weight              (1024, 3072)
  yield model.layers.0.mlp.gate_up_proj.weight           (6144, 1024)
  yield model.layers.0.post_attention_layernorm.weight   (1024,)
streaming 227 tensors: 0.16 s
qkv_proj after loading: (4096, 1024) torch.bfloat16 cpu

The model structure

The decoder layers of the four architectures are nearly identical, differing only in three switches:

Architecture q/k/v bias QK-Norm MLP
Llama depends on attention_bias (usually none) no SwiGLU
Qwen2 yes no SwiGLU
Qwen3 no yes SwiGLU
Qwen3-MoE no yes MoE (chapter 20)

So we use one generic decoder.py:

minisgl/models/decoder.py · CausalLM
class CausalLM(BaseLLMModel):
    has_qk_norm = False
    moe = False

    def __init__(self, config: ModelConfig):
        self.model = DecoderModel(config, has_qk_norm=self.has_qk_norm, moe=self.moe)
        self.lm_head = ParallelLMHead(
            config.vocab_size, config.hidden_size,
            tie_word_embeddings=config.tie_word_embeddings,
            tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None,
        )

    def forward(self) -> torch.Tensor:
        hidden = self.model.forward(get_global_ctx().batch.input_ids)
        return self.lm_head.forward(hidden)
minisgl/models/decoder.py · Qwen3ForCausalLM
class Qwen3ForCausalLM(CausalLM):
    has_qk_norm = True
The sub-layers live in models/utils.py: RopeAttn (qkv_proj, attention layer, o_proj), GatedMLP (gate_up_proj, silu × up, down_proj) and MoEMLP.

minisgl/models/utils.py · GatedMLP
class GatedMLP(BaseOP):
    """A SwiGLU feed-forward network: gate and up merged into one matmul, then the activation, then down."""

    def __init__(self, config: ModelConfig):
        if config.hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {config.hidden_act}")
        self.gate_up_proj = LinearColParallelMerged(
            config.hidden_size, [config.intermediate_size, config.intermediate_size], has_bias=False
        )
        self.down_proj = LinearRowParallel(config.intermediate_size, config.hidden_size,
                                           has_bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(x)))
CausalLM.forward() takes no arguments, reads input_ids from the global context, and returns [batch_size, vocab] logits, one row per request, which is what comes out of the LM head taking only the last position during prefill.

Difference from upstream: the model files

Upstream writes one file per architecture (llama.py, qwen2.py, qwen3.py, qwen3_moe.py, mistral.py, each 80-odd nearly identical lines) and lazy-loads them through a registry keyed on architectures[0]. We merge them into one decoder.py plus three switches, and do not implement Mistral.

Streaming weight loading

minisgl/models/weight.py · load_weight
def load_weight(model_path: str, device: torch.device) -> Iterator[Tuple[str, torch.Tensor]]:
    config = ModelConfig.from_hf(cached_load_hf_config(model_path))
    tp = get_tp_info()
    folder = download_hf_weight(model_path)
    files = sorted(glob.glob(f"{folder}/*.safetensors"))
    merge_buf: Dict[str, Dict[str, torch.Tensor]] = {}
    expert_buf: Dict[str, Dict[int, torch.Tensor]] = {}
    for file in files:
        with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
            for name in f.keys():
                tensor = shard_tensor(name, f.get_tensor(name), tp.rank, tp.size,
                                      config.num_kv_heads)
                if (info := _merge_info(name)) is None:
                    out = (name, tensor)
                else:
                    merged, slot, slots = info
                    merge_buf.setdefault(merged, {})[slot] = tensor
                    if not all(s in merge_buf[merged] for s in slots):
                        continue  # merge only once the group's other projections have been read
                    parts = merge_buf.pop(merged)
                    out = (merged, torch.cat([parts[s] for s in slots], dim=0))

                if config.is_moe and (m := _EXPERT_PATTERN.match(out[0])) is not None:
                    packed = f"{m['prefix']}.{m['name'].removesuffix('.weight')}"
                    slots_e = expert_buf.setdefault(packed, {})
                    slots_e[int(m["idx"])] = out[1]
                    if len(slots_e) < config.num_experts:
                        continue  # wait for every expert before packing into [E, ...]
                    yield packed, torch.stack([slots_e[i] for i in range(config.num_experts)])
                    del expert_buf[packed]
                else:
                    yield out
    assert not merge_buf, f"Incomplete merge groups: {list(merge_buf)}"
    assert not expert_buf, f"Incomplete expert tensors: {list(expert_buf)}"
The loader is a generator that yields one fully processed (name, tensor) at a time:

  1. Read one at a time: safetensors.safe_open pulls a single tensor from the file on demand, so the whole checkpoint never has to be in memory. A 70B bf16 model is 140 GB, and reading it at once would blow out CPU memory; streaming adds only one tensor plus a small merge buffer at the peak.
  2. Shard for TP: shard_tensor takes this rank's slice (covered in chapter 16). Sharding before merging means the merged tensor is naturally this rank's [q_local; k_local; v_local].
  3. Merge projections: q_proj goes into a buffer when it is read, and once the same layer's k_proj and v_proj have also arrived they are concatenated in the fixed q, k, v order and yielded. That order must match how the attention layer splits them (qkv.split([q_dim, kv_dim, kv_dim])); the order in the file is alphabetical, which is why down_proj appears before the merged gate_up_proj in the output above, and has nothing to do with the concatenation order.
  4. Pack the experts: an MoE model stores each expert separately (experts.0.gate_proj, experts.1.gate_proj and so on), and once all of them have been read torch.stack turns them into one [E, ...] three-dimensional tensor.

The two assertions at the end make sure no merge group was left half-read.

Into the model

The engine (chapter 6) loads in three lines:

with torch.device("meta"), torch_dtype(config.dtype):   # an empty shell: shapes only, no memory
    self.model = create_model(config.model_config)
self.model.load_state_dict({k: v.to(self.dtype) for k, v in load_weight(path, device)})

torch_dtype temporarily changes the default dtype so that the torch.empty(...) calls during construction already use the target precision, which is what lets the dtype assertion in load_state_dict pass. load_state_dict replaces the meta tensors with the real ones from the checkpoint; Qwen3-0.6B's checkpoint carries an extra lm_head.weight (equal to embed_tokens.weight with tied embeddings) that ParallelLMHead.load_state_dict drops, while every other key must match exactly or it raises.

The official implementation

Tests

tests/test_ch03_models.py · test_weight_loader_produces_exactly_the_model_keys
def test_weight_loader_produces_exactly_the_model_keys():
    set_tp_info(0, 1)
    set_rope_device(torch.device("cpu"))
    config = ModelConfig.from_hf(cached_load_hf_config(QWEN3))
    with torch.device("meta"), torch_dtype(torch.bfloat16):
        model = create_model(config)
    expected = {k: tuple(v.shape) for k, v in model.state_dict().items()}
    loaded = {k: tuple(v.shape) for k, v in load_weight(QWEN3, torch.device("cpu"))}
    # Qwen3-0.6B ties its embeddings, yet the checkpoint still stores an lm_head.weight, which ParallelLMHead drops on load
    assert loaded.pop("lm_head.weight") == (151936, 1024)
    assert loaded == expected  # q/k/v merge into qkv_proj, gate/up into gate_up_proj
    assert expected["model.layers.0.self_attn.qkv_proj.weight"] == ((16 + 2 * 8) * 128, 1024)
This test makes sure the set of keys the loader produces is exactly the model's state_dict (apart from the dropped lm_head.weight), with matching shapes. test_model_config_from_hf checks the key fields of the Qwen3 and Qwen2.5 configs (Qwen2.5 only needs its config.json, kept in tests/configs/, so no model download is required).

How to explain it

On model loading: Llama, Qwen2, Qwen3 and Qwen3-MoE share one decoder structure and differ only in three switches, whether qkv carries a bias (Qwen2), whether there is QK-Norm (Qwen3), and whether the MLP is an MoE. The weights have to be loaded as a stream: reading the whole checkpoint at once doubles the peak memory, whereas streaming shards for tensor parallelism as it reads, concatenates q_proj / k_proj / v_proj into one qkv_proj in the same order the forward pass splits them, merges gate and up, and packs the MoE experts into a three-dimensional tensor. The config has to accept both the old and the new layout (transformers 5.x moved rope_theta into the RoPE parameters), and head_dim is not always hidden / heads.

Exercises

  1. Write a function on top of ModelConfig that computes how many bytes of KV cache one token takes (the next chapter needs it). What is it for Qwen3-0.6B in bf16?
  2. Add a tqdm progress bar to the loader, shown only on TP rank 0.
  3. If a checkpoint put q_proj and k_proj in two different safetensors files, would our loader still merge them correctly? Why?
Answers
  1. 2 * num_layers * num_kv_heads * head_dim * 2 bytes = 2 × 28 × 8 × 128 × 2 = 114688 bytes, about 112 KB.
  2. for file in tqdm(files, disable=not get_tp_info().is_primary()), which is what upstream does.
  3. Yes. The merge_buf persists across files, so a group merges as soon as all its projections have been read; only a group still unfinished after every file has been read raises.

Summary

  • ModelConfig collapses the various HF configs into the fields the engine needs; watch out for head_dim, the old and new homes of the RoPE parameters, and Qwen2's biases.
  • Llama, Qwen2, Qwen3 and Qwen3-MoE share one decoder structure, differing in the qkv bias, QK-Norm and MoE switches.
  • The streaming loader shards and merges as it reads: a small memory peak, q/k/v and gate/up merged into one matmul each, and MoE experts packed into a three-dimensional tensor.
  • The model is built as an empty shell on the meta device in the target dtype, then load_state_dict swaps in the real weights.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue