From 84b97f662001ed7a18172a87619ddd86cfaf98b0 Mon Sep 17 00:00:00 2001 From: Brian Bell Date: Fri, 24 Jul 2026 23:32:25 -0500 Subject: [PATCH] compat: transformers 5.x + custom-architecture loader shims, nemotron_h support Loading custom-architecture checkpoints (Nemotron Omni family) against transformers 5.14 surfaced API drift and wrapper mismatches: - flash_attention_2 pinned in config but flash-attn not installed: recursively fall back to eager attention, incl. nested sub-configs (multimodal llm_config). - all_tied_weights_keys shim: expected by the accelerate device-map integration and assigned by 5.x post_init, but absent on older remote code. Settable property aggregating legacy _tied_weights_keys, filtering keys that don't resolve on multimodal wrappers. - prepare_inputs_for_generation: 5.x may pass cache_position=None; older remote code assumes a tensor. Wrap and synthesize it. - Legacy list-style _tied_weights_keys normalized post-load so save_pretrained works (emptied when tie_word_embeddings is false). - Multimodal wrappers whose forward() requires media inputs are unwrapped to their language_model submodule for text-only surgery. - Register nemotron_h architecture (hybrid Mamba/attention/MoE, all layer content under layer.mixer) in strategies/utils.py. --- obliteratus/models/loader.py | 139 ++++++++++++++++++++++++++++++++ obliteratus/strategies/utils.py | 8 ++ 2 files changed, 147 insertions(+) diff --git a/obliteratus/models/loader.py b/obliteratus/models/loader.py index ab2e1b2..c623e5e 100644 --- a/obliteratus/models/loader.py +++ b/obliteratus/models/loader.py @@ -252,6 +252,51 @@ except Exception: pass +# ── 12. PreTrainedModel.all_tied_weights_keys ────────────────────── +# Referenced by the accelerate device-map integration in transformers 5.x, +# but older remote-code model classes never define it. Aggregate the +# legacy ``_tied_weights_keys`` (list of names or dict) across modules into +# the {param_name: target_name} mapping accelerate expects. +# Affected: Nemotron Omni (NemotronH_Nano_Omni_Reasoning_V3) +try: + from transformers import PreTrainedModel as _PTM + if not hasattr(_PTM, "all_tied_weights_keys"): + def _get_all_tied_weights_keys(self): + cached = self.__dict__.get("_all_tied_weights_keys_store") + if cached is not None: + return cached + tied: dict = {} + for module in self.modules(): + keys = getattr(module, "_tied_weights_keys", None) + if not keys: + continue + if isinstance(keys, dict): + candidates = keys + else: # legacy list of patterns/names — one tied group + names = list(keys) + candidates = {name: names[0] for name in names} + for name, target in candidates.items(): + # Only expose keys that actually resolve on this model — + # multimodal wrappers inherit "lm_head.weight" from their + # LM submodule's config but have no top-level lm_head. + try: + self.get_submodule(name.rsplit(".", 1)[0]) + except AttributeError: + continue + tied.setdefault(name, target) + return tied + + def _set_all_tied_weights_keys(self, value): + # transformers 5.x post_init() assigns this attribute. + self.__dict__["_all_tied_weights_keys_store"] = value + + _PTM.all_tied_weights_keys = property( + _get_all_tied_weights_keys, _set_all_tied_weights_keys, + ) +except Exception: + pass + + # ── Deferred shims ────────────────────────────────────────────────── # Some patches must wait until the _LazyModule has fully initialized # (it replaces its __dict__ during bootstrap). We apply these once, @@ -633,6 +678,41 @@ def load_model( f"Consider using quantization='4bit' or quantization='8bit'." ) + # Some repos pin attn_implementation="flash_attention_2" in config.json. + # If flash-attn isn't installed, loading would fail — fall back to eager. + # Multimodal wrappers pin it on nested sub-configs too (e.g. Nemotron + # Omni's llm_config), so walk those as well. + def _force_eager_cfg(cfg) -> None: + from transformers import PretrainedConfig as _PTC + if getattr(cfg, "_attn_implementation", None) == "flash_attention_2": + cfg._attn_implementation = "eager" + for sub in vars(cfg).values(): + if isinstance(sub, _PTC): + _force_eager_cfg(sub) + + def _has_flash_attention(cfg) -> bool: + from transformers import PretrainedConfig as _PTC + if getattr(cfg, "_attn_implementation", None) == "flash_attention_2": + return True + return any( + _has_flash_attention(sub) + for sub in vars(cfg).values() + if isinstance(sub, _PTC) + ) + + if _has_flash_attention(config): + try: + from transformers.utils import is_flash_attn_2_available + _fa2 = is_flash_attn_2_available() + except Exception: + _fa2 = False + if not _fa2: + logger.warning( + "config.json pins flash_attention_2 but flash-attn is not " + "installed — falling back to eager attention." + ) + _force_eager_cfg(config) + model_cls = _select_model_class(task, config) load_kwargs: dict = { "pretrained_model_name_or_path": model_name, @@ -761,6 +841,65 @@ def load_model( model.eval() + # Multimodal wrappers (e.g. Nemotron Omni) require pixel_values/audio in + # forward() — unusable for text-only surgery. Unwrap the language-model + # submodule; the pipeline only ever edits the LM. + if hasattr(model, "language_model"): + import inspect as _inspect + lm = getattr(model, "language_model", None) + try: + params = _inspect.signature(model.forward).parameters + needs_media = any( + p.default is _inspect.Parameter.empty + and p.kind in (_inspect.Parameter.POSITIONAL_ONLY, + _inspect.Parameter.POSITIONAL_OR_KEYWORD) + and p.name not in ("self", "input_ids") + for p in params.values() + ) + except (TypeError, ValueError): + needs_media = False + if needs_media and isinstance(lm, PreTrainedModel): + logger.warning( + "Multimodal wrapper %s requires media inputs in forward() — " + "unwrapping language_model (%s) for text-only surgery.", + type(model).__name__, type(lm).__name__, + ) + model = lm + config = model.config + + # transformers 5.x may call prepare_inputs_for_generation with + # cache_position=None; older remote code (Nemotron-H) assumes a tensor. + _pifg = getattr(model, "prepare_inputs_for_generation", None) + if _pifg is not None: + def _prepare_inputs_safe(input_ids, *args, cache_position=None, **kwargs): + if cache_position is None: + past_len = 0 + pkv = kwargs.get("past_key_values") or (args[0] if args else None) + if pkv is not None: + try: + past_len = pkv.get_seq_length() + except AttributeError: + try: + past_len = pkv[0][0].shape[-2] + except Exception: + past_len = 0 + cache_position = torch.arange( + past_len, past_len + input_ids.shape[-1], + device=input_ids.device, + ) + return _pifg(input_ids, *args, cache_position=cache_position, **kwargs) + + model.prepare_inputs_for_generation = _prepare_inputs_safe + + # transformers 5.x expects module._tied_weights_keys as a dict; older + # remote code (Nemotron-H) defines it as a list. Normalize: when word + # embeddings aren't actually tied, nothing should be dropped at save. + _tie_embeddings = bool(getattr(config, "tie_word_embeddings", False)) + for _m in model.modules(): + _twk = getattr(_m, "_tied_weights_keys", None) + if isinstance(_twk, list): + _m._tied_weights_keys = {k: k for k in _twk} if _tie_embeddings else {} + # Free accelerator cache after loading dev.empty_cache() diff --git a/obliteratus/strategies/utils.py b/obliteratus/strategies/utils.py index 7924a2e..7e22e52 100644 --- a/obliteratus/strategies/utils.py +++ b/obliteratus/strategies/utils.py @@ -47,6 +47,7 @@ _LAYER_ATTR_PATHS: dict[str, list[str]] = { "internlm2": ["model", "layers"], "granite": ["model", "layers"], "gemma3": ["model", "layers"], + "nemotron_h": ["backbone", "layers"], } _ATTENTION_ATTR: dict[str, str] = { @@ -89,6 +90,10 @@ _ATTENTION_ATTR: dict[str, str] = { "internlm2": "attention", "granite": "self_attn", "gemma3": "self_attn", + # Nemotron-H hybrid: attention (q/k/v/o_proj), Mamba (in/out_proj), and + # MoE experts all live under layer.mixer. Projection name matching + # picks out whichever of those actually exists on the layer. + "nemotron_h": "mixer", } _FFN_ATTR: dict[str, str] = { @@ -131,6 +136,9 @@ _FFN_ATTR: dict[str, str] = { "internlm2": "feed_forward", "granite": "mlp", "gemma3": "mlp", + # Nemotron-H: MoE experts (mixer.experts) and shared experts + # (mixer.shared_experts) hang off layer.mixer. + "nemotron_h": "mixer", }