Deploy architecture inspector

examples/01-structure.jsonc

@@ -0,0 +1,255 @@
+// 01 — structure: components, edges, dataflow + per-node facts (the annotated diagram).
+// Trimmed mistral (real values); "…" marks omitted repetition. Tiers 02–04 add to this.
+{
+  "schema_version": "architecture-template-v0",
+  "model_type": "mistral",
+  "architecture": {
+    "view": "decoder",
+    "family": "causal_lm",
+    "attention_variant": "GQA",
+    "positional": "rope",
+    "is_moe": false,
+    "sliding_window": 4096,
+    "tie_word_embeddings": false
+  },
+  "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,
+      "num_attention_heads": 32,
+      "num_key_value_heads": 8,
+      "head_dim": 128,
+      "vocab_size": 32000,
+      "sliding_window": 4096
+    }
+  },
+  "components": [
+    {
+      "id": "model",
+      "kind": "model",
+      "class_name": "MistralModel",
+      "path_pattern": "model",
+      "children": [
+        "embed_tokens",
+        "decoder_layers",
+        "norm",
+        "rotary_emb"
+      ]
+    },
+    // "decoder_layers" is the repeat id
+    {
+      "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"
+      }
+    }
+  ],
+  "templates": [
+    {
+      "id": "decoder_layer",
+      "kind": "transformer_block",
+      "class_name": "MistralDecoderLayer",
+      "path_pattern": "model.layers.{i}",
+      // {i} = layer index
+      "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.o_proj"
+      ],
+      // … k_proj, v_proj
+      "attributes": {
+        "variant": "GQA",
+        "n_heads": 32,
+        "n_kv_heads": 8,
+        "head_dim": 128,
+        "rope": true,
+        "sliding_window": 4096,
+        "pattern": "sliding"
+      }
+    },
+    {
+      "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.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.input_layernorm",
+      "kind": "normalization",
+      "class_name": "MistralRMSNorm",
+      "path_pattern": "model.layers.{i}.input_layernorm",
+      "attributes": {
+        "norm_type": "rms",
+        "kernel": "RMSNorm"
+      }
+    }
+    // … post_attention_layernorm, and the k/v/o + gate/up/down projection nodes
+  ],
+  "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"
+    }
+    // 32 identical layers → one symbolic entry
+  ],
+  "edges": [
+    {
+      "source": "embed_tokens",
+      "target": "decoder_layers",
+      "kind": "data"
+    },
+    {
+      "source": "decoder_layers",
+      "target": "norm",
+      "kind": "data"
+    },
+    {
+      "source": "decoder_layer.input_layernorm",
+      "target": "decoder_layer.self_attn",
+      "kind": "data",
+      "provenance": "observed_forward"
+    },
+    {
+      "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"
+    }
+    // … kinds: data | residual | mask | position | cross_attention | route | cache_read | cache_write
+  ],
+  "dataflow": {
+    "source": "observed_forward_meta",
+    "input": {
+      "name": "input_ids",
+      "shape": [
+        "B",
+        "S"
+      ]
+    },
+    // B=batch, S=sequence
+    "output": {
+      "shape": [
+        "B",
+        "S",
+        "config.hidden_size"
+      ]
+    },
+    "shapes": {
+      // per-node observed shapes (symbolized); node ORDER lives in edges
+      "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"
+        ]
+      }
+    }
+  },
+  "provenance": {
+    "config_class": "MistralConfig",
+    "config_module": "transformers.models.mistral.configuration_mistral",
+    "model_class": "MistralModel",
+    "model_module": "transformers.models.mistral.modeling_mistral"
+  }
+}

examples/02-capabilities.jsonc

@@ -0,0 +1,62 @@
+// 02 — everything in 01, plus `capabilities`: what the model can do / run with.
+// (structure/dataflow trimmed here to keep the focus on the new block — see 01 for the full diagram.)
+{
+  "schema_version": "architecture-template-v0",
+  "model_type": "mistral",
+  "architecture": {
+    "view": "decoder", "family": "causal_lm", "attention_variant": "GQA",
+    "positional": "rope", "is_moe": false, "sliding_window": 4096, "tie_word_embeddings": false
+  },
+
+  // ▼▼▼ the tier-02 addition ▼▼▼
+  "capabilities": {
+    "attention_backends": ["eager", "sdpa", "flash_attention", "flex_attention"],  // "can run with", not "installed"
+    "attention_patterns": ["sliding"],       // distinct per-layer mask kinds
+    "attention_schedule": null,              // raw config.layer_types when non-uniform, else null
+    "task_heads": ["causal_lm", "question_answering", "sequence_classification", "token_classification"],
+    "tensor_parallel": true,                 // → projection nodes carry attributes.tp (colwise/rowwise)
+    "kernels": {                             // kernelizable layers → compatible Hub kernel repos
+      "RMSNorm": ["kernels-community/liger-kernels", "kernels-community/rmsnorm", "kernels-community/mlx_rmsnorm"]
+    }
+  },
+  // ▲▲▲ nodes point in via attributes.kernel: "RMSNorm" ▲▲▲
+
+  "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, "num_key_value_heads": 8, "head_dim": 128 }
+  },
+  "components": [
+    { "id": "model", "kind": "model", "class_name": "MistralModel", "path_pattern": "model",
+      "children": ["embed_tokens", "decoder_layers", "norm"] },
+    { "id": "norm", "kind": "normalization", "class_name": "MistralRMSNorm", "path_pattern": "model.norm",
+      "attributes": { "norm_type": "rms", "kernel": "RMSNorm" } }  // ← joins to capabilities.kernels["RMSNorm"]
+    // … embed_tokens, rotary_emb (see 01)
+  ],
+  "templates": [
+    { "id": "decoder_layer.self_attn", "kind": "attention", "class_name": "MistralAttention",
+      "path_pattern": "model.layers.{i}.self_attn",
+      "attributes": { "variant": "GQA", "n_heads": 32, "n_kv_heads": 8, "head_dim": 128, "pattern": "sliding" } }
+    // … decoder_layer + its projection / norm / mlp children (see 01)
+  ],
+  "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" }
+  ],
+  "edges": [
+    { "source": "embed_tokens", "target": "decoder_layers", "kind": "data" },
+    { "source": "decoder_layers", "target": "norm", "kind": "data" }
+  ],
+  "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"] } }
+  },
+  "provenance": {
+    "config_class": "MistralConfig", "config_module": "transformers.models.mistral.configuration_mistral",
+    "model_class": "MistralModel", "model_module": "transformers.models.mistral.modeling_mistral"
+  }
+}

examples/03-modularity.jsonc

@@ -0,0 +1,61 @@
+// 03 — everything in 01+02, plus modular inheritance: `extends` + `patches`.
+// (structure/capabilities trimmed here — see 01/02.) Standalone models have extends=null and no patches.
+{
+  "schema_version": "architecture-template-v0",
+  "model_type": "mistral",
+
+  // ▼▼▼ the tier-03 addition ▼▼▼
+  "extends": "llama",           // dominant parent model_type (from modular_mistral.py)
+  "patches": [                  // per-class changes vs the parent; trivial (unchanged) classes are omitted
+    { "relation": "inherits", "target_class": "MistralAttention", "component_kind": "attention",
+      "parent_class": "LlamaAttention", "overridden": { "methods": ["__init__", "forward"] } },
+    { "relation": "inherits", "target_class": "MistralMLP", "component_kind": "feed_forward",
+      "parent_class": "LlamaMLP", "overridden": { "methods": ["__init__"] } },
+    { "relation": "new", "target_class": "MistralForQuestionAnswering", "component_kind": null,
+      "parent_class": "MistralPreTrainedModel" }
+    // … relation: "inherits" | "new"; buckets: overridden / added / deleted, each { methods?, attrs? }
+  ],
+  // ▲▲▲ the diff_size metric is NOT here — it lives per-node in modular_graph.json (see bottom) ▲▲▲
+
+  "architecture": { "view": "decoder", "family": "causal_lm", "attention_variant": "GQA", "positional": "rope" },
+  "capabilities": { "attention_backends": ["eager", "sdpa", "flash_attention", "flex_attention"],
+                    "kernels": { "RMSNorm": ["kernels-community/rmsnorm"] } },
+  "config": {
+    "class_name": "MistralConfig", "module": "transformers.models.mistral.configuration_mistral", "model_type": "mistral",
+    "referenced_fields": { "num_hidden_layers": 32 }
+  },
+  "components": [
+    { "id": "model", "kind": "model", "class_name": "MistralModel", "path_pattern": "model",
+      "children": ["embed_tokens", "decoder_layers", "norm"] }
+    // … see 01
+  ],
+  "templates": [
+    { "id": "decoder_layer.self_attn", "kind": "attention", "class_name": "MistralAttention",
+      "path_pattern": "model.layers.{i}.self_attn", "attributes": { "variant": "GQA", "pattern": "sliding" } }
+    // … see 01
+  ],
+  "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" }
+  ],
+  "edges": [ { "source": "embed_tokens", "target": "decoder_layers", "kind": "data" } ],
+  "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"] } }
+  },
+  "provenance": {
+    "config_class": "MistralConfig", "config_module": "transformers.models.mistral.configuration_mistral",
+    "model_class": "MistralModel", "model_module": "transformers.models.mistral.modeling_mistral"
+  }
+}
+
+// Companion library-wide file `modular_graph.json` (one for the whole library, not per model):
+//   { "schema_version": "modular-graph-v0",
+//     "roots": ["llama", "vit", "clip"],
+//     "nodes": { "mistral": { "extends": "llama", "children": ["mixtral"], "root": "llama",
+//                             "depth": 1, "is_modular": true, "diff_size": 9 } } }
+// → models sharing a `root` are the same lineage (align cleanly for comparison); diff_size = modularity metric.

examples/04-full.jsonc

@@ -0,0 +1,133 @@
+// 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, "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"
+  }
+}