Separate model "zoo" from infra in components/models/

RFC: Separate model "zoo" from infra in components/models/

Motivation

components/models/<name>/ is intended to be the model zoo — a thin home
for model definitions. In practice it is mostly that, but model-specific
infrastructure (parallelization plans, runtime patches,
preprocessing/collation) lives in cross-cutting framework files keyed off
model class names. The result:

• A new model touches 3–4 files outside its own directory.
• Infra files grow without bound (one is ~1.7k LoC, another ~2k LoC).
• It is hard to see, per model, "what is this model + what does it cost
  to onboard."

This RFC proposes a per-model layout that pulls the infra into the model
directory while keeping the framework's dispatcher thin.

Proposed per-model layout

components/models/<name>/
  model.py                        # nn.Module + forward only
  parallelizer.py                # TP plan, FSDP rules, AC policy, strategy class
  state_dict_adapter.py  # already there
  patches.py                    # HF monkeypatches (RoPE/attention/etc.)
  preprocessing.py          # collate_fn / processor wiring (VLMs/Omnis)
  flops.py                          # flops formula
  misc.py                           # tiny helpers that don't deserve their own file
  __init__.py

Framework side:

• components/distributed/parallelizer.py keeps only the abstract base,
  the default strategy, and the registry.
• components/distributed/optimized_tp_plans.py keeps only generic
  parallel styles (e.g. SequenceParallelAllGatherActivation,
  VocabParallelEmbedding) — per-model _parallelize_* functions move
  into each model's parallelizer.py.
• components/datasets/vlm/collate_fns.py keeps only the generic
  default_collate_fn / pad_collate_fn and the dispatch table —
  per-model collates move into each model's preprocessing.py.
• _transformers/v4_patches/ keeps only the dispatch hook —
  Nemotron-Flash RoPE etc. move into the model's patches.py.
• components/checkpoint/{checkpointing,conversion_mapping,addons}.py
  drop their hard-coded model-class / model_type branches in favor of
  small hooks the model-side adapter / patches.py implements.

Concrete evidence of the split today

1. Parallelization is scattered

components/distributed/parallelizer.py (1721 LoC) holds per-model
strategies and the registry that maps model class names to them:

# components/distributed/parallelizer.py:618
PARALLELIZATION_STRATEGIES: Dict[str, ParallelizationStrategy] = {
    "NemotronHForCausalLM":             NemotronHParallelizationStrategy(),
    "Qwen3_5ForConditionalGeneration":  Qwen3_5ParallelizationStrategy(),
    "Qwen3_5ForCausalLM":               Qwen3_5ParallelizationStrategy(),
    "WanTransformer3DModel":            WanParallelizationStrategy(),
    "HunyuanVideo15Transformer3DModel": HunyuanParallelizationStrategy(),
}

components/distributed/optimized_tp_plans.py (648 LoC) defines per-model
TP plans inline:

_parallelize_gemma3, _parallelize_gemma4, _parallelize_baichuan,
_parallelize_llama, _parallelize_ministral3, _parallelize_mistral3_vlm,
_parallelize_qwen, _parallelize_qwen_classification, _parallelize_phi,
_parallelize_phi3, _parallelize_qwen3_5_vlm,
get_llama_nemotron_super_tp_plan, get_decilm_nemotron_tp_plan, ...

The Qwen3.5 strategy already reaches back into the model dir to apply
patches, demonstrating the awkward layering:

# components/distributed/parallelizer.py:422
from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model

This would be from .patches import patch_hf_model if the strategy lived
in components/models/qwen3_5_moe/parallelizer.py.

2. Patches sit in framework code

_transformers/v4_patches/rotary.py carries Nemotron-Flash–specific
RoPE init (_compute_flash_inv_freq, fix_rotary_embeddings, gated by
_is_nemotron_flash_config). It belongs in
components/models/nemotron_v3/patches.py. The framework only needs a
generic "ask the model to register patches" hook.

3. Preprocessing leaks into the dataset layer

components/datasets/vlm/collate_fns.py is 1976 LoC dominated by
per-model collates:

# components/datasets/vlm/collate_fns.py:1967
COLLATE_FNS = {
    "KimiVLProcessor":         kimi_vl_collate_fn,         # L702
    "NemotronParseProcessor":  nemotron_parse_collate_fn,  # L1019
    "NemotronOmniProcessor":   nemotron_omni_collate_fn,   # L1560
    "LlavaOneVisionProcessor": llava_onevision_collate_fn, # L1915
}

Each entry is 50–400 lines of model-specific tensor packing.
Co-locating these in components/models/<name>/preprocessing.py cuts the
dataset layer down to dispatch + generic helpers.

4. Checkpoint loading carries per-model carve-outs

The base load path in components/checkpoint/ repeatedly branches on the
model class / model_type to work around model-specific quirks:

# components/checkpoint/checkpointing.py:629
is_nemotron_v2 = model_class == "NemotronHForCausalLM" and not getattr(model.config, "n_routed_experts", None)
is_nemotron_v3_hf = (
    model_class == "NemotronHForCausalLM"
    and getattr(model.config, "n_routed_experts", None)
    and hasattr(model, "backbone")
)
skip_initialize_weights = (
    model_class in ["Gemma3ForConditionalGeneration", "Gemma3ForCausalLM"]
    or is_nemotron_v2
    or is_nemotron_v3_hf
    or has_padding_idx
    or owns_weight_load
)

# components/checkpoint/checkpointing.py:705
if model_type == "nemotron_h" and hasattr(model, "backbone"):
    key_mapping = None  # skip backbone.* -> model.* conversion

The same pattern shows up in conversion_mapping.py as hard-coded sets
and dicts:

# components/checkpoint/conversion_mapping.py:60
MODELS_REQUIRING_TENSOR_MERGING = {
    "mixtral", "minimax", "phimoe", "qwen2_moe", "qwen3_moe",
    "deepseek_v2", "deepseek_v3", "jamba", "olmoe", "lfm2_moe",
    "dots1", "ernie4_5_moe", "glm4_moe", "glm4v_moe",
    "longcat_flash", "qwen3_omni_moe", "qwen3_next",
    "qwen3_vl_moe", "hunyuan_v1_moe", "flex_olmo",
}
# components/checkpoint/conversion_mapping.py:168
_VLM_KEY_MAPPINGS = {"gemma3": { ... }}

And in addons.py the PEFT path reaches into a specific model's
adapter to special-case it:

# components/checkpoint/addons.py:275
def _is_qwen3_moe(model):
    from nemo_automodel.components.models.qwen3_moe.state_dict_adapter import Qwen3MoeStateDictAdapter
    return isinstance(getattr(model, "state_dict_adapter", None), Qwen3MoeStateDictAdapter)

# components/checkpoint/addons.py:293
if _is_qwen3_moe(model):
    return ["mlp.experts.gate_up_proj", "mlp.experts.down_proj"]

Each of these is a tax the framework pays for one model. Under the
proposed split they belong on the model: a small per-model hook (e.g.
should_skip_init_weights(model) -> bool,
extra_target_parameters(model) -> list[str],
requires_tensor_merging advertised by the adapter) that the loader
queries, with the implementation in
components/models/<name>/{adapter,patches}.py. The loader stays
generic; the carve-outs travel with the model that needs them.

5. Today's per-model footprint is already compact

Most existing model dirs are small — even the biggest are dominated by
the actual nn.Module, not boilerplate:

model	model.py	adapter	layers	rope	total
llama	531	76	—	249	856
qwen3_moe	322	264	145	—	731
deepseek_v3	365	431	216	238	1250
mistral3_vlm	176	229	—	—	405
qwen3_omni_moe	464	104	—	—	568
qwen3_vl_moe	623	147	—	—	770

Pulling the per-model parallelizer + collate into each dir adds
~100–300 LoC per model but removes 10× as much from the shared infra
files (each branch in optimized_tp_plans.py / collate_fns.py /
parallelizer.py is 30–400 LoC), and gives a one-glance answer to
"what does this model need to run."

Migration sketch

1. Add parallelizer.py, patches.py, preprocessing.py slots to the
   model-onboarding skill (SKILL.md) — only required when non-default.
2. Move _parallelize_<model> from optimized_tp_plans.py and
   <Model>ParallelizationStrategy from distributed/parallelizer.py
   into each model's parallelizer.py. Keep register_parallel_strategy
   as the public API; the model's __init__.py registers on import.
3. Move per-processor collates from datasets/vlm/collate_fns.py into
   each model's preprocessing.py; replace COLLATE_FNS with a
   registry populated at import time (same pattern as
   MODEL_ARCH_MAPPING).
4. Move v4 / kernel patches keyed on a specific model into that model's
   patches.py.
5. Replace the model-class / model_type branches in
   components/checkpoint/{checkpointing,conversion_mapping,addons}.py
   with hooks (e.g. should_skip_init_weights,
   requires_tensor_merging, extra_target_parameters) that each
   model's adapter / patches.py implements; delete the hard-coded
   sets and dicts (MODELS_REQUIRING_TENSOR_MERGING,
   _VLM_KEY_MAPPINGS, the Gemma3 / NemotronH / Qwen3-MoE branches).
6. Land one model per PR (suggested order: nemotron_v3, qwen3_5_moe,
   mistral3_vlm, kimivl, llava_onevision) so the framework files
   shrink incrementally and each move is reviewable.

Where should the model dir live?

The structural split (above) is the main proposal. A natural follow-up
question is where the per-model dirs should sit. A few options:

1. nemo_automodel/_transformers/models/<name>/ (proposed)
   Co-locates with the HF bridge that already lives in _transformers/
   (auto_model.py, registry.py, capabilities.py, infrastructure.py,
   v4_patches/). The registry's import paths become local references
   (_transformers.registry → _transformers.models.<name>.model)
   instead of crossing packages. Symmetric with a future
   _diffusers/models/ for diffusion.
   Cost: _transformers/ is leading-underscore, which conventionally
   reads as "internal" — but model classes are arguably the most public
   surface users subclass.
2. nemo_automodel/models/<name>/
   Top-level, no underscore. Treats models as the primary artifact of
   this repo (which they are). Loses the "HF-bridge cluster" co-location.
3. nemo_automodel/zoo/<name>/
   Leans into the "zoo" framing. Same tradeoffs as (2), with a more
   evocative name.
4. Status quo: nemo_automodel/components/models/<name>/
   Keeps models as a peer of datasets/, distributed/, checkpoint/
   under components/. Fine, but does not reflect that models are a
   different kind of thing — they're the artifact, not a building block.

Either way the move is a one-shot rename + sweep of MODEL_ARCH_MAPPING
import paths in _transformers/registry.py and any string references
in tests / docs. Suggest doing it as the first PR in the series so
subsequent moves land on the new layout.

Out of scope

• Changing the state_dict_adapter.py contract or the registry contract
  — those are already in the right place.

Open questions

• Should parallelizer.py register itself on import, or should
  _transformers/registry.py learn to call into it lazily? Lazy is
  cleaner but loses the "import the model dir to use the model"
  property.
• For models that share a strategy (Qwen3_5 CausalLM + ConditionalGen
  both use Qwen3_5ParallelizationStrategy), where does the shared
  class live? Probably the more general dir (qwen3_5_moe) with the
  other importing it.