Add heterogeneous config support (per-layer configuration)

[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 (and the follow-up modeling PR #45332), heterogeneous support requires just a few lines.

This PR contains the configuration layer only. The modeling, cache, and masking changes that consume this config are in #45332.

How it works

Configuration (per_layer_config)

A new per_layer_config parameter 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"]},
    },
)

Under the hood, apply_heterogeneous_config validates the overrides, computes fallback values, and stores a HeterogeneitySpec on the config. Serialization keeps per_layer_config as a sparse mapping of layer indices to changed attributes only. At runtime, config.per_layer_config[i] returns a full layer config built from the current global config plus that layer's overrides. When a per-layer override matches the global config value, it is omitted from the serialized per_layer_config so the sparse config only records real differences. The config supports full save_pretrained / from_pretrained round-trips (keys are zero-padded for correct JSON sort order).

Accessing a per-layer attribute on the global config raises AttributeError and points callers to config.per_layer_config[i], since the global value can be misleading. Callers can set allow_global_per_layer_attribute_access=True to allow config.<per_layer_attribute> access, but it will output a warning_once message.

Key changes

• New src/transformers/heterogeneity/ package - configuration utilities (HeterogeneitySpec, validation, runtime per-layer config view, serialization)
• configuration_utils.py - per_layer_config property/view, is_heterogeneous, serialization hooks, and __getattribute__ guard for per-layer attributes

Tests

Test suite in tests/heterogeneity/test_configuration_utils.py covering per-layer overrides and fallback, skip normalization, validation, global per-layer attribute access, and save/load round-trip.

Who can review?

@ArthurZucker
@hmellor

[eladsegal]

@askliar
Related to vllm-project/vllm#36512

[ArthurZucker]

Reviewed this first: #45332 (review) !

I think we want stuff tto be explicit and "simple".

We do have 2 choices basically:

1. per_layer_config: List[PreTrainedConfig]

2. per_layer_hidden_size: List[int]

3. Per layer configs would be quite simple, we find a way to have a minimal serialization to on serialize the actual per layer. I like that a bit less? But its the most convenient for modeling changes / inheritance

4. With this, we can have PreTrainedConfig just init the per-layer-configs based on the lists. This means nice to understand serialization, and you just parse the per_layer_<key> into Config(<key> = value).

WDYT?

[eladsegal]

Thanks @ArthurZucker! I agree with the explicit and simple direction.

The code now keeps serialization minimal, and exposes full per-layer configs:

• On disk / in to_dict(), per_layer_config is a sparse mapping of layer index -> changed attributes only, e.g. {"1": {"hidden_size": 2048}}.
• At runtime, config.per_layer_config[i] returns a full PreTrainedConfig built from the current base config plus that layer's overrides. Using the current base config matters because model initialization can mutate config state, for example _attn_implementation_internal.
• Values matching the global config are omitted, empty layer entries are dropped, and direct global access to varying per-layer attrs raises AttributeError unless explicitly opted into.

So we keep the config file simple and minimal for both input and serialization, while the instantiated config object provides convenient and intuitive per-layer access.

[ArthurZucker]

Nice! Missing a tad but of doc / md example, and appart from that if we have a mixin would make it more isolated / we avoid polluting the main configuration_utils.py with more atomic calls!

Looks nice and simple. I would document the reasoning behind: "its a power feature and we'll thrive to merge arch that have this explicitly, but this allows modularity with little to no cost at tleast in the config side!"

something like this

[ArthurZucker]

arf not a super fan of this!

[eladsegal]

Yeah, me neither, but I didn't find a cleaner path here

[ArthurZucker]

motivates me to put a HeterogenousMixin, upon init sets itself to is_heterogenous if we detect per layer config -> then you unlock the different properties etc. If not is_heterogenous then you don't see this?

WDYT I hope its not a version you already had. its easier to be looking at the code like this!

[eladsegal]

Done!
To really make these properties invisible for non-heterogeneous configs would be pretty dirty and require dynamic attribute hiding, but I think using the mixin already makes it a lot cleaner

[ArthurZucker]

IDK if there is a way to disable the function of the inherited class / enable based on has_heterogenous_specs. If possible it would simplify a lot because we could overwrite all

[ArthurZucker]

when you init, via MRO can't this be done in init (Maybe order is wrong?)

[eladsegal]

You are right, I was able to remove it so from_dict is now untouched.
I did add allow_global_per_layer_attribute_access and serialize_explicit_per_layer_config as properties on the mixin, so behavior stays consistent whether they are passed through kwargs or included in the config_dict.

[ArthurZucker]

ah its hiding if not used?

[eladsegal]

Since by default we don't want per-layer attributes to be accessed from the global config, we hide them when enumerating the keys on the config. For a regular config this is a passthrough, identical to main

[ArthurZucker]

I still have a few concerns about injection of heterogenous checks etc but its not blocking anymore IMO

[eladsegal]

Thanks for the review @ArthurZucker!
I believe I addressed all of the issues raised in your comments.

Regarding disabling the function of the inherited class / enable based on is_heterogeneous - If I understand correctly, that's exactly the case. All of the functionality provided by HeterogeneousConfigMixin is enabled only if the config is indeed heterogeneous, and otherwise there's no effect other than the visibility of the methods and properties added by the mixin.

Overwriting the inherited methods entirely to keep PreTrainedConfig clean from changes isn't possible though: the mixin is a base of PreTrainedConfig, so PreTrainedConfig's own implementations win the MRO. I think the only way it would have been possible is by deconstructing PreTrainedConfig into smaller mixins that will be after HeterogeneousConfigMixin in the MRO.
However, where the MRO does allow it, we already rely on it: the per-layer attribute guard is called through super().__getattribute__, so PreTrainedConfig.__getattribute__ is unchanged compared to main.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 29015415548:2
Result: success | Jobs: 15 | Tests: 170,757 | Failures: 1 | Duration: 4h 3m