Views
No views yet
| File path | Size |
|---|---|
| model.safetensors | 277.3MB |
1# Not fully tested, please report any issues if you find any problems.
2model_id=tiny-random/deepseek-v4
3vllm serve $model_id \
4 --trust-remote-code \
5 --kv-cache-dtype fp8 \
6 --block-size 256 \
7 --tensor-parallel-size 2 \
8 --no-enable-flashinfer-autotune \
9 --tokenizer-mode deepseek_v4 \
10 --tool-call-parser deepseek_v4 \
11 --enable-auto-tool-choice \
12 --reasoning-parser deepseek_v4 \
13 --speculative-config '{"method":"mtp","num_speculative_tokens":2}'1# Tested on H20. Please report any issues if you find any problems.
2export NVCC_PREPEND_FLAGS=-allow-unsupported-compiler
3export NVCC_APPEND_FLAGS=-allow-unsupported-compiler
4export SGLANG_OPT_USE_TILELANG_MHC_PRE=0
5export SGLANG_OPT_USE_TILELANG_MHC_POST=0
6export SGLANG_OPT_DEEPGEMM_HC_PRENORM=0
7export SGLANG_OPT_USE_TILELANG_INDEXER=1
8
9model_id=tiny-random/deepseek-v4
10sglang serve \
11 --trust-remote-code \
12 --model-path $model_id \
13 --tp 2 \
14 --moe-runner-backend marlin \
15 --fp8-gemm-backend triton \
16 --speculative-algorithm EAGLE \
17 --speculative-num-steps 3 \
18 --speculative-eagle-topk 1 \
19 --speculative-num-draft-tokens 4 \
20 --mem-fraction-static 0.6 \
21 --disable-cuda-graph \
22 --disable-custom-all-reduce1import json
2import hashlib
3from pathlib import Path
4from typing import Any, Literal, TypedDict
5
6import torch
7from huggingface_hub import file_exists, hf_hub_download
8from safetensors.torch import save_file
9from transformers import AutoTokenizer, GenerationConfig
10
11source_model_id = "deepseek-ai/DeepSeek-V4-Pro"
12save_folder = "/tmp/tiny-random/deepseek-v4"
13config = {
14 "architectures": [
15 "DeepseekV4ForCausalLM"
16 ],
17 "attention_bias": True,
18 "attention_dropout": 0.0,
19 "bos_token_id": 0,
20 "eos_token_id": 1,
21 "expert_dtype": "fp4",
22 "hc_eps": 1e-06,
23 "hc_mult": 4,
24 "hc_sinkhorn_iters": 20,
25 # SGLang's DSV4 KV-cache layout fixes the non-RoPE portion at 448
26 # elements; together with qk_rope_head_dim=64 this must be 512.
27 "head_dim": 512,
28 "hidden_act": "silu",
29 # SGLang's Hopper MXFP4 Marlin path pads hidden_size to 256. Keeping
30 # the checkpoint at 128 leaves its per-32 scales at width 4 while the
31 # runtime allocates width 8, so make the checkpoint natively compatible.
32 "hidden_size": 256,
33 "index_head_dim": 128,
34 "index_n_heads": 32,
35 "index_topk": 1024,
36 "initializer_range": 0.02,
37 "max_position_embeddings": 1048576,
38 "model_type": "deepseek_v4",
39 "moe_intermediate_size": 256,
40 "n_routed_experts": 128,
41 "n_shared_experts": 1,
42 "norm_topk_prob": True,
43 "num_attention_heads": 4,
44 "num_experts_per_tok": 6,
45 "num_hidden_layers": 7,
46 "num_hash_layers": 3,
47 "num_key_value_heads": 1,
48 "num_nextn_predict_layers": 1,
49 "o_groups": 2,
50 "o_lora_rank": 128,
51 "q_lora_rank": 128,
52 "qk_rope_head_dim": 64,
53 "quantization_config": {
54 "activation_scheme": "dynamic",
55 "fmt": "e4m3",
56 "quant_method": "fp8",
57 "scale_fmt": "ue8m0",
58 "weight_block_size": [
59 128,
60 128
61 ]
62 },
63 "rms_norm_eps": 1e-06,
64 "rope_scaling": {
65 "beta_fast": 32,
66 "beta_slow": 1,
67 "factor": 16,
68 "original_max_position_embeddings": 65536,
69 "type": "yarn"
70 },
71 "rope_theta": 10000,
72 "routed_scaling_factor": 2.5,
73 "scoring_func": "sqrtsoftplus",
74 "sliding_window": 128,
75 "swiglu_limit": 10.0,
76 "tie_word_embeddings": False,
77 "topk_method": "noaux_tc",
78 "torch_dtype": "bfloat16",
79 "transformers_version": "4.57.1",
80 "use_cache": True,
81 "vocab_size": 129280,
82 "compress_rope_theta": 160000,
83 "compress_ratios": [128, 128, 4, 128, 4, 128, 4, 0]
84}
85
86def main():
87 torch.manual_seed(42)
88 Path(save_folder).mkdir(parents=True, exist_ok=True)
89 state_dict = generate(config)
90 save_file(state_dict, Path(save_folder) / "model.safetensors")
91 with open(Path(save_folder) / "model.safetensors", "rb") as f:
92 state_dict = f.read()
93 print("Hash: ", hashlib.sha256(state_dict).hexdigest())
94
95 with open(Path(save_folder) / "config.json", "w", encoding="utf-8") as f:
96 json.dump(config, f, indent=2, ensure_ascii=False)
97
98 tokenizer = AutoTokenizer.from_pretrained(
99 source_model_id, trust_remote_code=True,
100 )
101 if file_exists(filename="chat_template.jinja", repo_id=source_model_id, repo_type='model', revision="refs/pr/146"):
102 with open(hf_hub_download(
103 source_model_id,
104 filename="chat_template.jinja",
105 repo_type='model',
106 revision="refs/pr/146",
107 ), 'r', encoding='utf-8') as f:
108 tokenizer.chat_template = f.read()
109 tokenizer.save_pretrained(save_folder)
110
111 generation_config = GenerationConfig.from_pretrained(
112 source_model_id, trust_remote_code=True,
113 )
114 generation_config.save_pretrained(save_folder)
115
116BF16 = "torch.bfloat16"
117F32 = "torch.float32"
118FP8 = "torch.float8_e4m3fn"
119SCALE = "torch.float8_e8m0fnu"
120I8 = "torch.int8"
121I64 = "torch.int64"
122
123class TensorSpec(TypedDict):
124 shape: list[int]
125 dtype: str
126
127Config = dict[str, Any]
128State = dict[str, TensorSpec]
129TensorDict = dict[str, torch.Tensor]
130WeightKind = Literal["bf16", "fp8", "fp4"]
131
132def initialize(state: State, config: Config, init_bound: float) -> TensorDict:
133 tensors: TensorDict = {}
134 scale_dtype = torch.float8_e8m0fnu
135 fp4_boundaries = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0])
136
137 def scale_for(
138 amax: torch.Tensor, qmax: float, dtype: torch.dtype
139 ) -> torch.Tensor:
140 scale = amax.float().clamp_min(torch.finfo(torch.float32).tiny) / qmax
141 if dtype == scale_dtype:
142 scale = torch.pow(2.0, torch.round(torch.log2(scale)))
143 return scale.to(dtype)
144
145 def fp8_tensor(
146 shape: list[int], scale_shape: list[int], dtype: torch.dtype
147 ) -> tuple[torch.Tensor, torch.Tensor]:
148 block_out, block_in = config["quantization_config"]["weight_block_size"]
149 assert shape == [scale_shape[0] * block_out, scale_shape[1] * block_in]
150 value = torch.empty(shape, dtype=torch.bfloat16).uniform_(
151 -init_bound, init_bound
152 )
153 blocks = value.view(
154 scale_shape[0], block_out, scale_shape[1], block_in
155 ).transpose(1, 2)
156 scales = scale_for(blocks.abs().amax(dim=(-1, -2)), 448, dtype)
157 weight = (
158 (blocks.float() / scales.float()[..., None, None])
159 .clamp(-448, 448)
160 .to(torch.float8_e4m3fn)
161 .transpose(1, 2)
162 .reshape(shape)
163 .contiguous()
164 )
165 return weight, scales
166
167 def fp4_tensors(
168 count: int,
169 shape: list[int],
170 scale_shape: list[int],
171 dtype: torch.dtype,
172 ) -> tuple[torch.Tensor, torch.Tensor]:
173 out_dim, packed_in_dim = shape
174 block_out, block_in = config["quantization_config"]["weight_block_size"]
175 assert out_dim % block_out == 0
176 assert packed_in_dim * 2 % block_in == 0
177 assert packed_in_dim == scale_shape[1] * 16
178 value = torch.empty(
179 count, out_dim, packed_in_dim * 2, dtype=torch.bfloat16
180 ).uniform_(-init_bound, init_bound)
181 blocks = value.view(count, out_dim, scale_shape[1], 32)
182 scales = scale_for(blocks.abs().amax(dim=-1), 6, dtype)
183 normalized = (blocks.float() / scales.float()[..., None]).clamp(-6, 6)
184 code = torch.bucketize(normalized.abs(), fp4_boundaries)
185 code += normalized.signbit() * 8
186 code = code.view(count, out_dim, packed_in_dim * 2)
187 packed = (code[..., ::2] | (code[..., 1::2] << 4)).to(torch.uint8)
188 weight = packed.view(torch.int8).contiguous()
189 return weight, scales
190
191 dtype_map: dict[str, torch.dtype] = {
192 BF16: torch.bfloat16,
193 F32: torch.float32,
194 FP8: torch.float8_e4m3fn,
195 SCALE: scale_dtype,
196 I8: torch.int8,
197 I64: torch.int64,
198 }
199 fp4_groups: dict[
200 tuple[tuple[int, ...], tuple[int, ...], torch.dtype], list[str]
201 ] = {}
202 for name, spec in state.items():
203 if spec["dtype"] != I8:
204 continue
205 scale_spec = state[name.replace(".weight", ".scale")]
206 key = (
207 tuple(spec["shape"]),
208 tuple(scale_spec["shape"]),
209 dtype_map[scale_spec["dtype"]],
210 )
211 fp4_groups.setdefault(key, []).append(name)
212
213 fp4_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
214 max_batch_elements = 4 * 1024 * 1024
215 for (shape_tuple, scale_shape_tuple, dtype), names in fp4_groups.items():
216 shape = list(shape_tuple)
217 scale_shape = list(scale_shape_tuple)
218 logical_elements = shape[0] * shape[1] * 2
219 batch_size = max(1, max_batch_elements // logical_elements)
220 for start in range(0, len(names), batch_size):
221 batch_names = names[start: start + batch_size]
222 weights, scales = fp4_tensors(
223 len(batch_names), shape, scale_shape, dtype
224 )
225 for index, name in enumerate(batch_names):
226 fp4_cache[name] = (
227 weights[index].clone(),
228 scales[index].clone(),
229 )
230
231 for name, spec in state.items():
232 if name in tensors:
233 continue
234 shape, dtype = spec["shape"], dtype_map[spec["dtype"]]
235 scale_name = name.replace(".weight", ".scale")
236 if spec["dtype"] == FP8:
237 scale_spec = state[scale_name]
238 scale_type = dtype_map[scale_spec["dtype"]]
239 tensors[name], tensors[scale_name] = fp8_tensor(
240 shape, scale_spec["shape"], scale_type
241 )
242 elif spec["dtype"] == I8:
243 tensors[name], tensors[scale_name] = fp4_cache.pop(name)
244 elif spec["dtype"] == I64:
245 tensors[name] = torch.randint(
246 config["n_routed_experts"], shape, dtype=dtype
247 )
248 elif not name.endswith(".scale"):
249 tensors[name] = torch.empty(shape, dtype=dtype).uniform_(
250 -init_bound, init_bound
251 )
252 print(f"{name}: {shape} {dtype}", flush=True)
253 file_size_by_name = {name: tensor.numel() * tensor.element_size() for name, tensor in tensors.items()}
254 total_file_size = sum(file_size_by_name.values())
255 k = 20
256 topk = sorted(file_size_by_name.items(), key=lambda x: x[1], reverse=True)[:k]
257 print(f"File size: {total_file_size / 1024 / 1024:.2f} MB")
258 print(f"Top {k} largest tensors:")
259 for name, size in topk:
260 print(f" {name}: {size / 1024 / 1024:.2f} MB")
261 return tensors
262
263def generate(
264 config: Config, init_bound: float = 0.2
265) -> TensorDict:
266 assert init_bound > 0
267 state: State = {}
268 dim = config["hidden_size"]
269 inter = config["moe_intermediate_size"]
270 heads = config["num_attention_heads"]
271 head_dim = config["head_dim"]
272 q_rank = config["q_lora_rank"]
273 o_rank = config["o_lora_rank"]
274 o_groups = config["o_groups"]
275 experts = config["n_routed_experts"]
276 vocab = config["vocab_size"]
277 hc = config["hc_mult"]
278 quant = config["quantization_config"]
279 block_out, block_in = quant["weight_block_size"]
280 weight_kind: WeightKind = (
281 "fp8" if quant["quant_method"] == "fp8" else "bf16"
282 )
283 scale_dtype = SCALE if quant.get("scale_fmt") == "ue8m0" else F32
284
285 def add(name: str, shape: list[int], dtype: str) -> None:
286 state[name] = {"shape": shape, "dtype": dtype}
287
288 def linear(
289 name: str, out_dim: int, in_dim: int, kind: WeightKind = weight_kind
290 ) -> None:
291 if kind == "fp4":
292 assert out_dim % block_out == 0 and in_dim % block_in == 0, (
293 f"{name} shape [{out_dim}, {in_dim}] is not divisible by "
294 f"block size [{block_out}, {block_in}]"
295 )
296 add(f"{name}.weight", [out_dim, in_dim // 2], I8)
297 add(f"{name}.scale", [out_dim, in_dim // 32], scale_dtype)
298 elif kind == "fp8":
299 assert out_dim % block_out == 0 and in_dim % block_in == 0, (
300 f"{name} shape [{out_dim}, {in_dim}] is not divisible by "
301 f"block size [{block_out}, {block_in}]"
302 )
303 add(f"{name}.weight", [out_dim, in_dim], FP8)
304 add(
305 f"{name}.scale",
306 [out_dim // block_out, in_dim // block_in],
307 scale_dtype,
308 )
309 else:
310 add(f"{name}.weight", [out_dim, in_dim], BF16)
311
312 def compressor(name: str, ratio: int, size: int) -> None:
313 out_dim = size * (2 if ratio == 4 else 1)
314 add(f"{name}.ape", [ratio, out_dim], F32)
315 add(f"{name}.wkv.weight", [out_dim, dim], BF16)
316 add(f"{name}.wgate.weight", [out_dim, dim], BF16)
317 add(f"{name}.norm.weight", [size], BF16)
318
319 def attention(name: str, ratio: int) -> None:
320 add(f"{name}.attn_sink", [heads], F32)
321 linear(f"{name}.wq_a", q_rank, dim)
322 add(f"{name}.q_norm.weight", [q_rank], BF16)
323 linear(f"{name}.wq_b", heads * head_dim, q_rank)
324 linear(f"{name}.wkv", head_dim, dim)
325 add(f"{name}.kv_norm.weight", [head_dim], BF16)
326 linear(f"{name}.wo_a", o_groups * o_rank, heads * head_dim // o_groups)
327 linear(f"{name}.wo_b", dim, o_groups * o_rank)
328
329 if ratio:
330 compressor(f"{name}.compressor", ratio, head_dim)
331 if ratio == 4:
332 index_heads = config["index_n_heads"]
333 index_dim = config["index_head_dim"]
334 linear(f"{name}.indexer.wq_b", index_heads * index_dim, q_rank)
335 add(f"{name}.indexer.weights_proj.weight", [index_heads, dim], BF16)
336 compressor(f"{name}.indexer.compressor", ratio, index_dim)
337
338 def expert(name: str, kind: WeightKind) -> None:
339 linear(f"{name}.w1", inter, dim, kind)
340 linear(f"{name}.w2", dim, inter, kind)
341 linear(f"{name}.w3", inter, dim, kind)
342
343 def moe(name: str, layer_id: int) -> None:
344 add(f"{name}.gate.weight", [experts, dim], BF16)
345 if layer_id < config["num_hash_layers"]:
346 add(
347 f"{name}.gate.tid2eid",
348 [vocab, config["num_experts_per_tok"]],
349 I64,
350 )
351 else:
352 add(f"{name}.gate.bias", [experts], F32)
353
354 routed_kind = "fp4" if config.get("expert_dtype") == "fp4" else weight_kind
355 for expert_id in range(experts):
356 expert(f"{name}.experts.{expert_id}", routed_kind)
357 expert(f"{name}.shared_experts", weight_kind)
358
359 def block(name: str, layer_id: int) -> None:
360 attention(f"{name}.attn", config["compress_ratios"][layer_id])
361 moe(f"{name}.ffn", layer_id)
362 add(f"{name}.attn_norm.weight", [dim], BF16)
363 add(f"{name}.ffn_norm.weight", [dim], BF16)
364 for part in ("attn", "ffn"):
365 add(f"{name}.hc_{part}_fn", [(2 + hc) * hc, hc * dim], F32)
366 add(f"{name}.hc_{part}_base", [(2 + hc) * hc], F32)
367 add(f"{name}.hc_{part}_scale", [3], F32)
368
369 def hc_head(name: str = "") -> None:
370 prefix = f"{name}." if name else ""
371 add(f"{prefix}hc_head_fn", [hc, hc * dim], F32)
372 add(f"{prefix}hc_head_base", [hc], F32)
373 add(f"{prefix}hc_head_scale", [1], F32)
374
375 add("embed.weight", [vocab, dim], BF16)
376 for layer_id in range(config["num_hidden_layers"]):
377 block(f"layers.{layer_id}", layer_id)
378 add("norm.weight", [dim], BF16)
379 add("head.weight", [vocab, dim], BF16)
380 hc_head()
381
382 first_mtp_layer = config["num_hidden_layers"]
383 for mtp_id in range(config["num_nextn_predict_layers"]):
384 name = f"mtp.{mtp_id}"
385 block(name, first_mtp_layer + mtp_id)
386 linear(f"{name}.e_proj", dim, dim)
387 linear(f"{name}.h_proj", dim, dim)
388 for norm in ("enorm", "hnorm", "norm"):
389 add(f"{name}.{norm}.weight", [dim], BF16)
390 hc_head(name)
391
392 return initialize(state, config, init_bound)
393
394main()