yujiepan / inkling-tiny-random

import json
from pathlib import Path

import torch
from huggingface_hub import file_exists, hf_hub_download
from safetensors.torch import load_file, save_file
from transformers import (
    AutoConfig,
    AutoProcessor,
    GenerationConfig,
    InklingForConditionalGeneration,
    set_seed,
)

source_model_id = "thinkingmachines/Inkling"
save_folder = "/tmp/yujiepan/inkling-tiny-random"

processor = AutoProcessor.from_pretrained(source_model_id)
processor.save_pretrained(save_folder)

with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:
    config_json = json.load(f)

# Only shrink size-critical dims. Keep kernel-sensitive knobs (d_rel, rel_extent,
# sliding_window_size, num_experts_per_tok, n_shared_experts, ...) as upstream.
hidden_size = 8
num_mtp_layers = 1
config_json['text_config'].update({
    'hidden_size': hidden_size,
    'num_hidden_layers': 2,
    'num_attention_heads': 8,
    'num_key_value_heads': 4,
    'head_dim': 32,
    'swa_num_attention_heads': 8,
    'swa_num_key_value_heads': 4,
    'swa_head_dim': 32,
    'local_layer_ids': [0],  # keep 1 sliding + 1 global with 2 layers
    'dense_mlp_idx': 1,  # 1 dense + 1 sparse
    'dense_intermediate_size': 32,
    'intermediate_size': 32,
    'moe_intermediate_size': 32,
})
config_json['vision_config'].update({
    'decoder_dmodel': hidden_size,
    'n_layers': 2,
})
config_json['audio_config'].update({
    'decoder_dmodel': hidden_size,
})
config_json['mtp_config'].update({
    'num_nextn_predict_layers': num_mtp_layers,
    'local_layer_ids': [0],
})

with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
    json.dump(config_json, f, indent=2)

config = AutoConfig.from_pretrained(save_folder)
print(config)
torch.set_default_dtype(torch.bfloat16)
model = InklingForConditionalGeneration(config)
torch.set_default_dtype(torch.float32)
if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
    model.generation_config = GenerationConfig.from_pretrained(
        source_model_id, trust_remote_code=True,
    )
set_seed(42)
model = model.cpu()
num_params = sum(p.numel() for p in model.parameters())
with torch.no_grad():
    for name, p in sorted(model.named_parameters()):
        torch.nn.init.normal_(p, 0, 0.2)
        print(name, p.shape, f'{p.numel() / num_params:.2%}', f'{p.numel() * p.element_size() / 1024**2:.2f}MB')
# Upstream MoE gate bias / global_scale are F32; sconv stays BF16 in the checkpoint.
for name, module in model.named_modules():
    if hasattr(module, "e_score_correction_bias"):
        module.e_score_correction_bias = torch.nn.Parameter(
            module.e_score_correction_bias.detach().float()
        )
    if name.endswith(".mlp.gate") and hasattr(module, "global_scale"):
        module.global_scale = torch.nn.Parameter(module.global_scale.detach().float())
model.save_pretrained(save_folder)

# HF ignores `model.mtp.*` on main load; write them with original checkpoint naming.
set_seed(42)
path = Path(save_folder) / "model.safetensors"
state = load_file(str(path))
dense_prefix = "model.llm.layers.0."  # MTP blocks are dense
dense_keys = {k: v for k, v in state.items() if k.startswith(dense_prefix)}
for i in range(num_mtp_layers):
    block_prefix = f"model.mtp.layers.{i}.transformer_block."
    for src_key, tensor in dense_keys.items():
        dst_key = block_prefix + src_key[len(dense_prefix):]
        state[dst_key] = torch.empty_like(tensor)
        torch.nn.init.normal_(state[dst_key], 0, 0.2)
        print(dst_key, tuple(state[dst_key].shape))
    for name, shape in (
        (f"model.mtp.layers.{i}.embed_norm.weight", (hidden_size,)),
        (f"model.mtp.layers.{i}.hidden_norm.weight", (hidden_size,)),
        (f"model.mtp.layers.{i}.input_proj.weight", (hidden_size, hidden_size * 2)),
    ):
        state[name] = torch.empty(shape, dtype=torch.bfloat16)
        torch.nn.init.normal_(state[name], 0, 0.2)
        print(name, shape)
# Keep checkpoint key dtypes aligned even if save_pretrained downcasts.
for key, tensor in list(state.items()):
    if key.endswith(".mlp.gate.bias") or key.endswith(".mlp.gate.global_scale"):
        state[key] = tensor.float()
save_file(state, str(path))

InklingForConditionalGeneration(
  (model): InklingModel(
    (language_model): InklingTextModel(
      (embed_tokens): Embedding(201024, 8)
      (layers): ModuleList(
        (0): InklingDecoderLayer(
          (self_attn): InklingAttention(
            (q_proj): Linear(in_features=8, out_features=256, bias=False)
            (k_proj): Linear(in_features=8, out_features=128, bias=False)
            (v_proj): Linear(in_features=8, out_features=128, bias=False)
            (r_proj): Linear(in_features=8, out_features=128, bias=False)
            (o_proj): Linear(in_features=256, out_features=8, bias=False)
            (k_sconv): InklingShortConvolution(
              (conv1d): Conv1d(128, 128, kernel_size=(4,), stride=(1,), padding=(3,), groups=128, bias=False)
            )
            (v_sconv): InklingShortConvolution(
              (conv1d): Conv1d(128, 128, kernel_size=(4,), stride=(1,), padding=(3,), groups=128, bias=False)
            )
            (q_norm): InklingRMSNorm((32,), eps=1e-06)
            (k_norm): InklingRMSNorm((32,), eps=1e-06)
            (rel_logits_proj): InklingRelativeLogits()
          )
          (mlp): InklingMLP(
            (gate_proj): Linear(in_features=8, out_features=32, bias=False)
            (up_proj): Linear(in_features=8, out_features=32, bias=False)
            (down_proj): Linear(in_features=32, out_features=8, bias=False)
            (act_fn): SiLUActivation()
          )
          (input_layernorm): InklingRMSNorm((8,), eps=1e-06)
          (post_attention_layernorm): InklingRMSNorm((8,), eps=1e-06)
          (attn_sconv): InklingShortConvolution(
            (conv1d): Conv1d(8, 8, kernel_size=(4,), stride=(1,), padding=(3,), groups=8, bias=False)
          )
          (mlp_sconv): InklingShortConvolution(
            (conv1d): Conv1d(8, 8, kernel_size=(4,), stride=(1,), padding=(3,), groups=8, bias=False)
          )
        )
        (1): InklingDecoderLayer(
          (self_attn): InklingAttention(
            (q_proj): Linear(in_features=8, out_features=256, bias=False)
            (k_proj): Linear(in_features=8, out_features=128, bias=False)
            (v_proj): Linear(in_features=8, out_features=128, bias=False)
            (r_proj): Linear(in_features=8, out_features=128, bias=False)
            (o_proj): Linear(in_features=256, out_features=8, bias=False)
            (k_sconv): InklingShortConvolution(
              (conv1d): Conv1d(128, 128, kernel_size=(4,), stride=(1,), padding=(3,), groups=128, bias=False)
            )
            (v_sconv): InklingShortConvolution(
              (conv1d): Conv1d(128, 128, kernel_size=(4,), stride=(1,), padding=(3,), groups=128, bias=False)
            )
            (q_norm): InklingRMSNorm((32,), eps=1e-06)
            (k_norm): InklingRMSNorm((32,), eps=1e-06)
            (rel_logits_proj): InklingRelativeLogits()
          )
          (mlp): InklingMoE(
            (gate): InklingTopkRouter()
            (experts): InklingExperts(
              (act_fn): SiLUActivation()
            )
            (shared_experts): InklingSharedExperts(
              (act_fn): SiLUActivation()
            )
          )
          (input_layernorm): InklingRMSNorm((8,), eps=1e-06)
          (post_attention_layernorm): InklingRMSNorm((8,), eps=1e-06)
          (attn_sconv): InklingShortConvolution(
            (conv1d): Conv1d(8, 8, kernel_size=(4,), stride=(1,), padding=(3,), groups=8, bias=False)
          )
          (mlp_sconv): InklingShortConvolution(
            (conv1d): Conv1d(8, 8, kernel_size=(4,), stride=(1,), padding=(3,), groups=8, bias=False)
          )
        )
      )
      (norm): InklingRMSNorm((8,), eps=1e-06)
      (embed_norm): InklingRMSNorm((8,), eps=1e-06)
    )
    (audio_tower): InklingAudioModel(
      (embed_audio_tokens): InklingAudioModelEmbeddings(
        (embed_audio_tokens): Embedding(1280, 8)
      )
      (norm): InklingRMSNorm((8,), eps=1e-06)
    )
    (vision_tower): InklingVisionModel(
      (encoder_layers): ModuleList(
        (0): InklingVisionEncoderLayer(
          (projection): Linear(in_features=300, out_features=320, bias=False)
          (layer_norm): InklingRMSNorm((320,), eps=1e-06)
        )
        (1): InklingVisionEncoderLayer(
          (projection): Linear(in_features=10240, out_features=8, bias=False)
        )
      )
      (final_norm): InklingRMSNorm((8,), eps=1e-06)
    )
  )
  (lm_head): Linear(in_features=8, out_features=201024, bias=False)
)