[eladsegal]

What does this PR do?

Adds heterogeneous model support - the ability for individual layers to differ from the global config (e.g., different intermediate_size, num_key_value_heads) and to skip sub-modules entirely (MLP, attention, etc.). This enables models where layers are not uniform, as in pruned, distilled, or NAS-derived architectures.

Examples of such models:

Model	Derived from
nvidia/Llama-3_3-Nemotron-Super-49B-v1_5	meta-llama/Llama-3.3-70B-Instruct
nvidia/Llama-3_1-Nemotron-Ultra-253B-v1	meta-llama/Llama-3.1-405B-Instruct
nvidia/gpt-oss-puzzle-88B	openai/gpt-oss-120b
nvidia/NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-BF16	nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16

These models previously required trust_remote_code=True to modify the classes of the model they are derived from. With this PR, checkpoints derived from the architectures with built-in support (see supported_models.py) work out of the box, and any other model - including custom models - can opt in with a few lines.

This PR adds the generic modeling integration, skip handling, a built-in architecture registry, layer-indexed attention masks, and cache handling for skipped cache writers. It builds on per_layer_config introduced in PR #45333, per-layer MTP configuration from PR #48264, and per-layer cache configuration from PR #48178.

How it works

per_layer_config describes how individual layers differ from the global configuration. For each layer,
config.per_layer_config[layer_idx] provides the resolved configuration, including attribute overrides and skip
entries. For example:

from transformers import LlamaConfig

config = LlamaConfig(
    ...,
    per_layer_config={
        0: {"intermediate_size": 64},
        2: {"intermediate_size": 96, "skip": ["attention"]},
    },
)

This PR makes supported model layers consume that resolved configuration during construction.

Modeling (automatic layer patching)

Architecture support is declared with a HeterogeneousModelingSpec. Built-in specs live in
supported_models.py; any other model opts in by setting _heterogeneous_modeling_spec on its PreTrainedModel base class. A spec targets one layer class. A model that builds its layers from several classes, like Jamba's attention and Mamba layers, sets a tuple with a spec for each.

At model init time, the framework monkey-patches the __init__ of each spec's layer_cls to:

1. Resolve the layer index using the spec's layer_idx_resolver
2. Pass config.per_layer_config[layer_idx] to the original __init__
3. For layers with skip entries, replace the corresponding sub-modules with parameter-free no-op replacements
4. Register a forward pre-hook that selects the layer's attention mask by its layer index

The layer-class wrapper remains installed, but it is active only when a matching model is being initialized. A ContextVar scopes the per-model state, which is restored automatically after model initialization by the wrappers installed through __init_subclass__. config.per_layer_config[i] creates each layer's config on first access and returns the same object from then on, so the layers are built with these objects and keep sharing them. Setting or deleting an attribute on the global config, for example the attention implementation, carries over to layer configs that still have the global value. What the skips leave behind (layers without attention mask consumers or cache writers) is attached to the config only after the root model's initialization succeeds, and initialization fails unless every layer is constructed from one of the specs' layer classes.

The model __init__ wrapper is installed on every PreTrainedModel subclass through __init_subclass__ because any of them, including a user's, can be the root model, and no other hook runs when the root model's __init__ returns or raises: post_init runs at the end of each nested model's __init__, such as LlamaModel inside LlamaForCausalLM, and not at all on failure.

The spec names the layer class and its layer-index resolver. Sub-module skipping also requires skip descriptors, and declarations of the layer members that consume the attention mask or write to the cache:

import torch

from transformers.integrations.heterogeneity import (
    HeterogeneousModelingSpec,
    LayerIdxFromArgument,
    ReturnEntry,
    get_skip_replacement_factory,
)
from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaMLP, LlamaRMSNorm


def identity(hidden_states):
    return hidden_states


spec = HeterogeneousModelingSpec(
    layer_cls=LlamaDecoderLayer,
    layer_idx_resolver=LayerIdxFromArgument("layer_idx"),
    skip_descriptors={
        "attention": {
            "input_layernorm": get_skip_replacement_factory(
                LlamaRMSNorm, ReturnEntry(arg_name="hidden_states", transform=identity)
            ),
            "self_attn": get_skip_replacement_factory(
                LlamaAttention, [ReturnEntry(arg_name="hidden_states", transform=torch.zeros_like), None]
            ),
        },
        "mlp": {
            "post_attention_layernorm": get_skip_replacement_factory(
                LlamaRMSNorm, ReturnEntry(arg_name="hidden_states", transform=identity)
            ),
            "mlp": get_skip_replacement_factory(LlamaMLP, ReturnEntry(arg_name="x", transform=torch.zeros_like)),
        },
    },
    attention_mask_consumers=["self_attn"],
    cache_writers={"self_attn": ["key_value"]},
)

The skip type names ("attention", "mlp") are what configurations reference in their skip lists. Layer members are named with member keys: a path, such as "self_attn" or "block.attn", or a (path, class) tuple, such as ("mixer", NemotronHAttention), which matches only if the member at that path is an instance of the class. (path, class) keys let one skip descriptor cover a member that has a different class in different layers.

attention_mask_consumers lists the members that consume the attention mask, and cache_writers maps the members that write to the cache to the caches they write ("key_value" or "linear_attention"). Both list every such member, including members that no skip replaces, with member keys. A (path, class) key can also name a member that only some layers build, like Bamba's self_attn, which only its attention layers have.

Note: a spec without skip_descriptors is sufficient for attribute heterogeneity (varying dimensions across layers).

Generic patching is opt-in. It is applied only when the configuration is heterogeneous and the model declares a HeterogeneousModelingSpec. Models without a spec, or with _heterogeneous_modeling_spec = None, are left unpatched, so architectures that consume per_layer_config directly can handle it in their own modeling code.

MTP models

When the MTP configuration is heterogeneous and the main model has a spec, MTP layers are built from their per-layer configs in mtp_per_layer_config, so per-layer attribute overrides apply to them. MTP layers don't support skip, and setting it raises an error.

Attention masks and caches

When generic patching is applied, the mask factories (causal, bidirectional, sliding-window, and chunked) return an AttentionMasksByLayerIdx mapping keyed by layer index, and each layer's pre-hook selects its own entry:

• Each mask is built from the layer's resolved config and its own cache layer
• Layers share one mask object when their mask settings (e.g., sliding_window) and cache geometry match
• With a static cache, full-attention layers share a mask on the same assumption upstream mask creation makes: all of them are at the same position
• A layer that skips all its attention mask consumers gets None instead of a mask
• This layer-indexed mapping composes with create_masks_for_generate's layer-type-keyed output

When a layer's skips replace every writer of one of its caches, its cache entry drops that cache. A layer with no cache left gets a DummyLayer, which holds nothing, so a static cache allocates no memory for it. A hybrid layer that keeps one of its two caches gets the layer type that holds only that cache, such as full_attention for a hybrid layer that keeps its key/value cache. Positions are read from the first layer with a key/value cache, so the first layer's attention can be skipped.

Key changes

• src/transformers/integrations/heterogeneity/ package - modeling specs and layer-index resolvers, generic layer patching, skip replacements (get_skip_replacement_factory / ReturnEntry), layer-indexed mask construction, cache layer types for skipped cache writers, built-in architecture registry, per-layer configs owned by the config (created once, following global changes), one spec per layer class for models with several layer classes, and configuration validation
• modeling_utils.py - applies generic heterogeneous modeling during model initialization and scopes the initialization context by wrapping subclass __init__ through __init_subclass__
• cache_utils.py - drops the caches that skips leave without writers when choosing each layer's cache type, and lets QuantizedCache build a DummyLayer for a layer without any cache. As a result, QuantizedCache also accepts models with MoE/MLP layers or KV-sharing layers, which it rejected before, because those layers get a DummyLayer too.
• modeling_layers.py - builds heterogeneous MTP layers from their per-layer configs, and rejects skip for them
• masking_utils.py - enables layer-indexed dispatch in the mask factories, and adds a layer_idx argument to the bidirectional ones
• generation/utils.py - handles layer-indexed masks nested in generation's mask dicts, and None attentions from skipped layers
• exporters/utils.py - keeps layer-indexed masks when merging the masks of decode steps
• utils/output_capturing.py - records None for skipped attention layers, so attentions[i] still belongs to layer i
• Docs - a generic modeling and spec-author guide (generic_heterogeneous_modeling.md), linked from heterogeneous_configurations.md

Tests

Comprehensive test suite in tests/heterogeneity/ covering 4 architectures (Llama, GptOss, Llama4, NemotronH) with their built-in specs, and Jamba and Bamba with custom specs, for models with several layer classes and for members that only some layers build. It covers structure verification, forward/generate equivalence against manually constructed reference models, model save/load round-trips, cached decoding (including with the first layer's attention skipped), static-cache and assisted generation with skipped attention, cache layer types for skipped cache writers, layer-indexed mask construction and sharing (including no recompiles under torch.compile as the cache grows), MTP models, built-in and custom specs, layer-index resolution, skip replacements, configuration validation, and edge cases.

---

Concrete example: Llama-3.1-8B-Instruct with 4 layers of attention skipped

Code snippet (click to expand)

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline


base_model_id = "meta-llama/Llama-3.1-8B-Instruct"

# 1. Load model with per-layer overrides (Llama support is built in - no extra setup needed)
model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    device_map="auto",
    dtype="auto",
    per_layer_config={
        19: {"skip": ["attention"]},
        20: {"skip": ["attention"]},
        22: {"skip": ["attention"]},
        23: {"skip": ["attention"]},
    },
)

# 2. Generate
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)
generation_config = GenerationConfig.from_pretrained(base_model_id)
generation_config.max_new_tokens = 256

messages = [
    {"role": "user", "content": "Hey, how are you doing today?"},
]

outputs = pipe(
    messages,
    generation_config=generation_config,
)

print(outputs[0]["generated_text"][-1]["content"])

Who can review?

@ArthurZucker
@hmellor

[eladsegal]

@askliar
Related to vllm-project/vllm#36512

[ArthurZucker]

Hey! thanks for the PR!
One thing that's super important to stay aligned with transformers philo is that is not somthing we want to ship to old models! We'll be super happy to create new models with modular that do have heterogenous configs per layer, but adding this to llama break llama. GPT OSS does NOT have this, as such it should not support it. Happy to have heterogen_gpt_oss !

This removes code paths, and the if is_heterogeneous etc.
Given this, I think it would mostly be a change in the config to have per_layer_config.

We can find a good way to serialize all args in order to reduce the bloat potentially (meaning all attributes are list, at init time this creates a list of config per-layer) this simplifies code changes in modeling_heterogenous_gpt_oss for example, that would just get per_layer_config[i]

Furthermore I don't understand the skip idea can you give a concrete example? you want to skip an entire layer, but you can just set it to nn.Identity for example? or just re-index them?

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: gpt_oss