diff --git a/optimum/exporters/executorch/tasks/causal_lm.py b/optimum/exporters/executorch/tasks/causal_lm.py index 7100ef1..bb04bbb 100644 --- a/optimum/exporters/executorch/tasks/causal_lm.py +++ b/optimum/exporters/executorch/tasks/causal_lm.py @@ -20,6 +20,7 @@ from ..integrations import CausalLMExportableModule from ..quantization import quantize_model_ from ..task_registry import register_task +from ..utils import disable_dynamic_rope_for_export # NOTE: It’s important to map the registered task name to the pipeline name in https://github.com/huggingface/transformers/blob/main/utils/update_metadata.py. @@ -69,11 +70,9 @@ def load_causal_lm_model(model_name_or_path: str, **kwargs) -> CausalLMExportabl ) config = kwargs.get("config") or AutoConfig.from_pretrained(model_name_or_path, gguf_file=gguf_file) - if hasattr(config, "rope_scaling") and config.rope_scaling is not None: - # NOTE: To make the model exportable we need to set the rope scaling to default to avoid hitting - # the data-dependent control flow in _longrope_frequency_update. Alternatively, users should rewrite - # that function to avoid the data-dependent control flow. - config.rope_scaling["type"] = "default" + # Downgrade longrope/dynamic RoPE to the static "default" RoPE, otherwise export fails on the + # data-dependent control flow in the RoPE frequency update. See helper for version details. + disable_dynamic_rope_for_export(config) if hasattr(config, "use_cache") and config.use_cache is False: config.use_cache = True diff --git a/optimum/exporters/executorch/tasks/multimodal_text_to_text.py b/optimum/exporters/executorch/tasks/multimodal_text_to_text.py index 7fc7811..e7035db 100644 --- a/optimum/exporters/executorch/tasks/multimodal_text_to_text.py +++ b/optimum/exporters/executorch/tasks/multimodal_text_to_text.py @@ -23,6 +23,7 @@ from ..integrations import MultiModalTextToTextExportableModule from ..quantization import quantize_model_ from ..task_registry import register_task +from ..utils import disable_dynamic_rope_for_export # NOTE: It's important to map the registered task name to the pipeline name in https://github.com/huggingface/transformers/blob/main/utils/update_metadata.py. @@ -90,9 +91,9 @@ def load_multimodal_text_to_text_model(model_name_or_path: str, **kwargs): if not (hasattr(config, "text_config")): raise ValueError(f"The model {model_name_or_path} does not have a `text_config`.") - if hasattr(config, "rope_scaling") and config.rope_scaling is not None: - # NOTE: Avoid hitting the data-dependent control flow in _longrope_frequency_update. - config.rope_scaling["type"] = "default" + # Avoid hitting the data-dependent control flow in the longrope/dynamic RoPE frequency update + # during export. See helper for version details. + disable_dynamic_rope_for_export(config) if hasattr(config, "use_cache") and config.use_cache is False: config.use_cache = True diff --git a/optimum/exporters/executorch/utils.py b/optimum/exporters/executorch/utils.py index d6c41f0..38a25b2 100644 --- a/optimum/exporters/executorch/utils.py +++ b/optimum/exporters/executorch/utils.py @@ -64,6 +64,39 @@ def save_config_to_constant_methods( return combined_metadata +def disable_dynamic_rope_for_export(config: PretrainedConfig) -> None: + """ + Force RoPE variants that use data-dependent control flow to the static ``"default"`` RoPE so the + model can be exported. + + ``longrope`` and ``dynamic`` RoPE recompute their frequencies inside the forward pass based on the + current sequence length (e.g. ``if seq_len > original_max_position_embeddings``). ``torch.export`` + /Dynamo cannot trace this data-dependent branching, so export fails. These RoPE types are downgraded + to ``"default"`` (which is static and traceable); other RoPE types (``linear``/``llama3``/``yarn``/...) + are static already and are left untouched. + + transformers>=5 renamed ``rope_scaling`` to ``rope_parameters`` (``rope_scaling`` is kept as an alias) + and dispatches on the ``"rope_type"`` key, whereas older versions read ``"type"``. Both keys are set so + the workaround is effective regardless of the installed transformers version. In transformers>=5 the + params may also be a nested dict keyed by layer type, which is handled here too. + """ + rope_params = getattr(config, "rope_scaling", None) + if not rope_params: + return + + # `rope_params` is either a single global dict or, in transformers>=5, a dict of per-layer-type dicts. + if all(isinstance(v, dict) for v in rope_params.values()): + param_dicts = list(rope_params.values()) + else: + param_dicts = [rope_params] + + for params in param_dicts: + rope_type = params.get("rope_type", params.get("type")) + if rope_type == "longrope" or (isinstance(rope_type, str) and "dynamic" in rope_type): + params["type"] = "default" + params["rope_type"] = "default" + + def apply_chat_template_with_fallback(processor, conversation, **kwargs): """ Apply chat template with fallback for external processors.