Add tiny random checkpoint

README.md

@@ -0,0 +1,19 @@
+---
+library_name: transformers
+tags:
+- tiny-random
+- quantized
+---
+
+# tiny-random-Llama-NVFP4
+
+A tiny random model for testing, shrunk from [`nvidia/Llama-3.1-8B-Instruct-FP4`](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP4): the same
+architecture, quantization config and checkpoint layout at test sizes. Its key patterns, dtypes and tensor ranks match
+the real checkpoint's (`scripts/extract_layout.py`).
+
+modelopt NVFP4 on every decoder linear (packed U8, E4M3 `weight_scale` per 16, F32 `weight_scale_2`, calibrated F32 `input_scale`) plus the FP8 KV cache's `k_scale` / `v_scale`. Quantized with modelopt's `NVFP4QTensor`.
+
+`reference/` holds the same weights dequantized to bf16, under the unquantized model's keys: the reference to compare
+logits against, so a test measures what the load path and kernels add, not the quantization itself.
+
+The weights are random; the outputs mean nothing. `scripts/` rebuilds it from the real checkpoint's `config.json`.

config.json

@@ -0,0 +1,73 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": [
+    1
+  ],
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 512,
+  "max_position_embeddings": 4096,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 2,
+  "num_key_value_heads": 2,
+  "pad_token_id": null,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "factor": 8.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 500000.0,
+    "rope_type": "llama3"
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.18.0.dev0",
+  "use_cache": true,
+  "vocab_size": 1024,
+  "quantization_config": {
+    "config_groups": {
+      "group_0": {
+        "input_activations": {
+          "dynamic": false,
+          "num_bits": 4,
+          "type": "float",
+          "group_size": 16
+        },
+        "weights": {
+          "dynamic": false,
+          "num_bits": 4,
+          "type": "float",
+          "group_size": 16
+        },
+        "targets": [
+          "Linear"
+        ]
+      }
+    },
+    "ignore": [
+      "lm_head"
+    ],
+    "quant_algo": "NVFP4",
+    "kv_cache_scheme": {
+      "dynamic": false,
+      "num_bits": 8,
+      "type": "float"
+    },
+    "producer": {
+      "name": "modelopt",
+      "version": "0.37.0.dev5+g76fb12d47.d20250905"
+    },
+    "quant_method": "modelopt"
+  }
+}

generation_config.json

@@ -0,0 +1,11 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    1
+  ],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.18.0.dev0",
+  "use_cache": true
+}

hf_quant_config.json

@@ -0,0 +1,14 @@
+{
+    "producer": {
+        "name": "modelopt",
+        "version": "0.37.0.dev5+g76fb12d47.d20250905"
+    },
+    "quantization": {
+        "quant_algo": "NVFP4",
+        "kv_cache_quant_algo": "FP8",
+        "group_size": 16,
+        "exclude_modules": [
+            "lm_head"
+        ]
+    }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a96839c1a64d95f02e40b69ac4126a34f26832ea9c9836f9a4c5ee3f3e556a4
+size 1721712

reference/config.json

@@ -0,0 +1,38 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": [
+    1
+  ],
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 512,
+  "max_position_embeddings": 4096,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 2,
+  "num_key_value_heads": 2,
+  "pad_token_id": null,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "factor": 8.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 500000.0,
+    "rope_type": "llama3"
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.18.0.dev0",
+  "use_cache": true,
+  "vocab_size": 1024
+}

reference/generation_config.json

@@ -0,0 +1,11 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    1
+  ],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.18.0.dev0",
+  "use_cache": true
+}

reference/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d605548fe78a18c4d21b9b17c050dab60e3d8411a0b52ed270ac9b0daf70f973
+size 3412672

scripts/build_dense_bf16.py

@@ -0,0 +1,54 @@
+"""A tiny random dense model in bf16, from a real checkpoint's config.json shrunk in its own fields:
+Qwen3, Llama or Mistral3 (Ministral-3, with its vision tower)."""
+
+import argparse
+import json
+import tempfile
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--real-config", required=True)
+parser.add_argument("--out", required=True)
+args = parser.parse_args()
+
+text_sizes = dict(
+    vocab_size=1024,
+    hidden_size=256,
+    intermediate_size=512,
+    num_hidden_layers=2,
+    num_attention_heads=4,
+    num_key_value_heads=2,
+    head_dim=64,
+    max_position_embeddings=4096,
+)
+real = json.load(open(args.real_config))
+real.pop("quantization_config", None)
+for key in ("architectures", "transformers_version", "torch_dtype", "dtype"):
+    real.pop(key, None)
+if "text_config" in real:  # a VLM: shrink the decoder and the vision tower in their own sections
+    text = {k: v for k, v in real["text_config"].items() if k not in ("architectures", "quantization_config")}
+    shrunk = dict(
+        real,
+        text_config=dict(text, **text_sizes),
+        vision_config=dict(
+            real["vision_config"], hidden_size=64, intermediate_size=128, num_hidden_layers=2, num_attention_heads=2,
+            head_dim=32, image_size=112,
+        ),
+        image_token_index=1000,
+    )
+else:
+    shrunk = dict(real, **text_sizes)
+for token in ("bos_token_id", "eos_token_id", "pad_token_id"):
+    for section in (shrunk, shrunk.get("text_config", {})):
+        if section.get(token) is not None:
+            section[token] = [1] if isinstance(section[token], list) else 1
+with tempfile.TemporaryDirectory() as tmp:  # parsed from disk, the way the real config.json is
+    json.dump(shrunk, open(f"{tmp}/config.json", "w"))
+    config = AutoConfig.from_pretrained(tmp)
+model_cls = AutoModelForImageTextToText if "vision_config" in config.to_dict() else AutoModelForCausalLM
+torch.manual_seed(0)
+model = model_cls.from_config(config).to(torch.bfloat16)
+model.save_pretrained(args.out)
+print(type(model).__name__)

scripts/calibrate_experts.py

@@ -0,0 +1,60 @@
+"""Per-expert activation amax of a bf16 MoE, as a modelopt calibration run collects it: the routed
+hidden rows each expert's gate/up sees, and the intermediate its down sees. Writes
+{"<layer>.<expert>.gate_up": amax, "<layer>.<expert>.down": amax, "<linear name>": input amax, "<linear name>:out": output amax}."""
+
+import argparse
+import json
+import re
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--bf16", required=True)
+parser.add_argument("--out", required=True)
+args = parser.parse_args()
+
+vlm = "vision_config" in AutoConfig.from_pretrained(args.bf16).to_dict()
+model_cls = AutoModelForImageTextToText if vlm else AutoModelForCausalLM
+model = model_cls.from_pretrained(args.bf16, dtype=torch.bfloat16, device_map="cuda").eval()
+amax = {}
+
+
+def hook(name):
+    layer = re.search(r"layers\.(\d+)\.", name).group(1)
+
+    def record(module, inputs, _output):
+        hidden, top_k_index = inputs[0], inputs[1]
+        for expert in range(module.num_experts):
+            rows = hidden[(top_k_index == expert).any(-1)]
+            if rows.numel() == 0:
+                continue
+            gate, up = (rows @ module.gate_up_proj[expert].T).chunk(2, dim=-1)
+            inter = module.act_fn(gate) * up
+            for key, value in ((f"{layer}.{expert}.gate_up", rows), (f"{layer}.{expert}.down", inter)):
+                amax[key] = max(amax.get(key, 0.0), value.abs().max().item())
+
+    return record
+
+
+def linear_hook(name):
+    def record(module, inputs, output):
+        amax[name] = max(amax.get(name, 0.0), inputs[0].abs().max().item())
+        # the output too: a k/v projection's is what an FP8 KV cache's `k_scale`/`v_scale` cover
+        amax[f"{name}:out"] = max(amax.get(f"{name}:out", 0.0), output.abs().max().item())
+
+    return record
+
+
+for name, module in model.named_modules():
+    if name.endswith("mlp.experts"):
+        module.register_forward_hook(hook(name))
+    elif isinstance(module, torch.nn.Linear):  # dense linears, for exports that quantize them too
+        module.register_forward_hook(linear_hook(name))
+
+torch.manual_seed(0)
+with torch.no_grad():
+    for _ in range(8):
+        model(torch.randint(3, model.config.get_text_config().vocab_size, (4, 128), device="cuda"))
+json.dump(amax, open(args.out, "w"), indent=1)
+print(f"{len(amax)} amax values")

scripts/extract_layout.py

@@ -0,0 +1,59 @@
+"""The key layout of a safetensors checkpoint, read from shard headers only (no tensor data).
+
+Keys are normalized by replacing every integer path component with `N`, so a layout is the set of
+distinct patterns with their dtypes, the number of keys behind each, and the layer indices each appears
+at. `compare(real, tiny)` is what the tiny-checkpoint builders assert against.
+"""
+
+import json
+import re
+import struct
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+
+def read_header(path):
+    with open(path, "rb") as f:
+        (n,) = struct.unpack("<Q", f.read(8))
+        header = json.loads(f.read(n))
+    header.pop("__metadata__", None)
+    return header
+
+
+def normalize(key):
+    return re.sub(r"(?<=\.)\d+(?=\.|$)", "N", key)
+
+
+def headers(source):
+    """`{key: {"dtype", "shape"}}` of a local checkpoint directory, or of a hub repo read remotely."""
+    if Path(source).is_dir():
+        for shard in sorted(Path(source).glob("*.safetensors")):
+            yield from read_header(shard).items()
+        return
+    from huggingface_hub import HfApi
+
+    for shard in HfApi().get_safetensors_metadata(source).files_metadata.values():
+        for key, tensor in shard.tensors.items():
+            yield key, {"dtype": tensor.dtype, "shape": tensor.shape}
+
+
+def layout(source):
+    patterns = defaultdict(lambda: {"dtypes": set(), "count": 0, "layers": set(), "ndims": set()})
+    for key, info in headers(source):
+        entry = patterns[normalize(key)]
+        entry["dtypes"].add(info["dtype"])
+        entry["ndims"].add(len(info["shape"]))
+        entry["count"] += 1
+        if m := re.search(r"layers\.(\d+)\.", key):
+            entry["layers"].add(int(m.group(1)))
+    return {
+        p: {"dtypes": sorted(e["dtypes"]), "count": e["count"], "layers": sorted(e["layers"]), "ndims": sorted(e["ndims"])}
+        for p, e in sorted(patterns.items())
+    }
+
+
+if __name__ == "__main__":
+    out = layout(sys.argv[1])
+    json.dump(out, open(sys.argv[2], "w"), indent=1)
+    print(f"{len(out)} patterns")

scripts/quantize_dense_nvfp4.py

@@ -0,0 +1,56 @@
+"""A tiny bf16 Llama written as nvidia/Llama-3.1-8B-Instruct-FP4 ships it, with modelopt's own NVFP4
+quantizer: every decoder linear NVFP4 (packed U8 weight, E4M3 `weight_scale` per 16, F32 `weight_scale_2`),
+a calibrated F32 `input_scale` per linear, and the FP8 KV cache's `k_proj.k_scale` / `v_proj.v_scale`
+(output amax / 448). The real `quantization_config` and `hf_quant_config.json`; also <out>/dequantized,
+the logits reference. Run in the producer venv (nvidia-modelopt)."""
+
+import argparse
+import json
+import re
+import shutil
+from pathlib import Path
+
+import torch
+from modelopt.torch.quantization.qtensor import NVFP4QTensor
+from safetensors.torch import load_file, save_file
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--bf16", required=True)
+parser.add_argument("--real", required=True, help="directory with the real config.json + hf_quant_config.json")
+parser.add_argument("--calibration", required=True)
+parser.add_argument("--out", required=True)
+args = parser.parse_args()
+
+source = load_file(f"{args.bf16}/model.safetensors")
+calibration = json.load(open(args.calibration))
+linear = re.compile(r"model\.layers\.\d+\.(self_attn|mlp)\.\w+_proj\.weight")
+quantized, dequantized = {}, dict(source)
+for key, tensor in source.items():
+    if not linear.fullmatch(key):
+        quantized[key] = tensor
+        continue
+    module = key.removesuffix(".weight")
+    qtensor, block_scale, global_scale = NVFP4QTensor.quantize(tensor.float(), 16)
+    quantized[key] = qtensor._quantized_data
+    quantized[f"{module}.weight_scale"] = block_scale
+    quantized[f"{module}.weight_scale_2"] = global_scale.reshape(()).float()
+    quantized[f"{module}.input_scale"] = torch.tensor(calibration[module] / (6.0 * 448.0), dtype=torch.float32)
+    for proj in ("k", "v"):
+        if module.endswith(f".{proj}_proj"):
+            quantized[f"{module}.{proj}_scale"] = torch.tensor(calibration[f"{module}:out"] / 448.0, dtype=torch.float32)
+    dequantized[key] = NVFP4QTensor(tensor.shape, torch.float32, qtensor._quantized_data).dequantize(
+        dtype=torch.float32, scale=block_scale, double_scale=global_scale, block_sizes={-1: 16}
+    ).to(torch.bfloat16)
+
+out = Path(args.out)
+real = json.load(open(f"{args.real}/config.json"))
+for name, tensors in (("nvfp4", quantized), ("dequantized", dequantized)):
+    (out / name).mkdir(parents=True, exist_ok=True)
+    save_file({k: v.contiguous() for k, v in tensors.items()}, out / name / "model.safetensors", metadata={"format": "pt"})
+    config = json.load(open(f"{args.bf16}/config.json"))
+    if name == "nvfp4":
+        config["quantization_config"] = real["quantization_config"]
+        shutil.copy(f"{args.real}/hf_quant_config.json", out / name / "hf_quant_config.json")
+    json.dump(config, open(out / name / "config.json", "w"), indent=2)
+    shutil.copy(f"{args.bf16}/generation_config.json", out / name / "generation_config.json")
+print(f"wrote {len(quantized)} tensors")