File size: 6,671 Bytes
6eb3ab3
 
 
1b96d67
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
---
license: apache-2.0
---

```python

import json
import sys
from pathlib import Path

import torch
from transformers import (
    AutoConfig,
    AutoProcessor,
    GenerationConfig,
)
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForConditionalGeneration


SEED = 42
MODEL_ID = "openbmb/MiniCPM-OCR-Jev"
OUTPUT_DIR = Path(__file__).resolve().parent / "tiny-random-minicpm-ocr-jev"

# Marker written into the cached config so a stale/incompatible cache is rebuilt
# instead of silently reused.
CACHE_MARKER = "tiny-minicpm-ocr-jev-v1"

MAX_PARAMS = 40_000_000
MAX_BYTES = 180 * 1024 * 1024


def build_tiny_config():
    """Load the real config and shrink every scale parameter in place.

    The full ``vocab_size`` (248320) is preserved so the real tokenizer and the
    vision/image special-token ids (image/video/vision-start/vision-end) stay
    valid for the image-text-to-text path.
    """
    cfg = AutoConfig.from_pretrained(MODEL_ID)

    tc = cfg.text_config
    tc.hidden_size = 64
    tc.intermediate_size = 128
    tc.num_hidden_layers = 4
    tc.num_attention_heads = 4
    tc.num_key_value_heads = 2
    tc.head_dim = 32
    # Preserve the real linear/full attention mix (gated-delta-net + attention),
    # with one full_attention layer like the full_attention_interval=4 pattern.
    tc.layer_types = ["linear_attention", "linear_attention", "linear_attention", "full_attention"]
    tc.linear_num_key_heads = 2
    tc.linear_num_value_heads = 4
    tc.linear_key_head_dim = 16
    tc.linear_value_head_dim = 16
    # keep tc.linear_conv_kernel_dim (4)
    tc.max_position_embeddings = 4096
    # keep tc.use_cache (original ships False; generation still works at this scale)

    # partial_rotary_factor 0.25 * head_dim(32) -> rotary dim 8 -> freqs len 4;
    # interleaved mrope_section must sum to rotary_dim // 2 == 4.
    if isinstance(tc.rope_parameters, dict):
        tc.rope_parameters["mrope_section"] = [2, 1, 1]
        tc.rope_parameters["partial_rotary_factor"] = 0.25

    vc = cfg.vision_config
    vc.depth = 2
    vc.hidden_size = 64
    vc.intermediate_size = 128
    vc.num_heads = 4
    # out_hidden_size must match the text hidden_size (image embeds merge into
    # the decoder embedding space).
    vc.out_hidden_size = tc.hidden_size
    # num_position_embeddings must be a perfect square (num_grid_per_side =
    # int(sqrt(num_position_embeddings))); 256 -> 16 per side.
    vc.num_position_embeddings = 256
    # keep patch_size(16), spatial_merge_size(2), temporal_patch_size(2), in_channels(3)

    # Force float32 at every effective precision level for a robust CPU forward.
    cfg.dtype = "float32"
    tc.dtype = "float32"
    vc.dtype = "float32"

    return cfg


def _count_and_measure(model):
    n_params = 0
    n_bytes = 0
    for p in model.parameters():
        n_params += p.numel()
        n_bytes += p.numel() * p.element_size()
    for b in model.buffers():
        n_bytes += b.numel() * b.element_size()
    return n_params, n_bytes


def _cache_is_valid():
    cfg_path = OUTPUT_DIR / "config.json"
    weights = (OUTPUT_DIR / "model.safetensors").exists() or (OUTPUT_DIR / "pytorch_model.bin").exists()
    if not cfg_path.exists() or not weights:
        return False
    try:
        data = json.loads(cfg_path.read_text())
    except Exception:
        return False
    if data.get("_tiny_cache_marker") != CACHE_MARKER:
        return False
    if data.get("model_type") != "qwen3_5":
        return False
    if data.get("architectures") != ["Qwen3_5ForConditionalGeneration"]:
        return False
    return True


def create_tiny_model():
    if _cache_is_valid():
        print(f"Reusing valid cache at {OUTPUT_DIR}")
        return OUTPUT_DIR

    torch.manual_seed(SEED)

    config = build_tiny_config()

    # Instantiate the real architecture class with random weights.
    with torch.device("cpu"):
        model = Qwen3_5ForConditionalGeneration(config)
    model = model.to(torch.float32)
    model.eval()

    n_params, n_bytes = _count_and_measure(model)
    print(f"params={n_params:,} bytes={n_bytes / 1024 / 1024:.1f}MB")
    if n_params > MAX_PARAMS:
        raise RuntimeError(f"Parameter budget exceeded: {n_params:,} > {MAX_PARAMS:,}")
    if n_bytes > MAX_BYTES:
        raise RuntimeError(f"Memory budget exceeded: {n_bytes:,} > {MAX_BYTES:,}")

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(OUTPUT_DIR, safe_serialization=True)

    # Save processor assets (tokenizer, image processor, chat template).
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    processor.save_pretrained(OUTPUT_DIR)

    # Save a usable generation config.
    try:
        gen = GenerationConfig.from_pretrained(MODEL_ID)
    except Exception:
        gen = GenerationConfig()
    gen.save_pretrained(OUTPUT_DIR)

    # Stamp the cache marker and assert float32 precision fields survived.
    cfg_path = OUTPUT_DIR / "config.json"
    data = json.loads(cfg_path.read_text())
    data["_tiny_cache_marker"] = CACHE_MARKER
    data["dtype"] = "float32"
    if isinstance(data.get("text_config"), dict):
        data["text_config"]["dtype"] = "float32"
    if isinstance(data.get("vision_config"), dict):
        data["vision_config"]["dtype"] = "float32"
        # ``Qwen3_5VisionConfig.model_type`` class attribute is "qwen3_5", so
        # ``to_dict()`` drops the real "qwen3_5_vision" marker; restore it.
        data["vision_config"]["model_type"] = "qwen3_5_vision"
    cfg_path.write_text(json.dumps(data, indent=2))

    print(f"Saved tiny model to {OUTPUT_DIR}")
    return OUTPUT_DIR


def verify_identity():
    """Assert the saved config preserves the real architecture identity."""
    orig = AutoConfig.from_pretrained(MODEL_ID)
    tiny = AutoConfig.from_pretrained(OUTPUT_DIR)

    assert tiny.model_type == orig.model_type == "qwen3_5", tiny.model_type
    assert tiny.architectures == orig.architectures == ["Qwen3_5ForConditionalGeneration"]
    assert tiny.text_config.model_type == orig.text_config.model_type == "qwen3_5_text"
    assert tiny.vision_config.model_type == orig.vision_config.model_type == "qwen3_5_vision"
    assert tiny.image_token_id == orig.image_token_id
    assert tiny.vision_start_token_id == orig.vision_start_token_id
    assert tiny.text_config.vocab_size == orig.text_config.vocab_size
    # linear + full attention mix preserved
    assert "linear_attention" in tiny.text_config.layer_types
    assert "full_attention" in tiny.text_config.layer_types
    print("Architecture identity OK")


if __name__ == "__main__":
    out = create_tiny_model()
    verify_identity()
    print(out)

```