Nested per-layer `rope_parameters` is misparsed when `layer_types` omits one of its keys

[albertvillanova]

System Info

• transformers: main @ 83d46aa2a2 (reported version 5.16.0.dev0). The relevant code in modeling_rope_utils.py is identical at the v5.16.1 tag, so the released version is affected too.
• torch: 2.8.0+cu128
• Python: 3.12.9
• Platform: Linux-5.15.0-190-generic-x86_64-with-glibc2.31

Who can help?

@zucchini-nlp

Description

A config whose rope_parameters is a nested per-layer-type dict is misparsed as a flat dict whenever layer_types does not happen to contain every key of that dict.

Models such as Gemma3 and Olmo3 populate rope_parameters with one entry per possible layer type, unconditionally:

{"sliding_attention": {...}, "full_attention": {...}}

But layer_types lists the types the model actually instantiates, derived from num_hidden_layers and the sliding-window period. A short model can legitimately end up with no full_attention layer at all, for example Gemma3 with 2 layers and _sliding_window_pattern=6, or Olmo3 with 2 layers and its (i + 1) % 4 pattern.

Both classification points test membership against that instantiated list:

• standardize_rope_params, line 784
• validate_rope, line 828

set(rope_parameters.keys()).issubset(layer_types)

{'full_attention', 'sliding_attention'}.issubset({'sliding_attention'}) is False, so the nested dict takes the "one global dict" branch.

Three consequences:

1. A spurious warning, since the layer-type keys are then read as rope parameter names.
2. standardize_rope_params injects rope_type and rope_theta at the top level of the nested dict, and save_pretrained persists them.
3. validate_rope validates the wrapper instead of the real per-layer entries, so the actual rope parameters (yarn, linear, ...) are silently never validated.

The models still run correctly, because they index config.rope_parameters[layer_type] directly.

Reproduction

from transformers import Gemma3TextConfig

Gemma3TextConfig(num_hidden_layers=2)

Unrecognized keys in `rope_parameters` for 'rope_type'='default': {'full_attention', 'sliding_attention'}

>>> Gemma3TextConfig(num_hidden_layers=2).rope_parameters
{'sliding_attention': {'rope_type': 'default', 'rope_theta': 10000.0},
 'full_attention': {'rope_type': 'default', 'rope_theta': 1000000.0},
 'rope_type': 'default',      # injected
 'rope_theta': None}          # injected

With enough layers for the pattern to yield a full_attention layer, both the warning and the injected keys disappear:

>>> Gemma3TextConfig(num_hidden_layers=6).rope_parameters
{'sliding_attention': {'rope_type': 'default', 'rope_theta': 10000.0},
 'full_attention': {'rope_type': 'default', 'rope_theta': 1000000.0}}

The injected keys survive a round trip. Loading a 2-layer Olmo3 checkpoint and saving it writes:

"rope_parameters": {
  "full_attention": {"rope_type": "yarn", "factor": 8.0, "beta_fast": 32.0, "...": "..."},
  "sliding_attention": {"rope_type": "default", "rope_theta": 500000.0},
  "rope_theta": null,
  "rope_type": "default"
}

Expected behavior

rope_parameters should be recognized as a per-layer-type dict based on whether its keys are layer-type labels, not on whether the config happens to instantiate every one of them. No warning, and no extra top-level keys written into the nested dict.

Possible fix

Compare against the layer-type vocabulary (ALLOWED_LAYER_TYPES in configuration_utils.py) instead of self.layer_types, in both places. standardize_rope_params's per-layer loop then also needs to iterate over the entries actually present rather than set(layer_types), otherwise entries for non-instantiated types are left un-standardized:

-        elif layer_types is None or rope_parameters == {} or not set(rope_parameters.keys()).issubset(layer_types):
+        elif layer_types is None or rope_parameters == {} or not set(rope_parameters.keys()).issubset(ALLOWED_LAYER_TYPES):
...
-            for layer_type in set(layer_types):
+            for layer_type in set(rope_parameters.keys()):

-        if getattr(self, "layer_types", None) is not None and set(rope_parameters_dict.keys()).issubset(
-            self.layer_types
-        ):
+        if getattr(self, "layer_types", None) is not None and set(rope_parameters_dict.keys()).issubset(
+            ALLOWED_LAYER_TYPES
+        ):

I checked this by monkeypatching both methods at runtime on main: the warning and the injected keys go away for Gemma3TextConfig(num_hidden_layers=2) and Olmo3Config(num_hidden_layers=2), and Gemma3TextConfig(num_hidden_layers=6) is unchanged. I have not run the test suite against it, and there may be a reason to prefer a different shape here. Happy to open a PR if this direction looks right to you.

Context

Found in TRL CI, where the tiny 2-layer Gemma3 and Olmo3 test models hit this on every config load: huggingface/trl#6961

[zucchini-nlp]

Taking a look a bit later, fwiw we might move on to proper per_layer_config for all fields including the rope params but I didn't yet start designing it