quark-onnx-subgraph-partitioner

Files (3)

• SKILL.md
• inspect_model.py
• validate_partition.py

SKILL.md (first lines)

# quark-onnx-subgraph-partitioner

## Purpose

Produce a `subgraph_partition.json` that groups ONNX graph nodes into named
functional blocks (stem, backbone stage, attention layer, FFN, FPN neck, etc.).
The partition can serve as input to AMP sensitivity analysis, layer-wise
profiling, partial deployment, or graph documentation.

Two helper scripts live alongside this file:

```bash
SKILL_DIR=skills/_legacy_impl/l1-atomic/onnx/quark-onnx-subgraph-partitioner
```

## Inputs

- Model path — a single `.onnx` file supplied as `$ARGUMENTS`
- Optional architecture hints from the user (layer count, naming convention,
  block granularity preference)

## Outputs: subgraph_partition.json

Schema: [`subgraph_partition.schema.json`](../../shared/contracts/subgraph_partition.schema.json)

```json
{
  "quantized":     false,
  "num_subgraphs": 4,
  "subgraphs": [
    {
      "name": "backbone_stem",
      "description": "Initial stride-2 Conv → SiLU activation, produces P1 features",
      "start_nodes": ["/model.0/conv/Conv"],
      "end_nodes":   ["/model.0/act/Mul"]
    },
    {
      "name": "attention_layer_0",
      "description": "Full encoder layer 0: self-attention + FFN + LayerNorm + residuals",
      "start_nodes": ["/encoder/layer.0/attention/self/query/MatMul"],
      "end_nodes":   ["/encoder/layer.0/output/LayerNorm"]
    }
  ]
}
```

Every entry requires all four fields: `name`, `description`, `start_nodes`,
`end_nodes`. **Never include `resolved_nodes`** — node resolution happens at
runtime. `num_subgraphs` must equal `len(subgraphs)`; set it last.

## Analysis Procedure

### Step 1 — Load and inspect the model

```bash
python3 "$SKILL_DIR/inspect_model.py" <MODEL_PATH>
```

For quantized models (detected automatically via Q/DQ op presence), Q/DQ
wrapper nodes are suppressed so structural landmarks remain visible; Q/DQ
counts are reported separately.

**Text `.onnxtxt`** — use `grep` instead:

```bash
grep -n 'op_type:\|  name:\|input:\|output:' <MODEL_PATH> | head -400
```

**Quantized model rules:** use compute nodes (`Conv`, `MatMul`, `Add`, …) as
boundaries — never Q/DQ nodes. BFS from a compute node captures surrounding
Q/DQ wrappers automatically.

### Step 2 — Identify structural landmarks

| Signal | Likely boundary |
|--------|-----------------|
| First `Conv` consuming the graph input | Backbone / model stem |
| Repeated `/model.N/` or `/layer.N/` prefix groups | Stage or layer repetition |
| `MaxPool` or strided `Conv` (stride 2) | Resolution downsampling |
| `Resize` (upsample) + `Concat` on lateral path | FPN / PAN top-down merge |
| `Split` → N parallel `Conv` → `Concat` + 1×1 `Conv` | C2f / CSP bottleneck |
| `GlobalAveragePool` → FC → activation → `Mul` | SE block |
| `MatMul`×3 → `Softmax` → `MatMul` (QKV pattern) | Self-attention block |
| `LayerNormalization` or `InstanceNormalization` | Transformer / diffusion layer |
| `MatMul`/`Gemm` → activation → `MatMul`/`Gemm` (hidden × 4) | FFN / MLP block |
| `Add` immediately after attention + FFN | Residual add (layer end) |
| `GridSample` with sampling offsets | Deformable attention |
| `Gather` on index 0 → `LayerNorm` → `Gemm` | CLS token / pooler / head |

### Step 3 — Assign blocks

| Block type | Description |
|------------|-------------|
| `Stem` | Initial convs before first downsampling |
| `BackboneStage` | Stage at one resolution (N residual / bottleneck blocks) |
| `SEBlock` | Squeeze-and-Excitation (GAP → FC → activation → Mul) |
| `PatchEmbedding` | ViT patch projection (Conv with patch-stride kernel) |
| `SelfAttentionBlock` | Multi-head self-attention (QKV + Softmax + out proj) |
| `FFNBlock` | Feed-forward network (Linear → activation → Linear) |
| `TransformerLayer` | Full encoder/decoder layer (attn + FFN + norms + residuals) |
| `SPPF` | Spatial Pyramid Pooling Fast |
| `FPNNeck` | Feature Pyramid Network (lateral convs + upsample + merge) |
| `PANNeck` | Path Aggregation Network (down-path convs + concat + merge) |
| `DetectionHead` | Conv layers + decode logic for bbox / class output |
| `ClassificationHead` | MLP predicting class logits |

**Granularity rules:**

- No single subgraph > ~25 % of total nodes; split if needed.
- No trivial single-activation subgraphs; minimum unit is one compute op.
- Split every repeated layer individually (12 transformer layers → 12 entries).
- Merge inseparable op chains (Conv + BN + ReLU with no branch points).
- Assign `Concat` / `Add` / `Resize` to the subgraph that *produces* the feature map.

### Step 4 — Choose start and end nodes

**`start_nodes`**: first node(s) receiving data from outside the block. Only
list a node if it is **not already reachable** from another listed start node —
downstream-reachable additions are redundant. When a block has multiple

From amd/Quark, MIT.

More from amd/Quark