Veomni New Model

Phase 2: Create Model Patch

1. Create the model directory: veomni/models/transformers/<model_name>/

2. Required files:

   • __init__.py — model registration
   • modeling_<model_name>_patch.py — monkey-patches for HF model (flash attention, sequence parallel, etc.)
   • parallel_plan.py — FSDP/FSDP2 sharding plan
   • generated/ — auto-generated files from patchgen (do NOT edit manually)

3. Patch patterns — follow existing models:

   • Replace attention with flash attention (via veomni/ops/flash_attn/)
   • Add sequence parallel support (via veomni/distributed/sequence_parallel/)
   • Register model in veomni/models/auto.py

4. Run patchgen if applicable: make patchgen

Phase 3: Define Parallel Plan

1. Create parallel_plan.py in the model directory.

2. Define FSDP/FSDP2 sharding strategy:

   • Which layers to wrap (typically transformer blocks)
   • Activation checkpointing granularity
   • Parameter dtype policies

3. If the model is MoE, define expert parallelism plan in addition to FSDP.

4. Reference existing parallel plans for guidance (e.g., veomni/models/transformers/qwen3/parallel_plan.py).

Phase 4: 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 5: 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

Phase 6: Test

1. Create toy config: Add tests/toy_config/<model_name>_toy/config.json with minimal parameters for fast testing.

2. Unit tests: Add tests in tests/models/ to verify:

   • Model loads correctly via veomni.models.auto
   • Forward pass produces correct output shape
   • Model patch applies without errors

3. E2e tests (if feasible): Test a short training run using the toy config.

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.
• Generated files: Never edit files in generated/ directories — they are overwritten by patchgen. Edit the patch spec or the modeling_*_patch.py instead.
• Tokenizer compatibility: Some models require specific tokenizer versions or custom chat templates — verify in veomni/data/chat_template.py.
• Transformers version: New code must target v5. Use is_transformers_version_greater_or_equal_to() for v4 compat guards. v4-era APIs (e.g., AutoModelForVision2Seq) no longer exist in v5 — use v5 equivalents.