veomni-migrate-transformers-v5

name

veomni-migrate-transformers-v5

description

Use this skill when adding or refreshing a patchgen-generated modeling file for a VeOmni model under its generated directory — GPU-only or GPU+NPU, dense or MoE, text-only / VLM / Omni-thinker+talker. Covers: creating GPU and NPU patchgen configs, using patchgen decorators (replace_class/override_method/replace_function/modify_init/add_post_import_block/drop_import_names), reusing sibling-model patches via name_map, handling MoE weight-loading (CheckpointTensorConverter + fused gate_up_proj layout), multimodal/VLM forward with Ulysses SP, excluding speech/vocoder subtrees in Omni models (talker/token2wav/DiT/BigVGAN), wiring __init__.py for the patchgen-generated classes, running codegen, and adding test cases. Trigger: 'port a model to patchgen', 'add patchgen for a model', 'transformers v5 migration', 'add NPU patchgen'. Do NOT edit files under generated/ manually — always regenerate via patchgen.

VeOmni Transformers v5 Patchgen Protocol

Purpose: add or refresh a model's patchgen-generated modeling under veomni/models/transformers/<model>/generated/. VeOmni pins transformers==5.9.0 and ships patchgen-generated modeling for every supported model; legacy v4 monkey-patches have been retired.

References (read first, load on demand):

• docs/transformers_v5/index.md — overview of what v5 migration covers
• docs/design/patchgen.md — patchgen DSL, CLI, CI drift check
• docs/transformers_v5/transformers_v5_moe_weight_loading.md — MoE fused-expert layout + runtime converter
• docs/transformers_v5/veomni_flash_attention_kernel_adapter.md — FA custom-name adapter
• docs/transformers_v5/testing_new_model.md — v5 test case SOP

Working examples (copy the structure, do not edit generated/):

Examples grouped by complexity / capability — pick the closest one and adapt:

• Text LLM (dense) — veomni/models/transformers/qwen3/, veomni/models/transformers/llama/, veomni/models/transformers/qwen2/, veomni/models/transformers/seed_oss/
  • __init__.py — registers a patchgen-generated <Model>ForCausalLM / <Model>Model / <Model>ForSequenceClassification via MODELING_REGISTRY.
  • <m>_gpu_patch_gen_config.py — Liger + SP + fused-CE patches. Llama is the minimal reference (5 OpSlot patches: RMSNorm, MLP, RoPE, ForCausalLM, ForSequenceClassification — no SP or MoE specifics).
• Text LLM with NPU patchgen — veomni/models/transformers/seed_oss/
  • __init__.py — branches on IS_NPU_AVAILABLE between patched_modeling_seed_oss_{gpu,npu}.
  • Sibling configs produce separate generated/*_{gpu,npu}.py outputs.
• MoE — veomni/models/transformers/qwen3_moe/
  • __init__.py — attaches _create_checkpoint_tensor_converter as a staticmethod on every patchgen-generated class.
  • qwen3_moe_gpu_patch_gen_config.py — replaces Qwen3MoeExperts with the fused-MoE layout and overrides get_parallel_plan.
  • checkpoint_tensor_converter.py — HF per-expert → fused runtime converter.
  • parallel_plan.py — single get_parallel_plan() sharding the fused gate_up_proj.
• MoE + NPU patchgen — veomni/models/transformers/deepseek_v3/
  • Sibling deepseek_v3_{gpu,npu}_patch_gen_config.py; both generated files committed.
  • Runtime kernel choice (deterministic Triton RoPE + batch-invariant RMSNorm) is wired in __init__.py via apply_veomni_deepseek_v3_device_patch(gen_module) for actor/rollout numerical parity. No Liger kernels in the generated file itself.
• VLM (non-MoE) + GPU+NPU patchgen — veomni/models/transformers/qwen3_vl/
  • __init__.py — registers the patchgen-generated classes, branching on IS_NPU_AVAILABLE between patched_modeling_qwen3_vl_{gpu,npu}.
  • qwen3_vl_gpu_patch_gen_config.py — full VLM forward with Ulysses SP, async Ulysses text attention, deepstack, precomputed mrope via get_position_id_func, and a SP-aware dummy_forward.
  • qwen3_vl_npu_patch_gen_config.py — demonstrates the NPU-inherits-GPU pattern: a thin NPU config that extends gpu_config.helpers / gpu_config.post_import_blocks / gpu_config.additional_imports and only overrides RMSNorm / rotary with torch_npu.npu_rms_norm / torch_npu.npu_rotary_mul. Avoids duplicating ~1K lines of shared VLM SP/deepstack patches.
• Omni (thinker+talker subtree, non-MoE) — veomni/models/transformers/qwen2_5_omni/
  • __init__.py — imports Qwen2_5OmniForConditionalGeneration / Qwen2_5OmniThinkerForConditionalGeneration from the patchgen-generated module and Qwen2_5OmniTalkerModel / Qwen2_5OmniTalkerForConditionalGeneration directly from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni (talker classes are excluded from the generated file but the registry still needs to return them when architecture mentions Talker...). MODEL_CONFIG_REGISTRY applies the tie_word_embeddings=False config patch.
  • qwen2_5_omni_gpu_patch_gen_config.py — the canonical non-MoE Omni template: excludes talker + token2wav + DiT + BigVGAN subtrees, overrides _init_weights to drop excluded UpSample1d/DownSample1d branches, overrides ForConditionalGeneration.__init__ to force has_talker=False and pin _no_split_modules=[DecoderLayer, VisionBlock, AudioEncoderLayer] (use a list[str] to match the upstream HF convention — modeling_utils.py converts it to a set internally, so either works at runtime, but staying with list[str] keeps the patched class isomorphic with the upstream base class attr), registers a load-state-dict pre-hook to strip talker.*/token2wav.* keys, overrides enable_talker/generate to raise NotImplementedError, and forwards ForConditionalGeneration.forward to thinker only — minus all MoE/EP machinery (no replace_class("…Experts"), no parallel_plan.py, no checkpoint_tensor_converter.py). Thinker uses Qwen2_5OmniThinkerCausalLMOutputWithLogProbs from veomni.utils.model_outputs to carry log_probs/entropy as constructor fields (same FSDP2 unshard-hook rationale as qwen3_omni_moe). Audio encoder uses 1D convs (conv1/conv2) — pull dummy-forward dtype from self.conv1.weight.dtype, not self.conv2d1 (that's qwen3_omni_moe-specific).
  • No parallel_plan.py / no checkpoint_tensor_converter.py — qwen2.5-Omni's thinker text model is dense (Qwen2-class MLP, not MoE), so neither EP nor fused-expert weight conversion applies. If you start from the qwen3_omni_moe template and forget to delete these, you'll get import errors from dangling references.
• VLM + MoE + GPU+NPU patchgen — veomni/models/transformers/qwen3_vl_moe/
  • __init__.py — registers three classes (Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeModel, Qwen3VLMoeTextModel) and attaches _create_checkpoint_tensor_converter as a staticmethod on each (the inner text submodel is also loadable standalone and must carry the converter).
  • qwen3_vl_moe_gpu_patch_gen_config.py — minimal config that imports most VLM SP / deepstack / async-Ulysses / dummy_forward patches from qwen3_vl via name_map={"Qwen3VL": "Qwen3VLMoe"}, and only writes MoE-specific deltas: replace_class("Qwen3VLMoeExperts") with fused layout, override_method("Qwen3VLMoeModel.__init__") to propagate _moe_implementation into config.text_config, a hand-cloned Qwen3VLMoeModel.forward (see below), Qwen3VLMoeForConditionalGeneration.forward with fused loss + aux_loss, and get_parallel_plan. This is the canonical template for any new VLM+MoE migration. Exception — do NOT reuse Model.forward via name_map: Qwen3VLMoeModelOutputWithPast carries an extra router_logits field absent from the dense Qwen3VLModelOutputWithPast; rewriting class names at the AST level keeps the dense constructor's argument list, silently dropping router_logits and collapsing MoE routing. Clone the forward body and hand-author the return.
  • checkpoint_tensor_converter.py — HF ships fused expert tensors under the same key names as VeOmni but in transposed layout ([E, H, 2*I] vs [E, 2*I, H]). Uses dim-1 shape dispatch to recognize HF vs VeOmni layout, passes VeOmni-native tensors through untouched, and hard-errors on unrecognized shapes — see Phase 3 "round-trip safety".
• Text + linear attention (qwen3_5) / VLM + MoE (qwen3_5_moe) — veomni/models/transformers/qwen3_5/, qwen3_5_moe/
  • qwen3_5_moe_gpu_patch_gen_config.py — demonstrates config.drop_import_names(...), config.add_post_import_block(...), cross-config reuse via from ...qwen3_5.qwen3_5_gpu_patch_gen_config import <fn>, and name_map={"Qwen3_5": "Qwen3_5Moe"} on override_method to share patches between sibling configs.
• MLA + MoE (GLM) — veomni/models/transformers/glm_moe_dsa/
  • Sibling glm_moe_dsa_{gpu,npu}_patch_gen_config.py produces separate generated/*_{gpu,npu}.py outputs.

---

Phase 0: Environment + Reference Setup

0.1 Verify transformers venv

Patchgen runs against transformers==5.9.0. Before touching code:

source .venv/bin/activate
python -c "import transformers; print(transformers.__version__)"

If not 5.9.0, re-sync the default env:

uv sync --frozen --extra gpu --group dev
source .venv/bin/activate

0.2 (Strongly recommended) Drop HF reference source into .agents_workspace/

.agents_workspace/ is gitignored. Keeping the upstream HF source next to your patchgen config is the single biggest accelerator for catching subtle signature/contract drift while iterating.

mkdir -p .agents_workspace/hf_reference/<m>/v5_8_1

curl -sL -o .agents_workspace/hf_reference/<m>/v5_8_1/modeling_<m>.py \
  "https://github.com/huggingface/transformers/raw/v5.9.0/src/transformers/models/<m>/modeling_<m>.py"

For VLMs also grab processing_<m>.py / image_processing_<m>.py / configuration_<m>.py if you expect processor-side or config-shape work.

If you are refreshing an existing patchgen-generated file across a transformers minor bump (e.g. the current pin 5.9.0 → 5.9.0), pull both versions side-by-side and diff to spot contract drift — substitute the <old_ver> / <new_ver> tags with the actual versions you are migrating between:

mkdir -p .agents_workspace/hf_reference/<m>/{old,new}
curl -sL -o .agents_workspace/hf_reference/<m>/old/modeling_<m>.py \
  "https://github.com/huggingface/transformers/raw/<old_ver>/src/transformers/models/<m>/modeling_<m>.py"
curl -sL -o .agents_workspace/hf_reference/<m>/new/modeling_<m>.py \
  "https://github.com/huggingface/transformers/raw/<new_ver>/src/transformers/models/<m>/modeling_<m>.py"
diff -u .agents_workspace/hf_reference/<m>/{old,new}/modeling_<m>.py | less

Things to watch for in upstream contracts:

• @can_return_tuple, @capture_outputs, @merge_with_config_defaults, @auto_docstring decorators → affect behavior of your override_method. When you override_method on a @auto_docstring-decorated method, every parameter you declare in the new signature must also appear in the patched docstring's Args: block — otherwise auto_docstring will emit warnings at import time about "undocumented parameter". For Omni-style overrides that add params like audio_feature_lengths, feature_lens, aftercnn_lens, rope_deltas, image_grid_thw, video_grid_thw, etc., copy the upstream docstring and append minimal one-line entries for every new param.
• Helper-method signatures (e.g. get_placeholder_mask takes inputs_embeds
  • image_features / video_features).
• Return-shape conventions: e.g. get_{image,video}_features.pooler_output is a tuple[per-image tensor] after torch.split, not a flat tensor.
• Packed position-ids contract ([4, bs, seq-len] with prepended text_position_ids).
• RoPE shape collapse — VLMs use apply_interleaved_mrope (and similar helpers) that collapse the leading 3-axis of mrope before layers see cos/sin, so the shape is (bs, seq_len, head_dim). Any SP path that gathers cos/sin across the sequence dim (async Ulysses, ring attention) must use the correct gather_dim. Grep upstream for interleaved_mrope, mrope_section, or any pre-attention RoPE reshape before writing the patch.
• attention_mask may be a dict — HF v5 routinely passes attention_mask={"full_attention": <tensor>, ...} keyed by attention type. Any patched forward that forwards attention_mask to compute_3d_position_ids / get_rope_index / other tensor-expecting helpers must defensively unwrap attention_mask.get("full_attention", None) when it's a dict.

Keep this directory around through commit; delete it after the PR merges (it's already gitignored so it won't leak into the repo).

---

Before You Start: Create Todos

Use TodoWrite to track phases. Suggested plan:

Phase 0: Verify venv + drop HF reference files       -> in_progress
Phase 1: Scope & audit upstream surface              -> pending
Phase 2: Draft <model>_gpu_patch_gen_config.py       -> pending
Phase 3: (MoE only) Add checkpoint converter         -> pending
Phase 4: Wire __init__.py to expose generated classes -> pending
Phase 5: Run patchgen + verify diff                   -> pending
Phase 6: Add test cases                               -> pending
Phase 7: Run tests (single-GPU + e2e)                 -> pending
Phase 8: Docs + /veomni-review + commit               -> pending

Drop phases that don't apply (e.g. Phase 3 for non-MoE models).

---

Phase 1: Scope & Audit

Input: model name <M> (e.g. qwen3_5, glm4_moe).

Operations:

1. Confirm model exists at veomni/models/transformers/<M>/. If not, the task is "add new model" — use /veomni-new-model instead.
2. If a patchgen-generated file already exists under veomni/models/transformers/<M>/generated/ you are refreshing an existing config (e.g. picking up upstream changes, adding NPU sibling, fixing a bug). Otherwise you are adding patchgen support to a model whose __init__.py previously imported HF classes directly. Either way, the rest of this protocol applies identically.
3. Decide backend coverage:
   • GPU only → one <m>_gpu_patch_gen_config.py + one generated/patched_modeling_<m>_gpu.py.
   • GPU + NPU → add sibling <m>_npu_patch_gen_config.py that writes generated/patched_modeling_<m>_npu.py; mirror the glm_moe_dsa or qwen3_vl layout.
4. Check model category: