Add heterogeneous model support (per-layer config and modeling)
Author: eladsegal
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=45332)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45332)
<!-- ci-dashboard-badge:end -->

# 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](https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1_5) | [meta-llama/Llama-3.3-70B-Instruct](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct) |
| [nvidia/Llama-3_1-Nemotron-Ultra-253B-v1](https://huggingface.co/nvidia/Llama-3_1-Nemotron-Ultra-253B-v1) | [meta-llama/Llama-3.1-405B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct) |
| [nvidia/gpt-oss-puzzle-88B](https://huggingface.co/nvidia/gpt-oss-puzzle-88B) | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b/tree/main) |

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`](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/heterogeneity/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:

```python
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`](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/heterogeneity/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:

```python
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 #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.)

```bash
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

<details>
<summary>Code snippet (click to expand)</summary>

```python
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"])
```

</details>

## Who can review?
@ArthurZucker
@hmellor

Comments:
@askliar
Related to https://github.com/vllm-project/vllm/pull/36512

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

run-slow: gpt_oss

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45332)
**Latest run:** [29281532849:2](https://github.com/huggingface/transformers/actions/runs/29313979421)
**Result:** `success` | **Jobs:** 6 | **Tests:** 136,873 | **Failures:** 1 | **Duration:** 43m 43s

<!-- ci-dashboard-recap:end -->