File size: 3,174 Bytes
906715b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """Model binding: where do the adapters attach?
Everything architecture-specific lives in binding/ and only here (the
substrate-shim discipline from the research line). Resolution order:
explicit binding object -> dotted path string -> registry by
model_type/class -> error listing candidate ModuleLists (never guess).
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable, Protocol, Sequence
import torch.nn as nn
class ModelBinding(Protocol):
name: str
def layers(self, model) -> Sequence[nn.Module]: ...
def set_layers(self, model, new: list[nn.Module]) -> None: ...
def hidden_size(self, model) -> int: ...
@dataclass
class PathBinding:
"""Generic binding around a dotted attribute path to a ModuleList."""
path: str
name: str = "path"
def _parent(self, model):
obj = model
parts = self.path.split(".")
for p in parts[:-1]:
obj = getattr(obj, p)
return obj, parts[-1]
def layers(self, model):
parent, leaf = self._parent(model)
return getattr(parent, leaf)
def set_layers(self, model, new):
parent, leaf = self._parent(model)
setattr(parent, leaf, nn.ModuleList(new))
def hidden_size(self, model) -> int:
cfg = model.config
for attr in ("text_config", None):
c = getattr(cfg, attr, cfg) if attr else cfg
if hasattr(c, "hidden_size"):
return int(c.hidden_size)
raise AttributeError("no hidden_size on model.config")
REGISTRY: dict[str, Callable[[], ModelBinding]] = {}
def register(key: str):
def deco(fn):
REGISTRY[key] = fn
return fn
return deco
@register("qwen3_5_vl")
def _qwen35() -> ModelBinding:
# dual-tower VLM: LLM tower at model.model.language_model.layers
return PathBinding("model.language_model.layers", name="qwen3_5_vl")
@register("generic_causal")
def _generic() -> ModelBinding:
# Llama/Qwen3/Mistral-style: model.model.layers
return PathBinding("model.layers", name="generic_causal")
def _candidates(model) -> list[str]:
out = []
for name, mod in model.named_modules():
if isinstance(mod, nn.ModuleList) and len(mod) >= 4:
out.append(f"{name} (len {len(mod)})")
return out
def resolve(model, binding: "ModelBinding | str | None" = None) -> ModelBinding:
if binding is None:
mt = getattr(model.config, "model_type", "")
if "qwen3_5" in mt and hasattr(model, "model") and \
hasattr(model.model, "language_model"):
return REGISTRY["qwen3_5_vl"]()
if hasattr(model, "model") and hasattr(model.model, "layers"):
return REGISTRY["generic_causal"]()
raise ValueError(
"amoe could not resolve an attach site. Pass "
"binding=<dotted path to the decoder ModuleList>. "
f"Candidates found: {_candidates(model)}")
if isinstance(binding, str):
if binding in REGISTRY:
return REGISTRY[binding]()
# dotted path relative to the model object
return PathBinding(binding)
return binding
|