File size: 8,790 Bytes
6f3c288 3335a58 6f3c288 | 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 | // 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"
}
}
|