将 Hugging Face 模型架构导入 MAX 时…

需要把一个 Hugging Face 模型 ID 接入 MAX 并验证输出与 Hugging Face 参考一致时。

适用场景

解决思路

1. 复制 references/template.md 到工作区，用于在工作过程中跟踪本次 port。
2. 所有命令都通过安装了 MAX 的 pixi 环境运行（pixi run python …、pixi run max serve …），并在 pixi.toml 所在的 skill 根目录下执行；不要使用裸 python 或 max 命令。例如：cd <path-to-skill>; pixi install; pixi run python scripts/inspect_hf.py <HF_MODEL_ID>; 或 pixi run test-scripts。
3. 把 skill 的 scripts/ 目录中的 helper scripts 复制或 vendor 到你的 repo；也可以用统一 dispatcher：pixi run python scripts/import_model.py inspect <HF_MODEL_ID>、scaffold <HF_MODEL_ID> --start-from llama3 --output-dir ./、list-archs --match LlamaForCausalLM、check-walls <HF_MODEL_ID>、list-keys <HF_MODEL_ID> --summary、gates <HF_MODEL_ID> --port-dir <port_dir>/、compare <HF_MODEL_ID> --slug <slug> --port 8000。
4. 确定 port layout：<port_dir> 是包含 arch.py 和 __init__.py 中 ARCHITECTURES 的 slug 文件夹（通常 <output_dir>/<slug>/）；把该路径传给 --custom-architectures 和 run_oss_gates.py --port-dir。MAX 解析 --custom-architectures 时会添加 dirname(<port_dir>) 到 sys.path 并 import basename(<port_dir>)；传父目录会导入错误模块名。
5. 编辑时遇到 Import/API 错误：复制 donor arch 到 modular/max/python/max/pipelines/architectures/<donor>/，并参考 pitfalls-config.md 的 Import and config API traps。
6. Phase 1 Guard：在写代码前检查架构是否已在 MAX 注册。检查 config.json::architectures[0]；运行 pixi run python list_native_archs.py --match <Class>。如果返回 slug，则运行 pixi run max serve --model <HF_MODEL_ID> 并停止，不需要 port。完整流程见 native-arch-check.md。
7. 读 config.json：运行 pixi run python -c "from transformers import AutoConfig; print(AutoConfig.from_pretrained('<HF_MODEL_ID>', trust_remote_code=True))"，或运行 pixi run python inspect_hf.py <HF_MODEL_ID>。该 helper 会从 Hub 获取原始 config.json、运行 native-arch 检查并打印每个 key 到 MAX API 的映射。
8. 列出 safetensors metadata（keys、shapes、dtypes，不下载权重）：运行 pixi run python list_checkpoint_keys.py <HF_MODEL_ID> --summary。每个 row 是 config.json key → pipeline_config.model.huggingface_config（或 arch.py 中 SupportedArchitecture 对应 architectures、torch_dtype）。无法通过 MyConfig.initialize() 接线的 key 就是需要在 graph 中实现的 delta。字段含义见 read-config-json.md。
9. 扫描硬阻塞：运行 pixi run python check_walls.py <HF_MODEL_ID>。Exit 0 继续；Exit 1 查看 recognize-walls.md；Exit 2 停止，直到 wall 被解决或明确排除。
10. 读 model card：打开 https://huggingface.co/<HF_MODEL_ID>，读 paper 或 blog post 的架构部分，注意作者强调的修改（QK-norm、MLA、sliding-window attention、MoE routing）；注意 'we introduce'、'unlike prior models'、'this is the first model to' 等短语标记的 delta。如果来自已知 family（Llama、Mistral、Qwen、Gemma），记下，donor 比较从最接近的已移植变体开始。如果提到 custom CUDA kernels、custom attention with no public reference、FP8/FP4-only released weights、ALiBi、recurrence 或 state-space layers，先看 recognize-walls.md。
11. 提出计划并接受否决：在写代码前写一段短文说明默认要做什么，然后等用户确认或否决。覆盖四个轴：distribution shape、quantization variants、validation depth、hardware target；给默认值让用户推回，不要问空问题。见 plan-and-veto.md。如果估计权重字节不适合一个 GPU，在选 --start-from 前读 distributed-transformer.md。
12. 与其他 MAX 架构比较以选 donor：最接近 = 相同 attention shape（dense vs. GQA vs. MLA vs. MoE）、相同 MLP shape（gated vs. non-gated、dense vs. routed）、相同 head layout（tied vs. untied、single Linear vs. multi-step）。运行 pixi run python list_native_archs.py 列出已安装 MAX 注册的架构（不要硬编码 slug 列表）。用 map-to-max.md 的启发式 HF 信号→donor slug 提示选择起点。
13. 打开所选 MAX donor arch 目录，读其顶层模型文件（通常 <slug>.py），回答哪些函数/类需要改、哪些不变。
14. 读对应 Hugging Face modeling 文件：运行 pixi run python -c "from transformers.models.<model_type> import modeling_<model_type>; print(modeling_<model_type>.__file__)"。读 __init__、attention forward、MLP forward、block class、final head，并与 MAX 等价实现逐项比较。见 read-modeling-code.md。
15. 输出 delta list：每个 HF 与 donor MAX arch 的真实差异一行（attention、MLP/MoE、block wiring、head、RoPE、masks）。每个 delta 都要在 Phase 2 实现。<=3 个结构 delta 说明 donor 选得好；很多 delta 则换更近的 donor 或计划重写整个类。不要在空或 'looks Llama-ish' 的 delta list 下进入验证。
16. Phase 2 Scaffold：运行 pixi run python scaffold.py <HF_MODEL_ID> --start-from <max_arch_slug> --output-dir <output_dir>。该命令读取 Hub config.json 的 architectures[0] 作为 arch.py::name，然后复制选定的 native MAX 架构到 <output_dir>/<slug>/，生成五个文件：arch.py、model_config.py、model.py、weight_adapters.py、<slug>.py。注意 scaffold.py 只复制文件，不实现模型。
17. 检查 scaffold 生成的 arch.py：确认 name= 匹配 architectures[0]、encoding 匹配、memory_planner= 从 donor 原样复制，包括配置的 planner 如 PagedMemoryPlanner.with_activation_reservation(...)。如果 donor 没有 planner，字段应留出 TODO 而不是默认。
18. 注意 scaffold 会留下 donor 的 docstrings 和 code comments；重写这些 docstrings 是 graph 实现的必须部分，不是可选修饰。见 honest-docstrings.md。
19. 实现 graph：按 implement-graph.md 的顺序执行。首先编辑 model_config.py，接线 Phase 1 / inspect_hf.py 得到的每个 config.json key，并设置 construct_kv_params() 的 head counts 和 head_dim 以匹配 HF。
20. 编辑 weight_adapters.py：把 checkpoint 的 safetensor keys 映射到你要使用的 MAX 模块名。先运行 list_checkpoint_keys.py；见 rename-weights.md。加载后接入 coverage audit（见 state-dict-audit.md），尤其 MoE 和 strict=False tied embeddings。
21. 编辑 <slug>.py：对 delta list 的每一行，编辑或替换 donor 类，使 MAX forward() 镜像 HF forward()。覆盖：Attention（Q/K/V、RoPE、mask、GQA、softcap 等）；MLP 或 MoE（activation、routing、shared experts 等）；Decoder block（norm order 和 residual wiring，不能与 Llama 互换）；Final norm 和 LM head（tie、logit scale、softcap）。
22. 编辑 arch.py：确认 name= 匹配 architectures[0]，default_encoding 匹配 Hub torch_dtype。保留 scaffolded memory_planner=；没有它 MAX 不会为 KV cache 预算 activation memory，并会跳过 max_batch_size 推理。如果 scaffold 留下 memory_planner TODO，必须解决：KV-cache ports 需要 memory_planner=PagedMemoryPlanner，只有自己做内存估计的架构（diffusion、embedding）才可留空。仅当 port 的 batching 与 donor 不同时才加 batching=。
23. 编辑 model.py：仅当 HF 对 backbone 的包装不同（VL、multi-modal）时才修改。
24. 边编辑边读 HF modeling_<type>.py，不要等验证失败后再读。只在 HF 和 donor 匹配处子类化 donor；在 delta list 说不同处重写整个类。
25. 实现完成标准：implement-graph.md 的 completion criteria 全部勾选，尤其每个 delta 有对应代码改变，权重加载无 orphan keys，且 scaffold-comment audit（见 honest-docstrings.md）已运行，每个匹配分类为 OK / Lie / Stale。审计是强制性的。快速 grep：pixi run rg -i -n 'qwen|llama|mistral|cohere|gemma|phi|deepseek|olmo|granite|qwen3|mixtral|single-GPU|single GPU|RMSNorm|QK-norm' <port_dir>/。实现完成消息必须明确 attest 审计，例如："docstrings rewritten to the three-sentence pattern; rg returns N hits, all legitimate lineage references"。
26. Preflight（首次 serve 前）：运行 pixi run python run_oss_gates.py <HF_MODEL_ID> --port-dir <port_dir>/，检查 Hub config 和 arch registration。
27. Phase 3 前强制 local smoke gate：max serve 冷编译需要 5–25 分钟。在 serve 前运行 serve-and-iterate.md 中的四个本地检查：import smoke、graph dry-build、adapter⇄graph key diff、weights-format preflight。run_oss_gates.py 覆盖 walls、checkpoint metadata、arch.py name/encoding，但不替代这四个。
28. Phase 3 检查是否生成连贯文本：前置条件是 graph 实现完成。不要在实现期间 serve 来 '看会发生什么'——先修 config、adapters、graph。
29. 先 sanity-check HF 参考：在涉及 MAX 前，单独用 HF 在 model card 的 intended prompt template 上运行。如果 HF 本身产生 gibberish，你的 oracle 坏了，不要继续。
30. 然后 serve：运行 pixi run max serve --model-path <HF_MODEL_ID> --custom-architectures <port_dir>，用 model card 的 intended template 探测（不要只用 'The capital of France is'，这对 PrefixLMs 和 instruction-tuned models 是错误 prompt）。三种结果：server 加载时 crash → 修 config/adapters；server 启动但返回 garbage → 进入 divergence hunt；server 返回 plausible text → 用 max_tokens=64+ 再运行再庆祝。见 serve-and-iterate.md。
31. Parity/coherence failure 时：当 server 启动但输出错误（gibberish、wrong greedy token at index K、high logit cosine with wrong argmax、coherent for N tokens then diverges），停止 scalar-tap loop，加载 debug-model 技能。debug-model 强制：1) 同 prompt + checkpoint 做 HF sanity-check；2) per-layer HF vs MAX tensor-dump comparators（不是 ops.print eyeballing）；3) parallel investigation agents，重新编译前做 numerical verification；4) 当 dumps 匹配但生成文本 diverges 时做 serve-vs-pipeline bisect。用 import-model 做 bring-up scaffolding 和 gates；用 debug-model 做 divergence hunt。
32. Quick logit probe（仅前 5 分钟）：在构建 dumpers 前运行 pixi run python compare_layers.py <HF_MODEL_ID> --slug <your_slug> --port 8000 --prompt "The capital of France is"。需要 pixi run max serve 带 --custom-architectures <port_dir> 在同一端口。该脚本打印 HF-only layer stats 并比较 prompt 的 top-1 logprob。见 layer-by-layer-debugging.md。如果 logprobs diverge 或输出 garbage，切换到 debug-model；不要只用 manual ops.output() taps 迭代。
33. 检查 against Hugging Face：端到端运行模型带预训练权重，然后运行 HF 同 prompt greedy sampling。MAX 侧使用匹配 weight encoding 的 dtype（大多模型发布 bfloat16）。输出应相同或几乎相同；小 BF16/FP16 舍入可能导致超过十几个 token 后 diverge。如果 divergence hunt 通过后 first tokens 仍持续 diverge，通常是 tokenizer/chat-template mismatch、dtype mismatch with released weights、或 nonzero MAX sampling。当匹配文本出来，port done for greedy text。真正 'done' 取决于 planning 时选的 validation depth——从 tier 1 (smoke) 到 6 (logit parity)。见 validation-tiers.md。
34. Common pitfalls：用 pitfalls.md 作为索引——找到 symptom 后加载一个 category 文件（config、weights、graph、serving）——docstring audit 看 honest-docstrings.md。两大坑：Scaffold ≠ port；Sed-rename 留下 donor docstrings intact。
35. Tests and CI：为 ported model 加 pytest 测试时，最小化每个文件的 MAX graph compilations。通过 module-scoped fixture 编译一次，并在 @pytest.mark.parametrize 用例中复用。对于必须编译不同 graph 的文件，用 Bazel shard_count 并行化，而不是拆分文件。见 tests-and-ci.md。

适用边界 · 注意事项

• 架构已在 MAX 注册时（list_native_archs.py --match <Class> 返回 slug），不要移植，直接运行 max serve --model <HF_MODEL_ID>。
• Hugging Face 参考实现本身在 intended prompt template 上输出乱码时，不要继续移植；先修复 oracle。
• server 启动但输出错误时，不要用 scalar-tap 或手动 ops.output() 迭代，应切换到 debug-model 技能。