veomni-new-model

> The hard part of a new transformers-family model — the patchgen config, > parallel plan, MoE weight conversion, `__init__.py` registration, codegen — > lives in `/veomni-patchgen-model`. This skill is the wrapper around it: it > decides *what* you are adding, then hands off, then does the config, trainer > and data work that patchgen does not cover. ## Before You Start: Create a Plan Track the phases with whatever todo/plan tool the running agent provides: ``` Phase 1: Analyze HF model -> in_progress Phase 2: Modeling (/veomni-patchgen-model) -> pending Phase 3: Write training config -> pending Phase 4: Integrate with trainer -> pending Phase 5: Test and document -> pending ``` ## Phase 1: Analyze HuggingFace Model 1. **Identify the model** on HuggingFace. Read its `config.json`, `modeling_*.py`, and any processor configs. 2. **Determine model category**: - Text-only LLM -> `veomni/models/transformers/<model_name>/` - Vision-Language -> `veomni/models/transformers/<model_name>/` + `veomni/data/multimodal/` - MoE model -> additional `veomni/distributed/moe/` integration - Diffusion model -> `veomni/models/diffusers/<model_name>/` 3. **Check existing similar models**: Find the closest existing model in `veomni/models/transformers/` and use it as a reference. E.g., if adding a new Qwen variant, reference `qwen3/` or `qwen3_vl/`. 4. **Identify required patches**: VeOmni uses a patchgen system (`veomni/patchgen/`) to generate model patches from the HuggingFace modeling. Check whether a sibling model already has a config you can extend via `name_map` — that is usually the difference between a 60-line config and a 1000-line one. 5. **Compare checkpoint keys** against the supported upstream version and any existing VeOmni model. Apply the decision rule below before resolving a mismatch. ### Checkpoint key conflicts require a user decision When upstream model, VeOmni model, or checkpoint parameter keys disagree, show the concrete old/new keys and explain the impact on weight loading, export, and optimizer/DCP resume. Ask the user how to resolve the conflict before implementing a rename, alias, or compatibility mapping. Do not silently retain an obsolete model hierarchy just to preserve checkpoint keys. If the user has already chosen a resolution in the current task, apply it without asking again. When that choice is to follow current upstream keys, keep those keys in the model and handle approved legacy-key conversion in the checkpoint layer. Verify the chosen direction with strict loading and checkpoint round-trip tests; do not hide mismatches with `strict=False`. ## Phase 2: Modeling — hand off to `/veomni-patchgen-model` 1. **Create the model directory**: `veomni/models/transformers/<model_name>/`. 2. **Switch to `/veomni-patchgen-model`.** It owns the whole modeling surface — the `<model_name>_{gpu,npu}_patch_gen_config.py` files, ExtraParallel `parallel_plan.py`, any required MoE `checkpoint_tensor_converter.py`, `__init__.py` registration, `make patchgen`, and the model-level test cases — with the working examples and the pitfalls that cost the most time. Do not re-derive it from this file. Note that `parallel_plan.py` is **not** an FSDP wrapping policy: FSDP2 wraps generically in `build_parallelize_model()`, and `ParallelPlan` (`veomni/distributed/parallel_plan.py`) only describes ExtraParallel sharding, such as expert parallelism or embedding sharding. Add a plan whenever the model uses ExtraParallel, including dense models that shard embeddings; a model without ExtraParallel does not need one. 3. **Exception — non-transformers architectures.** Diffusion models under `veomni/models/diffusers/<model_name>/`, and the `flux` / `movqgan` / `wan` directories, have no `generated/` output and no patchgen config: they patch through `device_patch.py` or direct modeling. Copy the closest existing one and skip to Phase 3. Come back here once the model loads and its registry / patch tests pass. ## Phase 3: Write Training Config 1. **Model config**: Create `configs/model_configs/<model_family>/<ModelName>.json` matching HuggingFace format. 2. **Training config**: Create YAML in the appropriate directory: - Text: `configs/text/<model_name>.yaml` - Multimodal: `configs/multimodal/<model_name>/<model_name>.yaml` - DiT: `configs/dit/<model_name>.yaml` 3. Config must include: model path, data config, optimizer settings, parallelism config, checkpoint settings. 4. **Verify against existing configs** — match the structure of similar model configs. ## Phase 4: Integrate with Trainer 1. Verify the model works with the appropriate trainer: - Text -> `TextTrainer` (`veomni/trainer/text_trainer.py`) - VLM -> `VLMTrainer` (`veomni/trainer/vlm_trainer.py`) - DiT -> `DitTrainer` (`veomni/trainer/dit_trainer.py`) 2. If the model needs custom data preprocessing: - Add transform in `veomni/data/data_transform.py` or `veomni/data/multimodal/` - Register the transform for the model 3. If the model needs custom collator logic: - Extend `veomni/data/data_collator.py` 4. **VLM only — multimodal metadata precompute**: to keep the ViT forward free of host-device CUDA syncs, derive ViT `cu_seqlens` / `max_seqlen` in the collator rather than the forward. Follow the checklist in `.agents/knowledge/multimodal_metadata.md` ("Adding the hook to a new model"): a `collate_multimodal_metadata` patchgen helper + a `get_metadata_collate_func` override, the per-modality `vit_metadata` sub-dict threaded through Model.forward → ViT.forward (with a runtime fallback), and the model added to `_MM_METADATA_WIRED_CASES` in the sync gate test. ## Phase 5: Test and Document 1. **Create toy config**: Add `tests/toy_config/<model_name>_toy/config.json` with minimal parameters for fast testing. 2. **Unit tests**: add cases to the existing enumerated tables rather than new files — `tests/models/test_model_registry.py` and `tests/models/test_models_patch.py` (`TEST_CASES`) already cover loading via `veomni.models.auto`, forward output shape, and patch application. See `.agents/knowledge/testing.md` for the full landing-spot table and for why a new file outside `tests/ops/` / `tests/data/` will not run in CI unless it is wired into the unit-test workflows. 3. **E2e tests** (if feasible): add a `pytest.param` to `tests/e2e/test_e2e_parallel.py` using the toy config, rather than a new e2e file. 4. Run `make quality` and `pytest tests/models/`. 5. **Update documentation**: - Add usage example to `docs/` (training command, config reference). - Update `.agents/knowledge/architecture.md` if the model adds a new module or trainer path. - Update supported models table in project `README.md` if applicable. ## Common Pitfalls - **Model registry**: Registration must happen at import time in `__init__.py`. If the model's `AutoConfig` type is not registered, `build_foundation_model()` will fail. - **Tokenizer compatibility**: Some models require specific tokenizer versions or custom chat templates — verify in `veomni/data/chat_template.py`. - **Skipping the handoff**: the modeling pitfalls — never editing `generated/`, MoE expert layout, `name_map` reuse, Omni subtree exclusion — are in `/veomni-patchgen-model`, not here. This file deliberately does not restate them, so a summary read of Phase 2 is not enough to write a config.