File size: 2,467 Bytes
6e802a4 | 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 | # backbones.py: Puts the transformer (Pythia-1.4B) behind the same interface as Mamba-2, including cached stepping.
from __future__ import annotations
import torch
import torch.nn as nn
class _NeoXInner(nn.Module):
def __init__(self, neox):
super().__init__()
self.neox = neox
def __getattr__(self, name):
if name == "embedding":
return self.neox.embed_in
if name == "layers":
return self.neox.layers
return super().__getattr__(name)
def __setattr__(self, name, value):
if name == "embedding":
self._modules["neox"].embed_in = value
return
super().__setattr__(name, value)
def forward(self, input_ids: torch.Tensor, inference_params=None, **kw) -> torch.Tensor:
if inference_params is None:
return self.neox(input_ids=input_ids).last_hidden_state
past = getattr(inference_params, "hf_past", None)
if int(inference_params.seqlen_offset) == 0:
past = None
out = self.neox(input_ids=input_ids, past_key_values=past, use_cache=True)
inference_params.hf_past = out.past_key_values
return out.last_hidden_state
class TransformerBackbone(nn.Module):
def __init__(self, model_name: str, device=None, dtype=None):
super().__init__()
from transformers import GPTNeoXForCausalLM
m = GPTNeoXForCausalLM.from_pretrained(model_name, torch_dtype=dtype)
if device is not None:
m = m.to(device)
cfg = m.config
cfg.d_model = cfg.hidden_size
self.config = cfg
self.backbone = _NeoXInner(m.gpt_neox)
self.lm_head = m.embed_out
self.max_position_embeddings = getattr(cfg, "max_position_embeddings", None)
def forward(self, input_ids: torch.Tensor, **kw):
return self.lm_head(self.backbone(input_ids))
def load_backbone(kind: str, model_name: str, device=None, dtype=None):
if kind == "mamba":
from mamba_ssm import MambaLMHeadModel
return MambaLMHeadModel.from_pretrained(model_name, device=device,
dtype=dtype)
if kind == "transformer":
return TransformerBackbone(model_name, device=device, dtype=dtype)
raise ValueError(f"unknown backbone kind {kind!r} (mamba|transformer)")
DEFAULT_MODELS = {"mamba": "state-spaces/mamba2-1.3b",
"transformer": "EleutherAI/pythia-1.4b"}
|