| 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"}, |
| }, |
| ) |
| |
| |
| 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", |
| ) |
|
|
