import-model

What this skill does

# Import a model into MAX

**Input:** a Hugging Face model ID (`$ARGUMENTS`).

Copy [references/template.md](references/template.md) to track this port as
you work through the phases.

Porting a model to MAX means writing a MAX graph that performs the same
computation as the model's `modeling_<type>.py` in Hugging Face `transformers`,
then loading the released weights into that graph and verifying the outputs
match.

The workflow has three phases: **decide & plan**, **implement**, **verify**.
Phase 1 is reading and planning. Phase 2 is the port: implement every divergent
sublayer in the graph. Phase 3 is verification, only after implementation is
complete. Guards (preconditions that stop the line) gate the transitions
between activities; they are not steps of their own.

**Anti-pattern:** running `scaffold.py`, tweaking `arch.py`, and serving
while `<slug>.py` still implements the donor (`llama3`, `qwen3`, …). That is
not a bring-up — logit verification will fail because the wrong architecture
is running. Do not run verification scripts until
[implement-graph.md](references/implement-graph.md) completion criteria pass.

Each phase links to references with the details. Read the reference for the
activity you're on, not all of them upfront.

**Environment:** run **every** command through the pixi env that has MAX
installed (`pixi run python …`, `pixi run max serve …`), from the skill
root where `pixi.toml` lives (do not use bare `python` or `max` on the
shell PATH):

```bash
cd <path-to-skill>
pixi install
pixi run python scripts/inspect_hf.py <HF_MODEL_ID>
# Or: pixi run test-scripts   # smoke-test all scripts (no GPU)
```

Helper scripts live in this skill's `scripts/` directory (copy or vendor
them into your repo). All helpers are also reachable through a unified
dispatcher with the same argument names and exit codes:

```bash
pixi run python scripts/import_model.py inspect <HF_MODEL_ID>
pixi run python scripts/import_model.py scaffold <HF_MODEL_ID> --start-from llama3 --output-dir ./
pixi run python scripts/import_model.py list-archs --match LlamaForCausalLM
pixi run python scripts/import_model.py check-walls <HF_MODEL_ID>
pixi run python scripts/import_model.py list-keys <HF_MODEL_ID> --summary
pixi run python scripts/import_model.py gates <HF_MODEL_ID> --port-dir <port_dir>/
pixi run python scripts/import_model.py compare <HF_MODEL_ID> --slug <slug> --port 8000
```

Port layout:

- **`<port_dir>`** — slug folder containing `arch.py` and `ARCHITECTURES` in
  `__init__.py` (usually `<output_dir>/<slug>/`). Pass this path to both
  `--custom-architectures` and `run_oss_gates.py --port-dir`.

MAX resolves `--custom-architectures <port_dir>` by adding `dirname(<port_dir>)`
to `sys.path` and importing `basename(<port_dir>)` as the module. Passing the
parent directory imports the wrong module name (e.g. `custom-arch` instead of
your slug).

Import/API errors while editing: copy the donor arch under
`modular/max/python/max/pipelines/architectures/<donor>/`; see
[pitfalls-config.md § Import and config API traps](references/pitfalls-config.md#import-and-config-api-traps).

---

## Phase 1 — Decide & plan

> **Guard: is the architecture already registered in MAX?**
> Before writing any code, check whether MAX already registers the architecture
> class in your model's `config.json::architectures[0]`. If
> `pixi run python list_native_archs.py --match <Class>` returns a slug, run
> `pixi run max serve --model <HF_MODEL_ID>`
> and stop — no port needed. Full procedure:
> [native-arch-check.md](references/native-arch-check.md).

### Read `config.json`

Pull the config and read every field:

```bash
pixi run python -c "from transformers import AutoConfig; \
  print(AutoConfig.from_pretrained('<HF_MODEL_ID>', trust_remote_code=True))"
```

Or use the helper, which fetches raw `config.json` from the Hub, runs the
native-arch check, and prints every key mapped to the MAX API:

```bash
pixi run python inspect_hf.py <HF_MODEL_ID>
```

Then list safetensors metadata (keys, shapes, dtypes — no weight download):

```bash
pixi run python list_checkpoint_keys.py <HF_MODEL_ID> --summary
```

Each row is one `config.json` key → `pipeline_config.model.huggingface_config`
(or `SupportedArchitecture` in `arch.py` for `architectures`, `torch_dtype`).
Keys you cannot wire through `MyConfig.initialize()` are the deltas you
implement in the graph. Field meanings and common deltas:
[read-config-json.md](references/read-config-json.md).

Scan for hard blockers before you commit to a port:

```bash
pixi run python check_walls.py <HF_MODEL_ID>
```

Exit 0 → continue. Exit 1 → review
[recognize-walls.md](references/recognize-walls.md). Exit 2 → stop until the
wall is resolved or scoped out.

### Read the model card

Open `https://huggingface.co/<HF_MODEL_ID>` and read the model card for:

- **The paper or blog post.** Skim its architecture section — authors call out
  the *interesting* modifications (QK-norm, MLA, sliding-window attention,
  MoE routing) because those are what they want credit for.
- **"Tricks" mentioned in the card.** Phrases like "we introduce", "unlike
  prior models", "this is the first model to" mark deltas that will bite you
  during implementation if you miss them now.

If the card says the model is from a known family (Llama, Mistral, Qwen,
Gemma), note that; the donor-comparison activity below will start from the
closest already-ported variant of that family.

If the card mentions custom CUDA kernels, custom attention with no public
reference, FP8/FP4-only released weights, ALiBi, recurrence or state-space
layers; see [recognize-walls.md](references/recognize-walls.md) before going
further. Some models can't be ported with the public MAX surface alone.

### Propose a plan; accept a veto

Before any code, write a short paragraph stating what you'd do by default,
then wait for the user to confirm or veto. Cover four axes (distribution
shape, quantization variants, validation depth, hardware target) — all
derived from what you've already read. Don't ask blank questions; state a
default and let them push back.

Full guidance and an example paragraph:
[plan-and-veto.md](references/plan-and-veto.md).

If estimated weight bytes do not fit one GPU, read
[distributed-transformer.md](references/distributed-transformer.md) before
choosing `--start-from` — distribution shape matters more than attention
family alone.

### Compare with other MAX architectures

You're picking the closest already-ported MAX architecture to copy from.
"Closest" means: same attention shape (dense vs. GQA vs. MLA vs. MoE), same
MLP shape (gated vs. non-gated, dense vs. routed), same head layout (tied vs.
untied, single Linear vs. multi-step).

List what your installed MAX registers (do not hard-code a slug list):

```bash
pixi run python list_native_archs.py
```

Heuristic HF-signal → donor slug hints are in
[map-to-max.md](references/map-to-max.md). Quick version:

| Your model                                          | Start from               |
|-----------------------------------------------------|--------------------------|
| Llama 3-ish (GQA, RoPE, SwiGLU MLP)                 | `llama3`                 |
| Gemma-ish (RMSNorm scale, logit softcap, dual norm) | `gemma2` or `gemma3`     |
| Qwen-ish (GQA, RoPE, may have QK-norm)              | `qwen2` / `qwen3`        |
| Mistral-ish (sliding window)                        | `mistral`                |
| Phi-ish (partial RoPE)                              | `phi3`                   |
| MoE (sparse experts, top-k routing)                 | `mixtral` or `qwen3_moe` |
| MLA (latent KV)                                     | `deepseekV3`             |

Open the chosen MAX arch's directory and read its top-level model file
(usually `<slug>.py`). You're answering: which functions/classes need to
change vs. stay the same when I port my model?

Now read the corresponding Hugging Face modeling file:

```bash
pixi run python -c "from transformers.mode