feat: Add comprehensive metadata to save_quantized for load_quantized

Metadata now includes all fields needed by from_quantized() to reconstruct
a KbitLoraModel without the original HF model:

- Model config: hidden_size, num_attention_heads, num_key_value_heads,
  head_dim, intermediate_size, vocab_size, rms_norm_eps, rope_theta
- MoE config: expert_intermediate_size, has_shared_expert, has_qk_norm,
  dense_layer_indices
- Per-projection dims: N, K, N_padded, k for every attention/MLP/expert
  projection in every layer, plus LM head dims

Updated test_checkpoint.py to verify all metadata fields are present and
correct for a tiny Llama model.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/checkpoint.py

@@ -79,24 +79,86 @@ def save_quantized(model, path: str):
     if model.embed_tokens is not None:
         tensors["embed_tokens.weight"] = model.embed_tokens.weight.data
 
-    # Metadata
+    # Metadata — comprehensive, enables load_quantized without the HF model
     metadata = {
+        # Model architecture
         "model_type": model.model_type,
         "hidden_size": str(model.hidden_size),
         "num_layers": str(model.num_layers),
         "num_loaded_layers": str(model._num_loaded_layers),
         "layer_start": str(model._layer_start),
         "layer_end": str(model._layer_end),
+        "num_attention_heads": str(model.num_heads),
+        "num_key_value_heads": str(model.num_kv_heads),
+        "head_dim": str(model.head_dim),
+        "intermediate_size": str(model.intermediate_size),
+        "vocab_size": str(model.vocab_size),
+        "rms_norm_eps": str(model.rms_norm_eps),
+        "rope_theta": str(model.rope_theta),
+        # Quantization config
         "k_attention": str(model.k_attention),
         "k_mlp": str(model.k_mlp),
         "k_lm_head": str(model.k_lm_head),
         "k_experts": str(model.k_experts),
         "k_shared_expert": str(model.k_shared_expert),
+        # MoE config
         "is_moe": str(model.arch.is_moe),
         "num_experts": str(model.arch.num_experts),
         "num_active_experts": str(model.arch.num_active_experts),
+        "expert_intermediate_size": str(model.arch.expert_intermediate_size),
+        "has_shared_expert": str(model.arch.has_shared_expert),
+        "has_qk_norm": str(model.arch.has_qk_norm),
     }
 
+    # Dense layer indices (comma-separated, empty if None or all MoE)
+    if model.arch.dense_layer_indices is not None:
+        metadata["dense_layer_indices"] = ",".join(
+            str(i) for i in model.arch.dense_layer_indices
+        )
+    else:
+        metadata["dense_layer_indices"] = ""
+
+    # Per-projection dimensions (needed for LoRA initialization in load_quantized)
+    for i, layer_info in enumerate(model._layer_data):
+        prefix = f"layer.{i}"
+
+        # Attention projections
+        for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+            metadata[f"{prefix}.attn.{proj}.N"] = str(layer_info[proj]["N"])
+            metadata[f"{prefix}.attn.{proj}.K"] = str(layer_info[proj]["K"])
+            metadata[f"{prefix}.attn.{proj}.N_padded"] = str(layer_info[proj]["N_padded"])
+            metadata[f"{prefix}.attn.{proj}.k"] = str(layer_info[proj]["k"])
+
+        # MLP or MoE
+        if layer_info.get("is_moe"):
+            # Shared expert dims
+            if "shared_gate_proj" in layer_info:
+                for proj in ["shared_gate_proj", "shared_up_proj", "shared_down_proj"]:
+                    metadata[f"{prefix}.moe.{proj}.N"] = str(layer_info[proj]["N"])
+                    metadata[f"{prefix}.moe.{proj}.K"] = str(layer_info[proj]["K"])
+                    metadata[f"{prefix}.moe.{proj}.N_padded"] = str(layer_info[proj]["N_padded"])
+                    metadata[f"{prefix}.moe.{proj}.k"] = str(layer_info[proj]["k"])
+
+            # Expert dims (same for all experts — store once)
+            metadata[f"{prefix}.moe.experts.N"] = str(layer_info.get("expert_N", 0))
+            metadata[f"{prefix}.moe.experts.K"] = str(layer_info.get("expert_K", 0))
+            metadata[f"{prefix}.moe.experts.N_padded"] = str(layer_info.get("expert_N_padded", 0))
+            metadata[f"{prefix}.moe.experts.k"] = str(layer_info.get("expert_k", 0))
+        else:
+            for proj in ["gate_proj", "up_proj", "down_proj"]:
+                metadata[f"{prefix}.mlp.{proj}.N"] = str(layer_info[proj]["N"])
+                metadata[f"{prefix}.mlp.{proj}.K"] = str(layer_info[proj]["K"])
+                metadata[f"{prefix}.mlp.{proj}.N_padded"] = str(layer_info[proj]["N_padded"])
+                metadata[f"{prefix}.mlp.{proj}.k"] = str(layer_info[proj]["k"])
+
+    # LM head dims
+    if model._lm_head_info is not None:
+        lm = model._lm_head_info
+        metadata["lm_head.N"] = str(lm["N"])
+        metadata["lm_head.K"] = str(lm["K"])
+        metadata["lm_head.N_padded"] = str(lm["N_padded"])
+        metadata["lm_head.k"] = str(lm["k"])
+
     # Move all tensors to CPU for saving
     cpu_tensors = OrderedDict()
     for k, v in tensors.items():

tests/test_checkpoint.py

@@ -81,9 +81,59 @@ def test_metadata_present(self, kbit_model):
             save_quantized(kbit_model, path)
             sf = safe_open(path, framework="pt", device="cpu")
             meta = sf.metadata()
+
+            # Model architecture
             assert meta["model_type"] == "llama"
-            assert meta["k_attention"] == "4"
+            assert int(meta["hidden_size"]) == 256
             assert int(meta["num_layers"]) == 2
+            assert int(meta["num_attention_heads"]) == 4
+            assert int(meta["num_key_value_heads"]) == 2
+            assert int(meta["head_dim"]) == 64  # 256 / 4
+            assert int(meta["intermediate_size"]) == 512
+            assert int(meta["vocab_size"]) == 1000
+            assert float(meta["rms_norm_eps"]) > 0
+            assert float(meta["rope_theta"]) > 0
+
+            # Quantization config
+            assert meta["k_attention"] == "4"
+            assert meta["k_mlp"] == "4"
+            assert meta["k_lm_head"] == "4"
+            assert meta["k_experts"] == "4"
+            assert meta["k_shared_expert"] == "4"
+
+            # MoE config
+            assert meta["is_moe"] == "False"
+            assert meta["has_shared_expert"] == "False"
+            assert meta["has_qk_norm"] == "False"
+            assert meta["dense_layer_indices"] == ""
+
+            # Per-projection dims for layer 0 attention
+            assert int(meta["layer.0.attn.q_proj.N"]) == 256  # q_dim = 4 * 64
+            assert int(meta["layer.0.attn.q_proj.K"]) == 256  # hidden_size
+            assert int(meta["layer.0.attn.q_proj.N_padded"]) == 256  # already mult of 128
+            assert int(meta["layer.0.attn.q_proj.k"]) == 4
+
+            assert int(meta["layer.0.attn.k_proj.N"]) == 128  # kv_dim = 2 * 64
+            assert int(meta["layer.0.attn.k_proj.K"]) == 256
+
+            # MLP dims
+            assert int(meta["layer.0.mlp.gate_proj.N"]) == 512  # intermediate
+            assert int(meta["layer.0.mlp.gate_proj.K"]) == 256  # hidden
+
+            # LM head dims
+            assert int(meta["lm_head.N"]) == 1000  # vocab_size
+            assert int(meta["lm_head.K"]) == 256  # hidden_size
+
+            # Check all layers have dims
+            for i in range(2):
+                for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+                    assert f"layer.{i}.attn.{proj}.N" in meta
+                    assert f"layer.{i}.attn.{proj}.K" in meta
+                    assert f"layer.{i}.attn.{proj}.N_padded" in meta
+                    assert f"layer.{i}.attn.{proj}.k" in meta
+                for proj in ["gate_proj", "up_proj", "down_proj"]:
+                    assert f"layer.{i}.mlp.{proj}.N" in meta
+                    assert f"layer.{i}.mlp.{proj}.K" in meta
         finally:
             os.unlink(path)