diff --git a/src/paddlefleet/transformer/hf_export.py b/src/paddlefleet/transformer/hf_export.py new file mode 100644 index 000000000..84aaae6bb --- /dev/null +++ b/src/paddlefleet/transformer/hf_export.py @@ -0,0 +1,419 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fleet <-> open-source (HuggingFace) config bridge: naming map + model-fact helpers. + +Single source of truth for the pieces of HF ``config.json`` export/import that are +**Fleet domain knowledge** -- i.e. facts about how a Fleet model's fields name and +structure map to the open-source (HF) side: + +- naming map ``FLEET_HF_FIELD_MAPPING`` / ``ROPE_SCALING_KEYMAP`` and the derived + ``HF_EXPORT_RULES`` / ``HF_IMPORT_RULES`` (+ ``rule_target``); +- rope structure: ``uses_yarn`` / ``pack_rope_scaling`` / ``unpack_rope_scaling``; +- window semantics: ``is_active_window`` / ``check_window_export_conflict`` / + ``swa_aware_import_rules``; +- MTP per-layer trimming ``trim_mtp_layers``; mHC injection ``inject_mhc_from_provider``. + +Pure Python, no ``paddle`` dependency. Sibling of ``TransformerConfig.transform_rules`` +(inbound HF->Fleet rename table in ``transformer_config.py``); this module is the +outbound / bidirectional counterpart. The export *strategy* (whitelist/blacklist, +orchestration, trainer I/O) lives above this layer, in the caller (erniebot). +""" + +import functools + + +def hidden_act_to_hf(act): + """Fleet ``hidden_act`` (callable or str) -> HF ``hidden_act`` name. + + Reverses ``TransformerConfig._process_attribute`` (``transformer_config.py``), + which turns an HF string into a callable: ``"gelu_pytorch_tanh"`` becomes + ``functools.partial(F.gelu, approximate=True)``, ``"situ"`` a named function, + and any other name ``n`` becomes ``getattr(F, n)`` (whose ``__name__`` is + ``n``). A ``partial`` has no ``__name__``, so reading it directly would raise + ``AttributeError``; the gelu-tanh partial is matched structurally instead. + """ + if isinstance(act, str): + return act + if isinstance(act, functools.partial): + func_name = getattr(act.func, "__name__", None) + if func_name == "gelu" and act.keywords.get("approximate") is True: + return "gelu_pytorch_tanh" + if func_name is not None: + return func_name + return act + return getattr(act, "__name__", act) + + +# rope_scaling structural map: HF nested key -> Fleet flat field name. +ROPE_SCALING_KEYMAP = { + "type": "rope_type", + "factor": "rotary_scaling_factor", + "original_max_position_embeddings": "original_max_position_embeddings", + "beta_fast": "beta_fast", + "beta_slow": "beta_slow", + "mscale": "mscale", + "mscale_all_dim": "mscale_all_dim", +} + + +# Single source of truth for bidirectional field mapping: +# (fleet_key, hf_key, Fleet -> HF converter, HF -> Fleet converter). +FLEET_HF_FIELD_MAPPING = [ + ("multi_latent_attention", "use_mla", None, None), + ("gated_attention", "use_gated_attn", None, None), + ("rotary_interleaved", "rope_interleave", None, None), + ("csa_compress_ratios", "compress_ratios", None, None), + ("csa_window_size", "sliding_window", None, None), + ("csa_compress_rotary_base", "compress_rope_theta", None, None), + ( + "params_dtype", + "torch_dtype", + lambda v: str(v).replace("paddle.", ""), + lambda v: v.replace("paddle.", "") if isinstance(v, str) else v, + ), + ( + "hidden_act", + "hidden_act", + hidden_act_to_hf, + None, + ), + ("window_attn_skip_freq", "hybrid_layer_pattern", None, None), + ( + "multimax_modules", + "multimax", + lambda v: v[0] if isinstance(v, (list, tuple)) and v else v, + lambda v: list(v) + if isinstance(v, (list, tuple)) + else ([v] if v is not None else v), + ), + ("num_residual_streams", "hc_mult", None, None), + ("mhc_sinkhorn_iterations", "hc_sinkhorn_iters", None, None), +] + +# Fleet -> HF config.json rename / value rules. +HF_EXPORT_RULES = { + fleet_key: (hf_key, export_fn) if export_fn else hf_key + for fleet_key, hf_key, export_fn, _ in FLEET_HF_FIELD_MAPPING +} + +# HF -> Fleet config.json reverse rename / value rules. +HF_IMPORT_RULES = { + hf_key: (fleet_key, import_fn) if import_fn else fleet_key + for fleet_key, hf_key, _, import_fn in FLEET_HF_FIELD_MAPPING +} + + +def rule_target(key, rules): + """The HF-side name a raw source key maps to under ``rules``. + + Keys without a rename rule pass through unchanged (target == source). + Used to translate a set of raw model_config keys into the HF names they + appear as in the built config, so they can be unioned with the whitelist. + """ + spec = rules.get(key) + if spec is None: + return key + return spec[0] if isinstance(spec, tuple) else spec + + +def source_or_provider(raw, provider, key): + """Value of ``key`` from the export source, falling back to the provider. + + Single provider-fallback accessor shared by the rope / mHC paths: the + pristine export source (``raw``) wins; when it does not declare ``key`` the + resolved provider config supplies the default. Returns ``None`` if neither + has it. + """ + val = raw.get(key, None) + if val is None and provider is not None: + val = getattr(provider, key, None) + return val + + +# dsv4_hybrid per-layer compress-ratio conventions. Source of truth: +# paddlefleet/transformer/dsv4_hybrid_attention.py (``compress_ratio == 128`` -> +# HCA; ``2 <= ratio < 128`` -> CSA; ``-1``/``0`` -> window; ``-2`` -> MLA). +# Compressed layers (ratio > 1) default to YaRN; window layers stay plain RoPE; +# MLA (-2) layers are built on the plain MLA path and follow global ``rope_type``. +HCA_COMPRESS_RATIO = 128 +CSA_COMPRESS_RATIO_MIN = ( + 2 # CSA range is [CSA_COMPRESS_RATIO_MIN, HCA_COMPRESS_RATIO) +) +MLA_COMPRESS_RATIO = -2 + + +def uses_yarn(raw, provider=None): + """Whether the model actually applies YaRN RoPE on any attention layer. + + Decided from the model *structure*, not from ``rotary_scaling_factor``: + those flat fields (``rotary_scaling_factor`` / ``original_max_position_embeddings`` + / ``beta_fast`` / ...) are YaRN-only parameters that sit in the config even + when RoPE runs in plain ``"rope"`` mode, so keying off ``factor != 1.0`` + yields false positives. + + The global ``rope_type`` is authoritative only for **non-hybrid** models + (plain MLA / DSA, which read it directly in ``multi_latent_attention.py`` / + ``dsa_attention.py``). For ``dsv4_hybrid`` the RoPE mode is decided *per + layer* (``dsv4_hybrid_attention.py``), so the global value must NOT + short-circuit -- otherwise ``rope_type="yarn"`` with every layer overridden + to plain RoPE would wrongly export a YaRN ``rope_scaling``: + + - HCA (ratio 128) / CSA (ratio in [2, 128)) layers default to YaRN and are + overridden per type by ``hca_rope_type`` / ``csa_rope_type``; a layer uses + YaRN unless its override forces ``"rope"``; + - MLA (ratio -2) layers are built on the plain MLA path and follow the + global ``rope_type``; + - window (ratio -1 / 0) layers stay plain RoPE. + """ + + def _pick(key): + return source_or_provider(raw, provider, key) + + global_rope_type = _pick("rope_type") + + if _pick("experimental_attention_variant") != "dsv4_hybrid": + return global_rope_type == "yarn" + + ratios = _pick("csa_compress_ratios") or [] + hca_rope_type = _pick("hca_rope_type") + csa_rope_type = _pick("csa_rope_type") + for ratio in ratios: + if not isinstance(ratio, (int, float)): + continue + if ratio == HCA_COMPRESS_RATIO: + per_type = hca_rope_type # default YaRN unless overridden to "rope" + elif CSA_COMPRESS_RATIO_MIN <= ratio < HCA_COMPRESS_RATIO: + per_type = csa_rope_type # default YaRN unless overridden to "rope" + elif ratio == MLA_COMPRESS_RATIO: + if global_rope_type == "yarn": # MLA follows global rope_type + return True + continue + else: + continue # window (-1 / 0): plain RoPE + if per_type != "rope": + return True + return False + + +def pack_rope_scaling(raw, provider=None): + """Fleet flat YARN fields -> HF nested ``rope_scaling`` dict; None when no YARN. + + Whether YaRN is active is decided structurally by :func:`uses_yarn` (global + ``rope_type == "yarn"`` or a ``dsv4_hybrid`` compressed layer that defaults + to YaRN), *not* from ``rotary_scaling_factor``. ``raw`` is the pristine + export source; it carries the model_config values (``rotary_scaling_factor`` + -> ``factor``, ``original_max_position_embeddings``). ``provider`` is the + resolved provider config, used to fill the YARN sub-fields the pristine + source does not declare (``beta_fast`` / ``beta_slow`` / ...). When YaRN is + active the ``type`` is normalized to ``"yarn"``. + + The neutral mscale defaults (``mscale`` == 1.0 / ``mscale_all_dim`` == 0.0) + are omitted so the emitted dict stays minimal (type / factor / + original_max_position_embeddings / beta_fast / beta_slow). + """ + if not uses_yarn(raw, provider): + return None + + out = {} + for hf_key, fleet_key in ROPE_SCALING_KEYMAP.items(): + val = source_or_provider(raw, provider, fleet_key) + if val is None: + continue + if hf_key == "mscale" and val == 1.0: + continue + if hf_key == "mscale_all_dim" and val == 0.0: + continue + out[hf_key] = val + if out.get("type") in (None, "default", "rope"): + out["type"] = "yarn" + return out or None + + +def unpack_rope_scaling(rope_scaling): + """HF nested rope_scaling dict -> Fleet flat fields dict. + + The RoPE type accepts both the current HF canonical key ``rope_type`` and + its legacy alias ``type`` (``ROPE_SCALING_KEYMAP`` only carries ``type``); + ``rope_type`` wins when both are present. Without this a modern config such + as ``{"rope_type": "yarn", "factor": 4}`` would drop the YaRN type and fall + back to plain RoPE. + """ + if not rope_scaling or not isinstance(rope_scaling, dict): + return {} + out = { + fleet_key: rope_scaling[hf_key] + for hf_key, fleet_key in ROPE_SCALING_KEYMAP.items() + if hf_key in rope_scaling + } + if "rope_type" in rope_scaling: + out["rope_type"] = rope_scaling["rope_type"] + return out + + +def is_active_window(sw): + """Whether a window-size field declares an active (non-zero) window. + + Neutral naming: used both for the native SWA ``sliding_window`` field and + the CSA ``csa_window_size`` field. A scalar is active when truthy; a + per-layer list/tuple is active when any entry is truthy. + """ + if sw in (None, 0, (), []): + return False + if isinstance(sw, (list, tuple)): + return any(x for x in sw) + return True + + +def check_window_export_conflict(raw): + """SWA and CSA both target the HF ``sliding_window`` field; reject coexistence. + + The SWA path exports its native ``sliding_window`` directly, while the CSA + path maps ``csa_window_size`` -> ``sliding_window`` (see ``HF_EXPORT_RULES``). + Emitting both into the single HF ``sliding_window`` field is not supported + yet, so a config declaring both active windows is rejected up front instead + of silently letting the CSA rename overwrite the SWA value. + """ + if is_active_window(raw.get("sliding_window")) and is_active_window( + raw.get("csa_window_size") + ): + raise ValueError( + "Both 'sliding_window' (SWA) and 'csa_window_size' (CSA) are set; " + "exporting both to the HF 'sliding_window' field is not supported yet." + ) + + +SWA_MARKER_HF_KEYS = [ + "add_swa_attention_sink_bias", + "swa_head_dim", + "swa_v_head_dim", + "swa_num_attention_heads", + "swa_num_key_value_heads", + "swa_rope_theta", + "swa_qk_nope_head_dim", + "swa_qk_rope_head_dim", + "head_wise_swa_ratio", +] + + +def is_swa_config(hf_config): + """Whether an HF config declares SWA-specific companion fields.""" + return any( + key in SWA_MARKER_HF_KEYS or key.startswith("swa_") for key in hf_config + ) + + +# HF markers that identify a DSv4 CSA config, whose ``sliding_window`` is the +# CSA window (renamed to ``csa_window_size`` on import). ``compress_ratios`` is +# the HF name of ``csa_compress_ratios``; ``csa_compress_ratios`` covers a +# Fleet-native input; ``experimental_attention_variant == "dsv4_hybrid"`` is the +# structural flag. Any other ``sliding_window`` is a native SWA window. +CSA_MARKER_HF_KEYS = ("compress_ratios", "csa_compress_ratios") + + +def is_csa_config(hf_config): + """Whether an HF config's ``sliding_window`` is a DSv4 CSA window. + + Only a dsv4_hybrid / compress-ratio config routes HF ``sliding_window`` to + Fleet ``csa_window_size``. A standard HF ``sliding_window`` (e.g. Mistral's + bare ``sliding_window=4096``) carries no CSA markers and must keep its + native name on import. + """ + if hf_config.get("experimental_attention_variant") == "dsv4_hybrid": + return True + return any(key in hf_config for key in CSA_MARKER_HF_KEYS) + + +def swa_aware_import_rules(hf_config, rules): + """Keep the ``sliding_window`` -> ``csa_window_size`` rename only for CSA. + + On import HF ``sliding_window`` has two possible owners: DSv4 CSA, where it + is the compressed-attention window and must become ``csa_window_size``, and + everything else (native SWA -- Mistral's bare ``sliding_window``, or a Fleet + SWA config carrying ``swa_*`` companions), which must keep its own name. + The rename therefore applies *only* when the config carries CSA markers; + otherwise ``sliding_window`` passes through unchanged, so a standard HF SWA + config is no longer silently re-homed onto ``csa_window_size``. + """ + if "sliding_window" not in rules or is_csa_config(hf_config): + return rules + return { + hf_key: spec + for hf_key, spec in rules.items() + if hf_key != "sliding_window" + } + + +# Per-layer list fields whose trailing MTP-layer entries are trimmed on import. +# Both Fleet-native and HF-renamed names are listed so the same key set works +# whether the input is an HF ``config.json`` or a Fleet ``model_config.json``: +# ``compress_ratios`` <- ``csa_compress_ratios`` and ``hybrid_layer_pattern`` <- +# ``window_attn_skip_freq``. Missing the HF name leaves an over-long list that +# fails ``TransformerConfig`` validation once ``num_nextn_predict_layers`` is 0. +MTP_TRIM_KEYS = ( + "window_attn_skip_freq", + "hybrid_layer_pattern", + "csa_compress_ratios", + "compress_ratios", +) + + +def trim_mtp_layers(out): + """Drop trailing MTP-layer entries from per-layer list fields, in place. + + When ``num_nextn_predict_layers`` > 0 the per-layer lists carry extra + trailing entries for the MTP layer(s), which the inference-side base config + should not include. Trims each present list in ``MTP_TRIM_KEYS`` and zeroes + ``num_nextn_predict_layers``. + """ + mtp_layers = out.get("num_nextn_predict_layers", 0) + if mtp_layers and mtp_layers > 0: + for key in MTP_TRIM_KEYS: + val = out.get(key) + if isinstance(val, list) and len(val) > mtp_layers: + out[key] = val[:-mtp_layers] + out["num_nextn_predict_layers"] = 0 + + +# mHC (Hyper-Connections) fields are Fleet TransformerConfig defaults that live +# only on the resolved provider (model.config at train time), not on the +# pristine export ``source``. When hyper-connections are active they are sourced +# off the provider here (Fleet names, later renamed to HF names via +# ``FLEET_HF_FIELD_MAPPING``); ``hc_eps`` has no config field (a module constant +# in hyper_connection.py) and is injected directly. +MHC_PROVIDER_FIELDS = ("num_residual_streams", "mhc_sinkhorn_iterations") +MHC_HC_EPS = 1e-6 + + +def inject_mhc_from_provider(raw, provider): + """Merge mHC fields off the resolved provider into ``raw`` (Fleet names). + + No-op unless ``provider.enable_hyper_connections`` is truthy. Uses the same + provider-fallback semantics as the rope path (:func:`source_or_provider`, + source wins), but applies it as an *early mutation* of ``raw`` -- rather + than a pack-time read -- because these fields must flow through the rename + rules (``num_residual_streams`` -> ``hc_mult`` / + ``mhc_sinkhorn_iterations`` -> ``hc_sinkhorn_iters``) to land as top-level + HF keys. ``hc_eps`` has no config field (a module constant in + hyper_connection.py) and is injected directly. + """ + if provider is None or not getattr( + provider, "enable_hyper_connections", False + ): + return + for name in MHC_PROVIDER_FIELDS: + val = source_or_provider(raw, provider, name) + if val is not None: + raw.setdefault(name, val) + raw.setdefault("hc_eps", MHC_HC_EPS) diff --git a/tests/single_card_tests/transformer/test_hf_export.py b/tests/single_card_tests/transformer/test_hf_export.py new file mode 100644 index 000000000..03ec85251 --- /dev/null +++ b/tests/single_card_tests/transformer/test_hf_export.py @@ -0,0 +1,401 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Fleet<->HF config bridge (paddlefleet.transformer.hf_export). + +Pure-Python (no paddle): covers the naming map / rules, structural YaRN +detection, rope_scaling pack/unpack, window semantics, MTP trimming, mHC +injection, and the provider-fallback accessor. +""" + +import functools +import unittest +from types import SimpleNamespace + +from paddlefleet.transformer.hf_export import ( + FLEET_HF_FIELD_MAPPING, + HF_EXPORT_RULES, + HF_IMPORT_RULES, + check_window_export_conflict, + hidden_act_to_hf, + inject_mhc_from_provider, + is_active_window, + is_csa_config, + pack_rope_scaling, + rule_target, + source_or_provider, + swa_aware_import_rules, + trim_mtp_layers, + unpack_rope_scaling, + uses_yarn, +) + + +class TestRulesAndMapping(unittest.TestCase): + def test_export_import_rules_are_inverse_on_renames(self): + for fleet_key, spec in HF_EXPORT_RULES.items(): + hf_key = spec[0] if isinstance(spec, tuple) else spec + self.assertIn(hf_key, HF_IMPORT_RULES) + back = HF_IMPORT_RULES[hf_key] + back_key = back[0] if isinstance(back, tuple) else back + self.assertEqual(back_key, fleet_key) + + def test_mapping_entry_count(self): + self.assertEqual(len(FLEET_HF_FIELD_MAPPING), len(HF_EXPORT_RULES)) + + def test_rule_target(self): + self.assertEqual( + rule_target("csa_window_size", HF_EXPORT_RULES), "sliding_window" + ) + self.assertEqual( + rule_target("unknown_key", HF_EXPORT_RULES), "unknown_key" + ) + + def test_value_converters(self): + # params_dtype -> torch_dtype strips the "paddle." prefix on export. + spec = HF_EXPORT_RULES["params_dtype"] + self.assertEqual(spec[0], "torch_dtype") + self.assertEqual(spec[1]("paddle.bfloat16"), "bfloat16") + # multimax_modules list -> scalar on export. + spec = HF_EXPORT_RULES["multimax_modules"] + self.assertEqual(spec[1](["lm_head"]), "lm_head") + + def test_dtype_import_returns_canonical_string(self): + # HF torch_dtype "bfloat16" must import to a bare Paddle-accepted + # dtype string, NOT "paddle.bfloat16" (create_parameter rejects it). + spec = HF_IMPORT_RULES["torch_dtype"] + self.assertEqual(spec[0], "params_dtype") + self.assertEqual(spec[1]("bfloat16"), "bfloat16") + # A Fleet-native value still normalises to the bare name. + self.assertEqual(spec[1]("paddle.bfloat16"), "bfloat16") + + def test_hidden_act_export_handles_partial(self): + # gelu_pytorch_tanh round-trips through functools.partial (no __name__). + act_fn = HF_EXPORT_RULES["hidden_act"][1] + partial = functools.partial( + lambda x, approximate=False: x, approximate=True + ) + partial.func.__name__ = "gelu" + self.assertEqual(hidden_act_to_hf(partial), "gelu_pytorch_tanh") + # exposed via the export rule too. + self.assertEqual(act_fn(partial), "gelu_pytorch_tanh") + + def test_hidden_act_export_named_callable(self): + def silu(x): + return x + + self.assertEqual(hidden_act_to_hf(silu), "silu") + + def test_hidden_act_export_passthrough_string(self): + self.assertEqual(hidden_act_to_hf("relu"), "relu") + + +class TestSourceOrProvider(unittest.TestCase): + def test_source_wins(self): + prov = SimpleNamespace(k=2) + self.assertEqual(source_or_provider({"k": 1}, prov, "k"), 1) + + def test_provider_fallback(self): + prov = SimpleNamespace(k=2) + self.assertEqual(source_or_provider({}, prov, "k"), 2) + + def test_none_when_absent(self): + self.assertIsNone(source_or_provider({}, None, "k")) + + +class TestUsesYarn(unittest.TestCase): + def test_global_rope_type_yarn(self): + self.assertTrue(uses_yarn({"rope_type": "yarn"})) + + def test_plain_rope_not_yarn(self): + # leftover rotary_scaling_factor must NOT trigger yarn on plain rope. + self.assertFalse( + uses_yarn({"rope_type": "rope", "rotary_scaling_factor": 16}) + ) + + def test_dsv4_hybrid_hca_defaults_yarn(self): + self.assertTrue( + uses_yarn( + { + "rope_type": "rope", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [128, -2], + } + ) + ) + + def test_dsv4_hybrid_hca_override_rope(self): + self.assertFalse( + uses_yarn( + { + "rope_type": "rope", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [128, -2], + "hca_rope_type": "rope", + } + ) + ) + + def test_dsv4_hybrid_csa_layer_defaults_yarn(self): + self.assertTrue( + uses_yarn( + { + "rope_type": "rope", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [64, -2], + } + ) + ) + + def test_dsv4_hybrid_no_compressed_layer(self): + self.assertFalse( + uses_yarn( + { + "rope_type": "rope", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [-2, 0, -2], + } + ) + ) + + def test_dsv4_hybrid_global_yarn_but_all_layers_overridden_rope(self): + # Global rope_type="yarn" must NOT short-circuit for dsv4_hybrid: with + # HCA/CSA forced to "rope" and only a window layer, no layer uses YaRN. + self.assertFalse( + uses_yarn( + { + "rope_type": "yarn", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [128, 64, 0], + "hca_rope_type": "rope", + "csa_rope_type": "rope", + } + ) + ) + + def test_dsv4_hybrid_mla_layer_follows_global_yarn(self): + # An MLA (-2) layer is built on the plain MLA path and follows the + # global rope_type, so global "yarn" activates YaRN through it. + self.assertTrue( + uses_yarn( + { + "rope_type": "yarn", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [-2, 0], + "hca_rope_type": "rope", + "csa_rope_type": "rope", + } + ) + ) + + def test_dsv4_hybrid_mla_layer_global_rope_not_yarn(self): + # MLA layer with a non-yarn global rope_type stays plain RoPE. + self.assertFalse( + uses_yarn( + { + "rope_type": "rope", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [-2, 0], + } + ) + ) + + def test_provider_fallback_rope_type(self): + prov = SimpleNamespace(rope_type="yarn") + self.assertTrue(uses_yarn({}, prov)) + + +class TestRopeScaling(unittest.TestCase): + def test_pack_yarn(self): + rs = pack_rope_scaling( + { + "rope_type": "yarn", + "rotary_scaling_factor": 4.0, + "original_max_position_embeddings": 4096, + "beta_fast": 32, + "beta_slow": 1, + } + ) + self.assertEqual(rs["type"], "yarn") + self.assertEqual(rs["factor"], 4.0) + self.assertEqual(rs["beta_fast"], 32) + + def test_pack_none_when_no_yarn(self): + self.assertIsNone( + pack_rope_scaling( + {"rope_type": "rope", "rotary_scaling_factor": 16} + ) + ) + + def test_pack_normalizes_type_to_yarn(self): + # dsv4_hybrid compressed layer -> yarn active; global type "rope" normalized. + rs = pack_rope_scaling( + { + "rope_type": "rope", + "experimental_attention_variant": "dsv4_hybrid", + "csa_compress_ratios": [128], + "rotary_scaling_factor": 8, + } + ) + self.assertEqual(rs["type"], "yarn") + + def test_pack_omits_neutral_mscale(self): + rs = pack_rope_scaling( + {"rope_type": "yarn", "mscale": 1.0, "mscale_all_dim": 0.0} + ) + self.assertNotIn("mscale", rs) + self.assertNotIn("mscale_all_dim", rs) + + def test_unpack(self): + flat = unpack_rope_scaling( + {"type": "yarn", "factor": 4.0, "beta_fast": 32} + ) + self.assertEqual(flat["rope_type"], "yarn") + self.assertEqual(flat["rotary_scaling_factor"], 4.0) + self.assertEqual(flat["beta_fast"], 32) + + def test_unpack_canonical_rope_type(self): + # Current HF configs use "rope_type" (not the legacy "type" alias); + # the YaRN type must survive import. + flat = unpack_rope_scaling({"rope_type": "yarn", "factor": 4}) + self.assertEqual(flat["rope_type"], "yarn") + self.assertEqual(flat["rotary_scaling_factor"], 4) + + def test_unpack_rope_type_wins_over_legacy_type(self): + flat = unpack_rope_scaling({"rope_type": "yarn", "type": "linear"}) + self.assertEqual(flat["rope_type"], "yarn") + + def test_unpack_empty(self): + self.assertEqual(unpack_rope_scaling(None), {}) + + +class TestWindow(unittest.TestCase): + def test_is_active_window(self): + self.assertFalse(is_active_window(0)) + self.assertFalse(is_active_window(None)) + self.assertFalse(is_active_window([0, 0])) + self.assertTrue(is_active_window(128)) + self.assertTrue(is_active_window([0, 128, 0])) + + def test_conflict_raises(self): + with self.assertRaises(ValueError): + check_window_export_conflict( + {"sliding_window": 4096, "csa_window_size": 128} + ) + + def test_no_conflict_ok(self): + check_window_export_conflict({"sliding_window": 4096}) # no raise + check_window_export_conflict({"csa_window_size": 128}) # no raise + + def test_swa_aware_import_rules_drops_sliding_window_for_swa(self): + rules = {"sliding_window": "csa_window_size", "other": "x"} + swa_cfg = {"swa_num_attention_heads": 8} + out = swa_aware_import_rules(swa_cfg, rules) + self.assertNotIn("sliding_window", out) + self.assertIn("other", out) + + def test_swa_aware_import_rules_native_sliding_window_kept(self): + # Standard HF SWA (Mistral: bare sliding_window, no CSA markers) must + # keep its native name rather than be re-homed to csa_window_size. + rules = {"sliding_window": "csa_window_size", "other": "x"} + out = swa_aware_import_rules({"sliding_window": 4096}, rules) + self.assertNotIn("sliding_window", out) + self.assertIn("other", out) + + def test_swa_aware_import_rules_keeps_rename_for_csa(self): + rules = {"sliding_window": "csa_window_size", "other": "x"} + out = swa_aware_import_rules({"compress_ratios": [128, 64]}, rules) + self.assertEqual(out["sliding_window"], "csa_window_size") + self.assertIn("other", out) + + def test_swa_aware_import_rules_keeps_rename_for_dsv4_variant(self): + rules = {"sliding_window": "csa_window_size"} + out = swa_aware_import_rules( + {"experimental_attention_variant": "dsv4_hybrid"}, rules + ) + self.assertIn("sliding_window", out) + + def test_is_csa_config(self): + self.assertTrue(is_csa_config({"compress_ratios": [128]})) + self.assertTrue( + is_csa_config({"experimental_attention_variant": "dsv4_hybrid"}) + ) + self.assertFalse(is_csa_config({"sliding_window": 4096})) + + +class TestMtpTrim(unittest.TestCase): + def test_trim(self): + out = { + "num_nextn_predict_layers": 1, + "window_attn_skip_freq": [1, 2, 3], + "csa_compress_ratios": [4, 5, 6], + } + trim_mtp_layers(out) + self.assertEqual(out["window_attn_skip_freq"], [1, 2]) + self.assertEqual(out["csa_compress_ratios"], [4, 5]) + self.assertEqual(out["num_nextn_predict_layers"], 0) + + def test_trim_hf_renamed_keys(self): + # HF-side names: hybrid_layer_pattern <- window_attn_skip_freq and + # compress_ratios <- csa_compress_ratios must also be trimmed. + out = { + "num_nextn_predict_layers": 1, + "hybrid_layer_pattern": [1, 2, 3], + "compress_ratios": [4, 5, 6], + } + trim_mtp_layers(out) + self.assertEqual(out["hybrid_layer_pattern"], [1, 2]) + self.assertEqual(out["compress_ratios"], [4, 5]) + self.assertEqual(out["num_nextn_predict_layers"], 0) + + def test_no_trim_when_no_mtp(self): + out = {"num_nextn_predict_layers": 0, "csa_compress_ratios": [4, 5, 6]} + trim_mtp_layers(out) + self.assertEqual(out["csa_compress_ratios"], [4, 5, 6]) + + +class TestMhcInjection(unittest.TestCase): + def test_inject(self): + prov = SimpleNamespace( + enable_hyper_connections=True, + num_residual_streams=4, + mhc_sinkhorn_iterations=3, + ) + raw = {} + inject_mhc_from_provider(raw, prov) + self.assertEqual(raw["num_residual_streams"], 4) + self.assertEqual(raw["mhc_sinkhorn_iterations"], 3) + self.assertEqual(raw["hc_eps"], 1e-6) + + def test_disabled_noop(self): + prov = SimpleNamespace( + enable_hyper_connections=False, num_residual_streams=4 + ) + raw = {} + inject_mhc_from_provider(raw, prov) + self.assertEqual(raw, {}) + + def test_source_value_wins(self): + prov = SimpleNamespace( + enable_hyper_connections=True, + num_residual_streams=4, + mhc_sinkhorn_iterations=3, + ) + raw = {"num_residual_streams": 9} + inject_mhc_from_provider(raw, prov) + self.assertEqual(raw["num_residual_streams"], 9) + + +if __name__ == "__main__": + unittest.main()