A
weight-transferred small replica of the DeepSeek-V4 architecture. Tensors
were sliced (or, where shapes already matched, copied verbatim) from
deepseek-ai/DeepSeek-V4-Flash into a smaller architecture
matching
kshitijthakkar/deepseek-v4-mini-1B-init.
The exhaustive list of COPY (frozen-recommended) keys is in
frozen_keys.json.
1# 1) install + auth (private repo)
2# !pip install -q "transformers>=4.50" huggingface_hub safetensors
3from huggingface_hub import login, snapshot_download
4login()
5
6# 2) download repo (weights + tokenizer + modeling package)
7local = snapshot_download(repo_id="kshitijthakkar/deepseek-v4-mini-1B-from-flash")
8
9# 3) register the deepseek_v4 model_type with HF auto classes
10import sys, os
11sys.path.insert(0, os.path.join(local, "code"))
12import deepseek_v4 # noqa: F401 — side-effect: AutoConfig / AutoModelForCausalLM registration
13
14# 4) load via the standard Auto API
15import torch
16from transformers import AutoModelForCausalLM, AutoTokenizer
17
18tok = AutoTokenizer.from_pretrained(local)
19model = AutoModelForCausalLM.from_pretrained(local, torch_dtype=torch.bfloat16)
20model.eval()
21
22# 5) chat-templated forward pass — initial loss is much lower than the
23# random-init scaffold, but the compressor/indexer paths are still random
24# until you fine-tune (see below).
25messages = [{"role": "user", "content": "Hello"}]
26ids = tok.apply_chat_template(messages, return_tensors="pt",
27 add_generation_prompt=True, return_dict=True)
28with torch.no_grad():
29 print(model(input_ids=ids["input_ids"]).logits.shape)
Because most non-MoE-expert weights are already well-trained, freeze them and
train only the lossy / random-init subset:
1from training.warm_start import apply_warm_start_freeze
2import os, json
3
4frozen_path = os.path.join(local, "frozen_keys.json")
5summary = apply_warm_start_freeze(model, frozen_path)
6print(summary)
7# -> trainable=Y M frozen=X M total=N M train_ratio=...%
8
9# Build optimizer over the trainable subset (configure_optimizers already
10# respects requires_grad)
11from training.optimizer import configure_optimizers
12optimizer = configure_optimizers(
13 model, learning_rate=2e-4, weight_decay=0.1,
14 betas=(0.9, 0.95), device_type="cuda", optimizer_type="muon",
15)
16
17# ... standard training loop ...
1@misc{deepseek_v4_2026,
2 author = {DeepSeek-AI},
3 title = {DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence},
4 year = {2026},
5 url = {https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash}
6}