mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
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.
This commit is contained in:
Brian Bell
@@ -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()
|
||||
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user