Views
No views yet
1"""
2Create a tiny GigaChat3 model for testing .
3
4GigaChat3 uses DeepseekV3Config (no text_config/vision_config sub-objects).
5Key constraint: qk_head_dim == qk_nope_head_dim + qk_rope_head_dim
6"""
7import json
8import os
9
10import torch
11from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
12
13model_id = "ai-sage/GigaChat3-10B-A1.8B-bf16"
14output_dir = "./tiny-gigachat3"
15
16config = AutoConfig.from_pretrained(model_id)
17config.num_hidden_layers = 2
18config.num_attention_heads = 2
19config.num_key_value_heads = 2
20config.hidden_size = 32
21config.intermediate_size = 64
22config.moe_intermediate_size = 32
23config.n_routed_experts = 4
24config.n_shared_experts = 1
25config.num_experts_per_tok = 2
26config.kv_lora_rank = 8
27config.q_lora_rank = None
28
29# Attention head dims — MUST satisfy: qk_head_dim == qk_nope_head_dim + qk_rope_head_dim
30config.qk_nope_head_dim = 4
31config.qk_rope_head_dim = 2
32config.qk_head_dim = 6 # 4 + 2
33config.v_head_dim = 4
34config.head_dim = config.qk_rope_head_dim # used by RoPE
35
36TINY_VOCAB = 32000
37config.vocab_size = TINY_VOCAB
38
39assert config.qk_head_dim == config.qk_nope_head_dim + config.qk_rope_head_dim
40
41os.makedirs(output_dir, exist_ok=True)
42model = AutoModelForCausalLM.from_config(config)
43model.save_pretrained(output_dir)
44
45tokenizer = AutoTokenizer.from_pretrained(model_id)
46tokenizer.save_pretrained(output_dir)
47
48tok_path = os.path.join(output_dir, "tokenizer.json")
49with open(tok_path, encoding="utf-8") as f:
50 tok_data = json.load(f)
51
52if "model" in tok_data and "vocab" in tok_data["model"]:
53 tok_data["model"]["vocab"] = {
54 k: v for k, v in tok_data["model"]["vocab"].items() if v < TINY_VOCAB
55 }
56 tok_data["model"]["merges"] = []
57
58if "added_tokens" in tok_data:
59 tok_data["added_tokens"] = [t for t in tok_data["added_tokens"] if t["id"] < TINY_VOCAB]
60
61with open(tok_path, "w", encoding="utf-8") as f:
62 json.dump(tok_data, f, ensure_ascii=False)
63
64# ── Smoke test ───────────────────────────────────────────────────────────────
65tokens = tokenizer("Hello world", return_tensors="pt")
66tokens.pop("token_type_ids", None)
67with torch.no_grad():
68 out = model(**tokens)
69
70total_mb = sum(os.path.getsize(os.path.join(output_dir, fn)) for fn in os.listdir(output_dir)) / 1e6
71print(f"shape={out.logits.shape} params={sum(p.numel() for p in model.parameters())/1e6:.2f}M size={total_mb:.1f} MB")
72