Repository: huggingface/transformers
File: tests/test_modeling_common.py (Python)

Code before:
                config.attention_probs_dropout_prob = 0

            # Update the head dim and try to update hidden size as well if present in config
            # NOTE: some models may have none if the values in sub-config, thus we check for `Noneness`
            head_dim = None
            if hasattr(config, "head_dim") and config.head_dim is not None:
                head_dim = config.head_dim
                config.head_dim = max(requested_dim, config.head_dim)

            cross_head_dim = None
            if hasattr(config, "cross_head_dim") and config.cross_head_dim is not None:
                cross_head_dim = config.cross_head_dim
                config.cross_head_dim = max(requested_dim, config.cross_head_dim)

            if (
                getattr(config, "hidden_size", None) is not None
                and getattr(config, "num_attention_heads", None) is not None
            ):
                # For some models, num_attention_heads is a list of ints: we take the max to maximize the multiplier
                num_attn_heads = getattr(config, "num_attention_heads")
                num_attn_heads = num_attn_heads if isinstance(num_attn_heads, int) else max(num_attn_heads)
                head_dim = head_dim if head_dim is not None else config.hidden_size // num_attn_heads
                config.hidden_size *= max(requested_dim // head_dim, 1)

            if (
                getattr(config, "decoder_hidden_size", None) is not None
                and getattr(config, "decoder_num_attention_heads", None) is not None
            ):
                decoder_head_dim = config.decoder_hidden_size // config.decoder_num_attention_heads
                config.decoder_hidden_size *= max(requested_dim // decoder_head_dim, 1)

            if (
                getattr(config, "cross_hidden_size", None) is not None
                and getattr(config, "cross_num_attention_heads", None) is not None
            ):
                cross_head_dim = (
                    cross_head_dim
                    if cross_head_dim is not None
                    else config.cross_hidden_size // config.cross_num_attention_heads
                )
                config.cross_hidden_size *= max(requested_dim // cross_head_dim, 1)

            # 3d rope also depends on the head dim
            # (we assume easy shapes here where we get to the requested head dim at least)
            if (
                getattr(config, "rope_parameters", None) is not None
                and len(config.rope_parameters.get("mrope_section", [])) > 0
            ):
                scaling_factor = max(requested_dim // (sum(config.rope_parameters["mrope_section"]) * 2), 1)
                config.rope_parameters["mrope_section"] = [
                    section * scaling_factor for section in config.rope_parameters["mrope_section"]

Review comment (suggestion):
Interesting seeing this the first time with the rms norm etc - does it interfere with models like gemma3 or qwen next which center around 1 instead of 0

Code after:
                    if hasattr(module.config, "return_loss"):
                        module.config.return_loss = False
                    # disable mamba kernels for every submodel (mamba, jamba)
                    if hasattr(module.config, "use_mamba_kernels"):
                        module.config.use_mamba_kernels = False
                # disable classifier cast for nllb-moe
                if hasattr(module, "_cast_classifier"):
                    module._cast_classifier = lambda *args, **kwargs: None
                # disable mamba mask update for ssms
                if hasattr(module, "_update_mamba_mask"):
                    module._update_mamba_mask = lambda attention_mask, *args, **kwargs: attention_mask
                if hasattr(module, "_update_linear_attn_mask"):
                    module._update_linear_attn_mask = lambda attention_mask, *args, **kwargs: attention_mask

            return model, inputs_dict

        for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__):
                if hasattr(self.model_tester, "prepare_config_and_inputs_for_model_class"):
                    config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_model_class(model_class)
                else:
                    config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
                inputs_dict = self._prepare_for_class(inputs_dict, model_class)
                set_config_for_less_flaky_test(config)
                model = model_class(config).eval().to(torch_device)
                set_model_for_less_flaky_test(model)

                # Prepare model and inputs for export
                model, inputs_dict = _prepare_for_export(model, inputs_dict)

                with torch.no_grad():
                    # Running the eager inference before the export to catch model/inputs comatibility issues, also sometimes after
                    # the export, the model used for export will return FakeTensors instead of real ones (torch compiler issue).
                    # This happens on cuda for example with (codegen, clvp, esm, gptj, levit, wav2vec2_bert and wav2vec2_conformer)
                    set_seed(1234)
                    eager_outputs = model(**copy.deepcopy(inputs_dict))
                    eager_outputs = _get_leaf_tensors(eager_outputs)
                    self.assertTrue(eager_outputs, "Eager outputs is empty.")

                try:
                    exported_program = torch.export.export(model, args=(), kwargs=copy.deepcopy(inputs_dict))
                except Exception as e:
                    raise e

                with torch.no_grad():
                    set_seed(1234)
                    exported_outputs = exported_program.module()(**copy.deepcopy(inputs_dict))
                    exported_outputs = _get_leaf_tensors(exported_outputs)
                    self.assertTrue(exported_outputs, "Exported outputs is empty.")

                # Check outputs closeness: