| // 04 — the complete ArchitectureTemplate (all tiers together), lightly annotated. |
| // Faithful trimmed copy of artifacts/mistral.json (a modular, GQA, sliding-window decoder LLM). |
| // JSONC = JSON + // comments (not machine-parseable as-is); "…" marks omitted repetition. |
| // Regenerate the exact file: python utils/architecture_ir/generate_architecture_ir.py --architectures mistral --output-dir out |
| { |
| "schema_version": "architecture-template-v0", |
| "model_type": "mistral", |
|
|
| // modular inheritance (diff_size metric lives in modular_graph.json) |
| "extends": "llama", |
| "patches": [ |
| { "relation": "inherits", "target_class": "MistralAttention", "component_kind": "attention", |
| "parent_class": "LlamaAttention", "overridden": { "methods": ["__init__", "forward"] } }, |
| { "relation": "new", "target_class": "MistralForQuestionAnswering", "component_kind": null, |
| "parent_class": "MistralPreTrainedModel" } |
| // … other overridden classes (MistralMLP, MistralModel, …) |
| ], |
|
|
| // model-level facts (multimodal reads the text backbone) |
| "architecture": { |
| "view": "decoder", "family": "causal_lm", "attention_variant": "GQA", |
| "positional": "rope", "is_moe": false, "mixer": "attention", "sliding_window": 4096, "tie_word_embeddings": false |
| }, |
|
|
| // what it can do / run with |
| "capabilities": { |
| "attention_backends": ["eager", "sdpa", "flash_attention", "flex_attention"], |
| "attention_patterns": ["sliding"], "attention_schedule": null, |
| "task_heads": ["causal_lm", "question_answering", "sequence_classification", "token_classification"], |
| "tensor_parallel": true, |
| "kernels": { "RMSNorm": ["kernels-community/liger-kernels", "kernels-community/rmsnorm", "kernels-community/mlx_rmsnorm"] } |
| }, |
|
|
| // config identity + the parametric surface (full config NOT serialized) |
| "config": { |
| "class_name": "MistralConfig", "module": "transformers.models.mistral.configuration_mistral", "model_type": "mistral", |
| "referenced_fields": { "num_hidden_layers": 32 }, |
| "salient_fields": { "hidden_size": 4096, "intermediate_size": 14336, "vocab_size": 32000, |
| "num_attention_heads": 32, "num_key_value_heads": 8, "head_dim": 128, |
| "max_position_embeddings": 131072, "sliding_window": 4096, "hidden_act": "silu", |
| "tie_word_embeddings": false } |
| }, |
|
|
| // nodes outside any repeat body (leaf nodes omit `children`) |
| "components": [ |
| { "id": "model", "kind": "model", "class_name": "MistralModel", "path_pattern": "model", |
| "children": ["embed_tokens", "decoder_layers", "norm", "rotary_emb"] }, |
| { "id": "embed_tokens", "kind": "embedding", "class_name": "Embedding", "path_pattern": "model.embed_tokens", |
| "attributes": { "num_embeddings": "config.vocab_size", "embedding_dim": "config.hidden_size" } }, |
| { "id": "norm", "kind": "normalization", "class_name": "MistralRMSNorm", "path_pattern": "model.norm", |
| "attributes": { "norm_type": "rms", "kernel": "RMSNorm" } }, |
| { "id": "rotary_emb", "kind": "position", "class_name": "MistralRotaryEmbedding", "path_pattern": "model.rotary_emb", |
| "attributes": { "scheme": "rope", "rope_theta": 10000.0, "head_dim": "config.head_dim" } } |
| ], |
|
|
| // the repeated block body, serialized once ({i} = layer index) |
| "templates": [ |
| { "id": "decoder_layer", "kind": "transformer_block", "class_name": "MistralDecoderLayer", |
| "path_pattern": "model.layers.{i}", |
| "children": ["decoder_layer.input_layernorm", "decoder_layer.self_attn", |
| "decoder_layer.post_attention_layernorm", "decoder_layer.mlp"] }, |
| { "id": "decoder_layer.self_attn", "kind": "attention", "class_name": "MistralAttention", |
| "path_pattern": "model.layers.{i}.self_attn", |
| "children": ["decoder_layer.self_attn.q_proj", "decoder_layer.self_attn.k_proj", |
| "decoder_layer.self_attn.v_proj", "decoder_layer.self_attn.o_proj"], |
| "attributes": { "variant": "GQA", "n_heads": 32, "n_kv_heads": 8, "head_dim": 128, |
| "rope": true, "sliding_window": 4096, "pattern": "sliding" } }, |
| // GQA: q/o are hidden-sized; k/v are num_kv_heads*head_dim = 1024 |
| { "id": "decoder_layer.self_attn.q_proj", "kind": "projection", "class_name": "Linear", |
| "path_pattern": "model.layers.{i}.self_attn.q_proj", |
| "attributes": { "in_features": "config.hidden_size", "out_features": 4096, "tp": "colwise" } }, |
| { "id": "decoder_layer.self_attn.k_proj", "kind": "projection", "class_name": "Linear", |
| "path_pattern": "model.layers.{i}.self_attn.k_proj", |
| "attributes": { "in_features": "config.hidden_size", "out_features": 1024, "tp": "colwise" } }, |
| // … v_proj (1024, colwise), o_proj (4096 → config.hidden_size, rowwise) |
| { "id": "decoder_layer.mlp", "kind": "feed_forward", "class_name": "MistralMLP", |
| "path_pattern": "model.layers.{i}.mlp", |
| "children": ["decoder_layer.mlp.gate_proj", "decoder_layer.mlp.up_proj", "decoder_layer.mlp.down_proj"], |
| "attributes": { "hidden_size": 4096, "intermediate_size": 14336, "activation": "silu" } }, |
| { "id": "decoder_layer.mlp.gate_proj", "kind": "projection", "class_name": "Linear", |
| "path_pattern": "model.layers.{i}.mlp.gate_proj", |
| "attributes": { "in_features": "config.hidden_size", "out_features": "config.intermediate_size", "tp": "colwise" } }, |
| // … up_proj (colwise), down_proj (config.intermediate_size → config.hidden_size, rowwise) |
| { "id": "decoder_layer.input_layernorm", "kind": "normalization", "class_name": "MistralRMSNorm", |
| "path_pattern": "model.layers.{i}.input_layernorm", "attributes": { "norm_type": "rms", "kernel": "RMSNorm" } } |
| // … post_attention_layernorm |
| ], |
|
|
| // 32 identical layers collapsed to one symbolic entry, count kept parametric |
| "repeats": [ |
| { "id": "decoder_layers", "kind": "symbolic_repeat", "body": "decoder_layer", |
| "count_expr": "config.num_hidden_layers", "count": 32, "count_source": "config", "index_symbol": "i", |
| "container_path_pattern": "model.layers", "item_path_pattern": "model.layers.{i}", |
| "repeated_class_name": "MistralDecoderLayer" } |
| ], |
|
|
| // coarse dataflow. kinds: data | residual | mask | position | cross_attention | route | cache_read | cache_write |
| "edges": [ |
| { "source": "embed_tokens", "target": "decoder_layers", "kind": "data" }, |
| { "source": "decoder_layers", "target": "norm", "kind": "data" }, |
| // block-level flow (pre-norm), re-grounded from the observed forward: |
| { "source": "decoder_layer.input_layernorm", "target": "decoder_layer.self_attn", "kind": "data", "provenance": "observed_forward" }, |
| { "source": "decoder_layer.self_attn", "target": "decoder_layer.post_attention_layernorm", "kind": "data", "provenance": "observed_forward" }, |
| { "source": "decoder_layer.post_attention_layernorm", "target": "decoder_layer.mlp", "kind": "data", "provenance": "observed_forward" }, |
| // intra-module fan-out/fan-in (role-based; q/k/v parallel, not chained): |
| { "source": "decoder_layer.self_attn", "target": "decoder_layer.self_attn.q_proj", "kind": "data", "provenance": "intra_module" }, |
| { "source": "decoder_layer.self_attn.q_proj", "target": "decoder_layer.self_attn.o_proj", "kind": "data", "provenance": "intra_module" }, |
| { "source": "decoder_layer.mlp.gate_proj", "target": "decoder_layer.mlp.down_proj", "kind": "data", "provenance": "intra_module" }, |
| { "source": "decoder_layer", "target": "decoder_layer.self_attn", "kind": "residual" }, |
| { "source": "rotary_emb", "target": "decoder_layer.self_attn", "kind": "position" }, |
| { "source": "input:attention_mask", "target": "decoder_layer.self_attn", "kind": "mask" }, |
| { "source": "state:kv_cache", "target": "decoder_layer.self_attn", "kind": "cache_read" }, |
| { "source": "decoder_layer.self_attn", "target": "state:kv_cache", "kind": "cache_write" } |
| ], |
|
|
| // observed tensor shapes (symbolized), keyed by node id; order lives in edges |
| "dataflow": { |
| "source": "observed_forward_meta", |
| "input": { "name": "input_ids", "shape": ["B", "S"] }, |
| "output": { "shape": ["B", "S", "config.hidden_size"] }, |
| "shapes": { |
| "embed_tokens": { "in": ["B", "S"], "out": ["B", "S", "config.hidden_size"] }, |
| "decoder_layer.self_attn": { "in": ["B", "S", "config.hidden_size"], "out": ["B", "S", "config.hidden_size"] } |
| // … norm, mlp, the other block children |
| } |
| }, |
|
|
| // which classes this IR was introspected from (resolution strategy is invariant → in SPEC, not here) |
| "provenance": { |
| "config_class": "MistralConfig", "config_module": "transformers.models.mistral.configuration_mistral", |
| "model_class": "MistralModel", "model_module": "transformers.models.mistral.modeling_mistral" |
| } |
| } |
|
|
