"""Architecture adapters — one place for multimodal-wrapper unwrapping. Some HF checkpoints (Gemma-3, LLaVA-style) load as multimodal wrappers even via `AutoModelForCausalLM`. NLA only cares about the text side. Rather than scattering `getattr(x, "text_config", x)` / `getattr(x, "language_model", x)` across train_actor/models/extractors, these functions centralize the unwrap with explicit arch detection. Pass-through for plain text models (Qwen/Llama/Mistral): they have none of the wrapper attributes, so all functions return the input unchanged. Add new wrapped architectures by extending _WRAPPER_MODEL_ATTRS / _WRAPPER_CONFIG_ATTRS — don't duck-type new getattr fallbacks at callsites. """ from typing import Any import torch from transformers import AutoModelForCausalLM # Multimodal wrapper → text model attribute name. # Gemma3ForConditionalGeneration.language_model → Gemma3ForCausalLM # (LLaVA-style wrappers would go here too if we ever support them) _WRAPPER_MODEL_ATTRS = ("language_model",) # Multimodal wrapper config → text config attribute name. # Gemma3Config.text_config → Gemma3TextConfig (has hidden_size, num_hidden_layers) _WRAPPER_CONFIG_ATTRS = ("text_config",) def resolve_text_config(config: Any) -> Any: """Return the text-side config for multimodal wrappers; pass-through otherwise. Gemma3Config nests hidden_size/num_hidden_layers under .text_config. Qwen2Config/LlamaConfig/MistralConfig have those at top level and no .text_config attr → return as-is. """ for attr in _WRAPPER_CONFIG_ATTRS: nested = getattr(config, attr, None) if nested is not None: return nested return config def resolve_text_model(model: Any) -> Any: """Return the text-side CausalLM for multimodal wrappers; pass-through otherwise. Invariant: always returns a CausalLM-shaped model (has .model + .lm_head), so `save_pretrained()` / `AutoModelForCausalLM.from_pretrained()` roundtrip. Gemma3ForConditionalGeneration.language_model is a Gemma3TextModel (bare transformer, NO .lm_head, NO .model wrapper). Returning it directly means `save_pretrained` writes keys like `layers.0.*` but `from_pretrained` via AutoModelForCausalLM loads Gemma3ForCausalLM expecting `model.layers.0.*` → zero keys match → everything random-inits. Observed Mar 13 2026: pred_norm=507 vs gold_norm=75616 on same input, step-0 loss=2.0 (orthogonal). Qwen/Llama/Mistral have no .language_model → pass through unchanged (already CausalLM-shaped). """ nested = _find_text_submodel(model) if nested is None: return model if hasattr(nested, "lm_head"): return nested # already CausalLM-shaped # Bare TextModel — wrap in CausalLM so keys roundtrip. meta device # avoids materializing a throwaway random model. with torch.device("meta"): wrapper = AutoModelForCausalLM.from_config(nested.config) wrapper.model = nested # transplant pretrained weights # lm_head is on meta. Critic caller strips it to Identity (harmless). # Actor caller (train_actor.NLATextOnlyCausalLM) needs a real one to # generate — Gemma ties it to embed_tokens, so tie_weights() points # lm_head.weight at the real embedding tensor (no extra alloc). # Non-tied archs would need the caller to load lm_head separately; # cross that bridge when we hit one. if getattr(nested.config, "tie_word_embeddings", False): wrapper.tie_weights() return wrapper def _find_text_submodel(model: Any) -> Any: """Locate the bare text transformer (`*TextModel`) inside a multimodal wrapper. Layout drifts across transformers versions / archs: - Gemma-3 (older tf): ForConditionalGeneration.language_model - Gemma-4 (tf>=5.5): ForConditionalGeneration.model.language_model (`.model` is the multimodal `*Model`; `.language_model` is its sibling of `.vision_tower`). Plain causal LMs (Qwen/Llama/Gemma*ForCausalLM) have no `.language_model` at either level → return None (already CausalLM-shaped). """ lm = getattr(model, "language_model", None) if lm is not None: return lm inner = getattr(model, "model", None) if inner is not None and inner is not model: lm = getattr(inner, "language_model", None) if lm is not None: return lm return None def find_decoder_layers_owner(model: Any) -> Any: """Return the submodule that directly owns the decoder `.layers` ModuleList. Descends through PEFT wrappers (PeftModel/LoraModel), multimodal wrappers (Gemma-3/4: `*ForConditionalGeneration.model.language_model`), and causal-LM wrappers (`.model`). Prefers the text path (`.language_model` before `.model`) so it never lands on the vision tower's own `.layers`. """ visited: set[int] = set() cur = model for _ in range(12): if id(cur) in visited: break visited.add(id(cur)) layers = getattr(cur, "layers", None) if isinstance(layers, torch.nn.ModuleList) and len(layers) > 0: return cur nxt = None for attr in ("language_model", "model", "transformer", "base_model"): child = getattr(cur, attr, None) if isinstance(child, torch.nn.Module) and child is not cur and id(child) not in visited: nxt = child break if nxt is None: break cur = nxt raise AssertionError( f"find_decoder_layers_owner: could not locate decoder .layers under " f"{type(model).__name__} — extend the descent attr list." ) def resolve_decoder_layers(model: Any) -> torch.nn.ModuleList: """Find the decoder layers ModuleList, descending through PEFT / multimodal / causal-LM wrappers (handles Gemma-4's nested .model.language_model.layers and PEFT-wrapped models alike — see find_decoder_layers_owner).""" layers = find_decoder_layers_owner(model).layers assert isinstance(layers, torch.nn.ModuleList), ( f"resolved {type(layers).__name__}, expected nn.ModuleList." ) return layers def is_multimodal_wrapper(config_or_model: Any) -> bool: """True if this is a known multimodal wrapper (has nested text config/model).""" for attr in (*_WRAPPER_CONFIG_ATTRS, *_WRAPPER_MODEL_ATTRS): if getattr(config_or_model, attr, None) is not None: return True return False # ───────────────────────────────────────────────────────────────────────── # Embedding forward scale — for RL rollout injection. # # Some archs' embedding modules multiply by a scale in their forward(): # - Gemma: Gemma3TextScaledWordEmbedding × √d_model # - T5-style: similar √d_model scaling # Plain nn.Embedding (Qwen/Llama/Mistral): no scale. # # Training-side injection is unaffected (the forward hook captures POST-scale # output). RL rollout builds embeds manually from raw weights → must apply # the scale or injection magnitude is off by ~62× (Gemma). # # This is an EXPLICIT registry — don't try to auto-detect. If you add a # scaled-embed arch, add it here AND the callsite will visibly multiply. # ───────────────────────────────────────────────────────────────────────── # model_type → scale expression. √d_model is the common case; if an arch # uses a different formula, add a lambda here. _SCALED_EMBED_MODEL_TYPES: dict[str, str] = { "gemma4": "sqrt_d_model", "gemma4_text": "sqrt_d_model", "gemma3": "sqrt_d_model", "gemma3_text": "sqrt_d_model", "gemma2": "sqrt_d_model", "gemma": "sqrt_d_model", "t5": "sqrt_d_model", } def resolve_embed_scale(config: Any) -> float: """Return the scalar that the model's embedding forward() multiplies by. For plain nn.Embedding archs (Qwen/Llama/Mistral): 1.0. For Gemma/T5-style: √d_model. RL rollout uses this to scale raw-weight-lookup embeds to match what the model's forward would produce. Training doesn't need it (forward hook captures post-scale output). See `_SCALED_EMBED_MODEL_TYPES` and extend there for new scaled-embed archs. """ text_config = resolve_text_config(config) model_type = getattr(text_config, "model_type", "") rule = _SCALED_EMBED_MODEL_TYPES.get(model_type) if rule is None: return 1.0 if rule == "sqrt_d_model": return float(text_config.hidden_size) ** 0.5 raise AssertionError( f"unknown embed-scale rule {rule!r} for model_type={model_type!r} — " f"extend resolve_embed_scale in arch_adapters.py" )