Add heterogeneous model support (per-layer config and modeling)

[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

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 contains the modeling, cache, and masking layers. It builds on the configuration layer (per_layer_config) introduced in the configuration PR #45333.

How it works

Configuration (per_layer_config) - recap

Introduced in #45333, recapped here since everything below consumes it. per_layer_config on PreTrainedConfig maps layer indices to attribute overrides:

from transformers import LlamaConfig

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

config.per_layer_config[i] returns the resolved configuration for layer i, combining the global values with that layer's overrides. See #45333 for validation, serialization, and global-attribute-access behavior.

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.

At model init time, the framework monkey-patches layer_cls.__init__ to:

1. Resolve the layer index (from the constructor arguments, or from the construction loop's variable via the call stack)
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. When sliding_window / attention_chunk_size vary per layer, patch the layer forward to select the attention mask matching its own value

The patching uses a ContextVar to pass per-model context to the layer __init__ wrapper. Cleanup happens automatically after model init (via __init_subclass__ wrapping).

The spec names the layer class and layer-index argument; sub-module skipping additionally requires skip descriptors:

import torch

from transformers.integrations.heterogeneity import (
    HeterogeneousModelingSpec,
    ReturnEntry,
    SkipDescriptor,
    get_skip_replacement,
)
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_variable_name="layer_idx",
    skip_descriptors={
        "attention": SkipDescriptor(
            replacements={
                "input_layernorm": get_skip_replacement(
                    LlamaRMSNorm, ReturnEntry(arg_name="hidden_states", transform=identity)
                ),
                "self_attn": get_skip_replacement(
                    LlamaAttention,
                    [ReturnEntry(arg_name="hidden_states", transform=torch.zeros_like), None],
                ),
            },
            replaces_kv_cache_updater=True,
        ),
        "mlp": SkipDescriptor(
            replacements={
                "post_attention_layernorm": get_skip_replacement(
                    LlamaRMSNorm, ReturnEntry(arg_name="hidden_states", transform=identity)
                ),
                "mlp": get_skip_replacement(LlamaMLP, ReturnEntry(arg_name="x", transform=torch.zeros_like)),
            },
            replaces_kv_cache_updater=False,
        ),
    },
)

The skip type names ("attention", "mlp") are what configurations reference in their skip lists. Replacement keys are member names, optionally restricted to a member class - ("mixer", NemotronHAttention) - for layers whose member class varies. replaces_kv_cache_updater states whether the skip replaces the member that updates the layer's KV cache (see below).

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

Cache, masks, and KV state

When sliding_window or attention_chunk_size varies per layer:

• DynamicCache and StaticCache create per-layer sliding window layers with the correct window sizes
• create_sliding_window_causal_mask / create_chunked_causal_mask return an AttentionMasksByAttributeValue mapping keyed by distinct window sizes (deduplicated), and each layer's forward selects its own mask; this composes with create_masks_for_generate's layer-type-keyed output

A layer whose attention is skipped never updates its KV cache, which breaks the long-standing assumption that any layer can report cache metadata:

• Cache.get_seq_length() (no layer argument) and attention mask creation now read a representative layer that actually holds KV state
• replaces_kv_cache_updater resolves the KV-disabled layer indices onto the config at model construction, so StaticCache(config, ...) picks correct representative layers statically and compiled decoding stays a single graph
• torch.export with DynamicCache preserves skipped-layer positions in the pytree round-trip instead of compacting KV states into wrong layers

Key changes

• src/transformers/integrations/heterogeneity/ package - modeling spec and layer patching, skip replacements (SkipDescriptor, get_skip_replacement/ReturnEntry), heterogeneous mask construction, and the built-in spec registry (the package's configuration utilities come from Add heterogeneous config support (per-layer configuration) #45333)
• modeling_utils.py - auto-calls apply_heterogeneous_modeling at model init, wraps subclass __init__ for cleanup via __init_subclass__
• cache_utils.py - per-layer sliding window resolution in DynamicCache/StaticCache; representative-KV-layer selection for cache metadata
• masking_utils.py - heterogeneous dispatch for sliding/chunked mask creation; mask layer selection prefers layers holding KV state
• integrations/executorch.py - DynamicCache pytree export support preserves the layer layout of caches with skipped layers
• models/gpt_oss/modeling_gpt_oss.py and integrations/mxfp4.py - Fix homogeneity assumptions in order to support heterogeneous configurations
• Docs - a spec-author guide (heterogeneous_modeling.md)

Tests

Comprehensive test suite in tests/heterogeneity/ covering 4 architectures (Llama, GptOss, Llama4, NemotronH):
Covers structure verification, forward/generate equivalence against manually-constructed reference models, model save/load round-trips, heterogeneous cache behavior (including torch.compile single-graph decoding and torch.export round-trips), mask construction, and edge cases. (test_configuration_utils.py is part of #45333.)

python -m pytest tests/heterogeneity tests/utils/test_cache_utils.py tests/utils/test_masking_utils.py

---

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

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 28766754685:2
Result: failure | Jobs: 15 | Tests: 171,642 | Failures: 6 | Duration: 25h 51m