def create_model(source_model_id: str, source_revision: str, save_folder: str): import json from pathlib import Path import torch from huggingface_hub import hf_hub_download from safetensors import safe_open from safetensors.torch import save_file from transformers import AutoConfig, AutoModelForImageTextToText folder = Path(save_folder) folder.mkdir(parents=True, exist_ok=True) torch.manual_seed(42) def read_json(name: str): return json.loads( Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_text(encoding="utf-8") ) source_config, params = read_json("config.json"), read_json("params.json") text = source_config["text_config"] text.update( architectures=["MistralLarge4ForCausalLM"], tie_word_embeddings=True, hidden_size=128, intermediate_size=128, moe_intermediate_size=128, num_hidden_layers=12, num_attention_heads=4, num_key_value_heads=4, head_dim=32, qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128, kv_lora_rank=512, index_head_dim=128, index_n_heads=4, index_topk=128, n_routed_experts=32, n_shared_experts=1, num_experts_per_tok=4, first_k_dense_replace=1, layer_types=["indexed_attention"] * 12, mlp_layer_types=["dense"] + ["sparse"] * 11, no_rope_layers=[1, 1, 1, 0] * 3, indexer_types=["full"] * 4 + ["full", "full", "shared", "shared"] * 2, kv_source_layers=[0, 1, 2, 3, 4, 5, 6, 3, 8, 9, 10, 3], ) source_config["tie_word_embeddings"] = True source_config["vision_config"].update( hidden_size=64, intermediate_size=128, num_hidden_layers=4, num_attention_heads=2, head_dim=32, image_size=1540, patch_size=14, ) source_config["quantization_config"] = { "quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [64, 64], "modules_to_not_convert": [ "model.vision_tower", "model.multi_modal_projector", "lm_head", "kv_b_proj", "indexer.wk", "indexer.weights_proj", "kv_a_proj_with_mqa", ], } source_config["transformers_weights"] = "consolidated.safetensors" params.update( dim=128, tied_embeddings=True, n_layers=12, head_dim=32, hidden_dim=128, n_heads=4, n_kv_heads=4, kv_lora_rank=512, qk_rope_head_dim=64, qk_nope_head_dim=64, v_head_dim=128, nope="rope,rope,rope,nope", kv_sharing={"source_layers_str": "0,1,2,3,4,5,6,3,8,9,10,3"}, dsa={ "topk": 128, "n_heads": 4, "head_dim": 128, "rope_head_dim": 64, "match_nope": True, "ssmax_scaling": True, "norm_eps": 1e-6, "sharing": "full,full,full,full,full,full,shared,shared,full,full,shared,shared", "fp8_indexer": True, "fp8_kv_cache": True, }, quantization={ "qscheme_act": "BLOCK", "qformat_weight": "fp8_e4m3", "qdynamic_act": True, "qscheme_weight": "BLOCK", "use_ue8m0_for_qscales": True, "block_size_out": 64, "block_size_in": 64, "quantize_kv_a": False, }, moe={ "expert_parallel": 1, "expert_model_parallel": 1, "route_every_n": 1, "first_k_dense_replace": 1, "num_experts": 32, "num_experts_per_tok": 4, "num_expert_groups": 1, "num_expert_groups_per_tok": 1, "routed_scale": 1.0, "expert_hidden_dim": 128, "num_shared_experts": 1, }, vision_encoder={ **params["vision_encoder"], "intermediate_size": 128, "num_hidden_layers": 4, "num_attention_heads": 2, "hidden_size": 64, "image_size": 1540, "max_image_size": 1540, "patch_size": 14, }, mtp={ "n_layers": 3, "depth": 3, "ragged_attention": "32,32,32", "target_activation_layers": "1,4,7,10,11", "norm_strategy": "eagle", "dspark": {"markov_rank": 4, "head_type": "rnn", "draft_input": "position_specific"}, }, ) # The public HF config is used by Transformers while vLLM also needs this # native Mistral block to build and enable the in-checkpoint DSpark draft. text["mtp_config"] = params["mtp"] config_for_build = dict(source_config) config_for_build.pop("quantization_config") (folder / "config.json").write_text(json.dumps(config_for_build, indent=2) + "\n", encoding="utf-8") model = AutoModelForImageTextToText.from_config(AutoConfig.from_pretrained(folder), dtype=torch.bfloat16) with torch.no_grad(): for name, value in model.named_parameters(): if name.endswith("weight") and "norm" in name: value.fill_(1) elif name.endswith("bias"): value.zero_() else: value.normal_(mean=0.0, std=0.02) model.tie_weights() model.save_pretrained(folder, safe_serialization=True) state = {} with safe_open(folder / "model.safetensors", framework="pt", device="cpu") as reader: for name in reader.keys(): state[name] = reader.get_tensor(name).contiguous() (folder / "model.safetensors").unlink() def should_quantize(name: str, value: torch.Tensor) -> bool: return ( name.endswith(".weight") and value.ndim == 2 and value.shape[0] % 64 == 0 and value.shape[1] % 64 == 0 and not name.startswith( ( "vision_encoder.", "vision_language_adapter.", "patch_merger.", "pre_mm_projector_norm.", "tok_embeddings.", "output.", ) ) and not any( piece in name for piece in ( ".wkv_a_with_mqa.", ".wkv_b.", ".indexer.wk.", ".indexer.weights_proj.", ".attention.wo.", ".feed_forward.", ".shared_experts.", ".experts.", ) ) and not name.endswith(".gate.weight") ) def quantize(name: str, value: torch.Tensor) -> None: rows, columns = value.shape blocks = value.float().reshape(rows // 64, 64, columns // 64, 64) maximum = blocks.abs().amax(dim=(1, 3)).clamp_min(torch.finfo(torch.float32).tiny) inverse = torch.pow(2.0, torch.ceil(torch.log2(1.0 / maximum))) state[name] = ( (blocks * inverse.float()[:, None, :, None]) .clamp(-448, 448) .reshape(rows, columns) .to(torch.float8_e4m3fn) .contiguous() ) state[name.removesuffix(".weight") + ".qscale_weight"] = (1.0 / inverse).contiguous() for name, value in list(state.items()): if should_quantize(name, value): quantize(name, value) for layer in range(3): prefix = f"mtp_dspark.backbone.layers.{layer}" shapes = { "attention_norm.weight": (128,), "attention.wq.weight": (512, 128), "attention.wkv_a_with_mqa.weight": (576, 128), "attention.kv_a_norm.weight": (512,), "attention.wkv_b.weight": (768, 512), "attention.wo.weight": (128, 512), "ffn_norm.weight": (128,), "feed_forward.w1.weight": (128, 128), "feed_forward.w2.weight": (128, 128), "feed_forward.w3.weight": (128, 128), } for name, shape in shapes.items(): state[f"{prefix}.{name}"] = ( torch.ones(shape, dtype=torch.bfloat16) if name.endswith("norm.weight") else torch.randn(shape, dtype=torch.bfloat16) * 0.02 ) state.update( { "mtp_dspark.backbone.target_projection.weight": torch.randn(128, 640, dtype=torch.bfloat16) * 0.02, "mtp_dspark.backbone.target_norm.weight": torch.ones(128, dtype=torch.bfloat16), "mtp_dspark.backbone.final_norm.weight": torch.ones(128, dtype=torch.bfloat16), "mtp_dspark.backbone.mask_embedding": torch.randn(2, 128, dtype=torch.bfloat16) * 0.02, "mtp_dspark.autoregressive_head.embedding.weight": torch.randn(204800, 4, dtype=torch.bfloat16) * 0.02, "mtp_dspark.autoregressive_head.output.weight": torch.randn(204800, 4, dtype=torch.bfloat16) * 0.02, "mtp_dspark.autoregressive_head.joint_proj.weight": torch.randn(12, 136, dtype=torch.bfloat16) * 0.02, "mtp_dspark.autoregressive_head.joint_proj.bias": torch.zeros(12, dtype=torch.bfloat16), "lerope.alpha": torch.zeros(32, dtype=torch.bfloat16), } ) save_file(state, folder / "consolidated.safetensors", metadata={"format": "pt"}) (folder / "config.json").write_text(json.dumps(source_config, indent=2) + "\n", encoding="utf-8") (folder / "params.json").write_text(json.dumps(params, indent=2) + "\n", encoding="utf-8") for name in ( "tekken.json", "tokenizer.json", "tokenizer_config.json", "processor_config.json", "chat_template.jinja", "generation_config.json", "LICENSE", ): (folder / name).write_bytes(Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_bytes()) (folder / "vllm_mistral_large4_tied.py").write_text( "from vllm.distributed import get_pp_group\n" "from vllm.models.mistral_large4 import MistralLarge4ForCausalLM\n\n\n" "class MistralLarge4TiedForCausalLM(MistralLarge4ForCausalLM):\n" ' def __init__(self, *, vllm_config, prefix=""):\n' " super().__init__(vllm_config=vllm_config, prefix=prefix)\n" " if self.config.tie_word_embeddings and get_pp_group().is_last_rank:\n" " self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)\n", encoding="utf-8", ) (folder / "vllm_plugin.py").write_text( "def register():\n" " from vllm import ModelRegistry\n\n" " ModelRegistry.register_model(\n" ' "MistralLarge4ForCausalLM",\n' ' "vllm_mistral_large4_tied:MistralLarge4TiedForCausalLM",\n' " )\n", encoding="utf-8", ) (folder / "pyproject.toml").write_text( "[build-system]\n" "requires = ['setuptools>=61']\n" "build-backend = 'setuptools.build_meta'\n\n" "[project]\n" "name = 'mistral-large-4-tied-vllm'\n" "version = '0.1.0'\n\n" "[project.entry-points.'vllm.general_plugins']\n" "mistral_large4_tied = 'vllm_plugin:register'\n\n" "[tool.setuptools]\n" "py-modules = ['vllm_mistral_large4_tied', 'vllm_plugin']\n", encoding="utf-8", ) model.add_module( "mtp_dspark", torch.nn.ModuleDict( {"backbone": torch.nn.ModuleDict({"layers": torch.nn.ModuleList([torch.nn.Identity() for _ in range(3)])})} ), ) return str(model) create_model( source_model_id="mistralai/Mistral-Large-4-1T-A52B", source_revision="f12dabdc0c1e99e7f7924461a6f308859b05af4d", save_folder="/tmp/mistral-large-4", )