tiny-random/deepseek-v4

This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from deepseek-ai/DeepSeek-V4-Pro.

Note:

• This model follows the quantization scheme of "FP4 + FP8 Mixed": MoE expert parameters use FP4 precision; most other parameters use FP8.
• Chat template from this PR.

File path	Size
model.safetensors	144.0MB

Example usage:

• vLLM

# Not fully tested, please raise an issue if you find any problems.
model_id=tiny-random/deepseek-v4
vllm serve $model_id \
  --trust-remote-code \
  --kv-cache-dtype fp8 \
  --block-size 256 \
  --tensor-parallel-size 2 \
  --no-enable-flashinfer-autotune \
  --tokenizer-mode deepseek_v4 \
  --tool-call-parser deepseek_v4 \
  --enable-auto-tool-choice \
  --reasoning-parser deepseek_v4 \
  --speculative-config '{"method":"mtp","num_speculative_tokens":2}'

Codes to create this repo:

Click to expand

import json
import hashlib
from pathlib import Path
from typing import Any, Literal, TypedDict

import torch
from huggingface_hub import file_exists, hf_hub_download
from safetensors.torch import save_file
from transformers import AutoTokenizer, GenerationConfig

source_model_id = "deepseek-ai/DeepSeek-V4-Pro"
save_folder = "/tmp/tiny-random/deepseek-v4"
config = {
    "architectures": [
        "DeepseekV4ForCausalLM"
    ],
    "attention_bias": True,
    "attention_dropout": 0.0,
    "bos_token_id": 0,
    "eos_token_id": 1,
    "expert_dtype": "fp4",
    "hc_eps": 1e-06,
    "hc_mult": 4,
    "hc_sinkhorn_iters": 20,
    "head_dim": 128,
    "hidden_act": "silu",
    "hidden_size": 128,
    "index_head_dim": 128,
    "index_n_heads": 4,
    "index_topk": 1024,
    "initializer_range": 0.02,
    "max_position_embeddings": 1048576,
    "model_type": "deepseek_v4",
    "moe_intermediate_size": 256,
    "n_routed_experts": 128,
    "n_shared_experts": 1,
    "norm_topk_prob": True,
    "num_attention_heads": 4,
    "num_experts_per_tok": 6,
    "num_hidden_layers": 7,
    "num_hash_layers": 3,
    "num_key_value_heads": 1,
    "num_nextn_predict_layers": 1,
    "o_groups": 2,
    "o_lora_rank": 128,
    "q_lora_rank": 128,
    "qk_rope_head_dim": 64,
    "quantization_config": {
        "activation_scheme": "dynamic",
        "fmt": "e4m3",
        "quant_method": "fp8",
        "scale_fmt": "ue8m0",
        "weight_block_size": [
            128,
            128
        ]
    },
    "rms_norm_eps": 1e-06,
    "rope_scaling": {
        "beta_fast": 32,
        "beta_slow": 1,
        "factor": 16,
        "original_max_position_embeddings": 65536,
        "type": "yarn"
    },
    "rope_theta": 10000,
    "routed_scaling_factor": 2.5,
    "scoring_func": "sqrtsoftplus",
    "sliding_window": 128,
    "swiglu_limit": 10.0,
    "tie_word_embeddings": False,
    "topk_method": "noaux_tc",
    "torch_dtype": "bfloat16",
    "transformers_version": "4.57.1",
    "use_cache": True,
    "vocab_size": 129280,
    "compress_rope_theta": 160000,
    "compress_ratios": [128, 128, 4, 128, 4, 128, 4, 0]
}

def main():
    torch.manual_seed(42)
    Path(save_folder).mkdir(parents=True, exist_ok=True)
    state_dict = generate(config)
    save_file(state_dict, Path(save_folder) / "model.safetensors")
    with open(Path(save_folder) / "model.safetensors", "rb") as f:
        state_dict = f.read()
        print("Hash: ", hashlib.sha256(state_dict).hexdigest())

    with open(Path(save_folder) / "config.json", "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)

    tokenizer = AutoTokenizer.from_pretrained(
        source_model_id, trust_remote_code=True,
    )
    if file_exists(filename="chat_template.jinja", repo_id=source_model_id, repo_type='model', revision="refs/pr/146"):
        with open(hf_hub_download(
            source_model_id,
            filename="chat_template.jinja",
            repo_type='model',
            revision="refs/pr/146",
        ), 'r', encoding='utf-8') as f:
            tokenizer.chat_template = f.read()
    tokenizer.save_pretrained(save_folder)

    generation_config = GenerationConfig.from_pretrained(
        source_model_id, trust_remote_code=True,
    )
    generation_config.save_pretrained(save_folder)

BF16 = "torch.bfloat16"
F32 = "torch.float32"
FP8 = "torch.float8_e4m3fn"
SCALE = "torch.float8_e8m0fnu"
I8 = "torch.int8"
I64 = "torch.int64"

class TensorSpec(TypedDict):
    shape: list[int]
    dtype: str

Config = dict[str, Any]
State = dict[str, TensorSpec]
TensorDict = dict[str, torch.Tensor]
WeightKind = Literal["bf16", "fp8", "fp4"]

def initialize(state: State, config: Config, init_bound: float) -> TensorDict:
    tensors: TensorDict = {}
    scale_dtype = torch.float8_e8m0fnu
    fp4_boundaries = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0])

    def scale_for(
        amax: torch.Tensor, qmax: float, dtype: torch.dtype
    ) -> torch.Tensor:
        scale = amax.float().clamp_min(torch.finfo(torch.float32).tiny) / qmax
        if dtype == scale_dtype:
            scale = torch.pow(2.0, torch.round(torch.log2(scale)))
        return scale.to(dtype)

    def fp8_tensor(
        shape: list[int], scale_shape: list[int], dtype: torch.dtype
    ) -> tuple[torch.Tensor, torch.Tensor]:
        block_out, block_in = config["quantization_config"]["weight_block_size"]
        assert shape == [scale_shape[0] * block_out, scale_shape[1] * block_in]
        value = torch.empty(shape, dtype=torch.bfloat16).uniform_(
            -init_bound, init_bound
        )
        blocks = value.view(
            scale_shape[0], block_out, scale_shape[1], block_in
        ).transpose(1, 2)
        scales = scale_for(blocks.abs().amax(dim=(-1, -2)), 448, dtype)
        weight = (
            (blocks.float() / scales.float()[..., None, None])
            .clamp(-448, 448)
            .to(torch.float8_e4m3fn)
            .transpose(1, 2)
            .reshape(shape)
            .contiguous()
        )
        return weight, scales

    def fp4_tensors(
        count: int,
        shape: list[int],
        scale_shape: list[int],
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        out_dim, packed_in_dim = shape
        block_out, block_in = config["quantization_config"]["weight_block_size"]
        assert out_dim % block_out == 0
        assert packed_in_dim * 2 % block_in == 0
        assert packed_in_dim == scale_shape[1] * 16
        value = torch.empty(
            count, out_dim, packed_in_dim * 2, dtype=torch.bfloat16
        ).uniform_(-init_bound, init_bound)
        blocks = value.view(count, out_dim, scale_shape[1], 32)
        scales = scale_for(blocks.abs().amax(dim=-1), 6, dtype)
        normalized = (blocks.float() / scales.float()[..., None]).clamp(-6, 6)
        code = torch.bucketize(normalized.abs(), fp4_boundaries)
        code += normalized.signbit() * 8
        code = code.view(count, out_dim, packed_in_dim * 2)
        packed = (code[..., ::2] | (code[..., 1::2] << 4)).to(torch.uint8)
        weight = packed.view(torch.int8).contiguous()
        return weight, scales

    dtype_map: dict[str, torch.dtype] = {
        BF16: torch.bfloat16,
        F32: torch.float32,
        FP8: torch.float8_e4m3fn,
        SCALE: scale_dtype,
        I8: torch.int8,
        I64: torch.int64,
    }
    fp4_groups: dict[
        tuple[tuple[int, ...], tuple[int, ...], torch.dtype], list[str]
    ] = {}
    for name, spec in state.items():
        if spec["dtype"] != I8:
            continue
        scale_spec = state[name.replace(".weight", ".scale")]
        key = (
            tuple(spec["shape"]),
            tuple(scale_spec["shape"]),
            dtype_map[scale_spec["dtype"]],
        )
        fp4_groups.setdefault(key, []).append(name)

    fp4_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
    max_batch_elements = 4 * 1024 * 1024
    for (shape_tuple, scale_shape_tuple, dtype), names in fp4_groups.items():
        shape = list(shape_tuple)
        scale_shape = list(scale_shape_tuple)
        logical_elements = shape[0] * shape[1] * 2
        batch_size = max(1, max_batch_elements // logical_elements)
        for start in range(0, len(names), batch_size):
            batch_names = names[start: start + batch_size]
            weights, scales = fp4_tensors(
                len(batch_names), shape, scale_shape, dtype
            )
            for index, name in enumerate(batch_names):
                fp4_cache[name] = (
                    weights[index].clone(),
                    scales[index].clone(),
                )

    for name, spec in state.items():
        if name in tensors:
            continue
        shape, dtype = spec["shape"], dtype_map[spec["dtype"]]
        scale_name = name.replace(".weight", ".scale")
        if spec["dtype"] == FP8:
            scale_spec = state[scale_name]
            scale_type = dtype_map[scale_spec["dtype"]]
            tensors[name], tensors[scale_name] = fp8_tensor(
                shape, scale_spec["shape"], scale_type
            )
        elif spec["dtype"] == I8:
            tensors[name], tensors[scale_name] = fp4_cache.pop(name)
        elif spec["dtype"] == I64:
            tensors[name] = torch.randint(
                config["n_routed_experts"], shape, dtype=dtype
            )
        elif not name.endswith(".scale"):
            tensors[name] = torch.empty(shape, dtype=dtype).uniform_(
                -init_bound, init_bound
            )
        print(f"{name}: {shape} {dtype}", flush=True)
    file_size_by_name = {name: tensor.numel() * tensor.element_size() for name, tensor in tensors.items()}
    total_file_size = sum(file_size_by_name.values())
    k = 20
    topk = sorted(file_size_by_name.items(), key=lambda x: x[1], reverse=True)[:k]
    print(f"File size: {total_file_size / 1024 / 1024:.2f} MB")
    print(f"Top {k} largest tensors:")
    for name, size in topk:
        print(f"  {name}: {size / 1024 / 1024:.2f} MB")
    return tensors

def generate(
    config: Config, init_bound: float = 0.2
) -> TensorDict:
    assert init_bound > 0
    state: State = {}
    dim = config["hidden_size"]
    inter = config["moe_intermediate_size"]
    heads = config["num_attention_heads"]
    head_dim = config["head_dim"]
    q_rank = config["q_lora_rank"]
    o_rank = config["o_lora_rank"]
    o_groups = config["o_groups"]
    experts = config["n_routed_experts"]
    vocab = config["vocab_size"]
    hc = config["hc_mult"]
    quant = config["quantization_config"]
    block_out, block_in = quant["weight_block_size"]
    weight_kind: WeightKind = (
        "fp8" if quant["quant_method"] == "fp8" else "bf16"
    )
    scale_dtype = SCALE if quant.get("scale_fmt") == "ue8m0" else F32

    def add(name: str, shape: list[int], dtype: str) -> None:
        state[name] = {"shape": shape, "dtype": dtype}

    def linear(
        name: str, out_dim: int, in_dim: int, kind: WeightKind = weight_kind
    ) -> None:
        if kind == "fp4":
            assert out_dim % block_out == 0 and in_dim % block_in == 0, (
                f"{name} shape [{out_dim}, {in_dim}] is not divisible by "
                f"block size [{block_out}, {block_in}]"
            )
            add(f"{name}.weight", [out_dim, in_dim // 2], I8)
            add(f"{name}.scale", [out_dim, in_dim // 32], scale_dtype)
        elif kind == "fp8":
            assert out_dim % block_out == 0 and in_dim % block_in == 0, (
                f"{name} shape [{out_dim}, {in_dim}] is not divisible by "
                f"block size [{block_out}, {block_in}]"
            )
            add(f"{name}.weight", [out_dim, in_dim], FP8)
            add(
                f"{name}.scale",
                [out_dim // block_out, in_dim // block_in],
                scale_dtype,
            )
        else:
            add(f"{name}.weight", [out_dim, in_dim], BF16)

    def compressor(name: str, ratio: int, size: int) -> None:
        out_dim = size * (2 if ratio == 4 else 1)
        add(f"{name}.ape", [ratio, out_dim], F32)
        add(f"{name}.wkv.weight", [out_dim, dim], BF16)
        add(f"{name}.wgate.weight", [out_dim, dim], BF16)
        add(f"{name}.norm.weight", [size], BF16)

    def attention(name: str, ratio: int) -> None:
        add(f"{name}.attn_sink", [heads], F32)
        linear(f"{name}.wq_a", q_rank, dim)
        add(f"{name}.q_norm.weight", [q_rank], BF16)
        linear(f"{name}.wq_b", heads * head_dim, q_rank)
        linear(f"{name}.wkv", head_dim, dim)
        add(f"{name}.kv_norm.weight", [head_dim], BF16)
        linear(f"{name}.wo_a", o_groups * o_rank, heads * head_dim // o_groups)
        linear(f"{name}.wo_b", dim, o_groups * o_rank)

        if ratio:
            compressor(f"{name}.compressor", ratio, head_dim)
        if ratio == 4:
            index_heads = config["index_n_heads"]
            index_dim = config["index_head_dim"]
            linear(f"{name}.indexer.wq_b", index_heads * index_dim, q_rank)
            add(f"{name}.indexer.weights_proj.weight", [index_heads, dim], BF16)
            compressor(f"{name}.indexer.compressor", ratio, index_dim)

    def expert(name: str, kind: WeightKind) -> None:
        linear(f"{name}.w1", inter, dim, kind)
        linear(f"{name}.w2", dim, inter, kind)
        linear(f"{name}.w3", inter, dim, kind)

    def moe(name: str, layer_id: int) -> None:
        add(f"{name}.gate.weight", [experts, dim], BF16)
        if layer_id < config["num_hash_layers"]:
            add(
                f"{name}.gate.tid2eid",
                [vocab, config["num_experts_per_tok"]],
                I64,
            )
        else:
            add(f"{name}.gate.bias", [experts], F32)

        routed_kind = "fp4" if config.get("expert_dtype") == "fp4" else weight_kind
        for expert_id in range(experts):
            expert(f"{name}.experts.{expert_id}", routed_kind)
        expert(f"{name}.shared_experts", weight_kind)

    def block(name: str, layer_id: int) -> None:
        attention(f"{name}.attn", config["compress_ratios"][layer_id])
        moe(f"{name}.ffn", layer_id)
        add(f"{name}.attn_norm.weight", [dim], BF16)
        add(f"{name}.ffn_norm.weight", [dim], BF16)
        for part in ("attn", "ffn"):
            add(f"{name}.hc_{part}_fn", [(2 + hc) * hc, hc * dim], F32)
            add(f"{name}.hc_{part}_base", [(2 + hc) * hc], F32)
            add(f"{name}.hc_{part}_scale", [3], F32)

    def hc_head(name: str = "") -> None:
        prefix = f"{name}." if name else ""
        add(f"{prefix}hc_head_fn", [hc, hc * dim], F32)
        add(f"{prefix}hc_head_base", [hc], F32)
        add(f"{prefix}hc_head_scale", [1], F32)

    add("embed.weight", [vocab, dim], BF16)
    for layer_id in range(config["num_hidden_layers"]):
        block(f"layers.{layer_id}", layer_id)
    add("norm.weight", [dim], BF16)
    add("head.weight", [vocab, dim], BF16)
    hc_head()

    first_mtp_layer = config["num_hidden_layers"]
    for mtp_id in range(config["num_nextn_predict_layers"]):
        name = f"mtp.{mtp_id}"
        block(name, first_mtp_layer + mtp_id)
        linear(f"{name}.e_proj", dim, dim)
        linear(f"{name}.h_proj", dim, dim)
        for norm in ("enorm", "hnorm", "norm"):
            add(f"{name}.{norm}.weight", [dim], BF16)
        hc_head(name)

    return initialize(state, config, init_bound)

main()

Test environment:

• safetensors: 0.8.0
• torch: 2.11.0+cu128
• transformers: 5.14.0.dev0
• vllm: 0.24.0+cu129