File size: 6,561 Bytes
88e15cd | 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | from __future__ import annotations
import json
from contextlib import nullcontext
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
from torch import nn
from transformers import AutoModel, AutoTokenizer
from .config import ModelConfig
def resolve_torch_dtype(name: str, device: torch.device) -> torch.dtype:
if name == "float32":
return torch.float32
if name == "float16":
return torch.float16
if name == "bfloat16":
return torch.bfloat16
if device.type != "cuda":
return torch.float32
if torch.cuda.is_bf16_supported():
return torch.bfloat16
return torch.float16
class RouterModel(nn.Module):
"""Frozen language-model encoder plus a small worker-selection head."""
def __init__(
self,
config: ModelConfig,
worker_ids: list[str],
device: torch.device,
backbone_source: str | Path | None = None,
):
super().__init__()
self.router_config = config
self.worker_ids = list(worker_ids)
self.device_ref = device
dtype = resolve_torch_dtype(config.dtype, device)
source = str(backbone_source or config.base_model)
self.tokenizer = AutoTokenizer.from_pretrained(source, use_fast=True)
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = "right"
self.backbone = AutoModel.from_pretrained(source, torch_dtype=dtype)
hidden_size = getattr(self.backbone.config, "hidden_size", None)
if hidden_size is None:
hidden_size = getattr(self.backbone.config, "d_model", None)
if hidden_size is None:
raise ValueError("Could not infer backbone hidden size")
self.norm = nn.LayerNorm(hidden_size, dtype=torch.float32)
self.dropout = nn.Dropout(config.dropout)
self.classifier = nn.Linear(hidden_size, len(worker_ids), dtype=torch.float32)
self.backbone.requires_grad_(not config.freeze_backbone)
self.to(device)
if config.freeze_backbone:
self.backbone.eval()
def train(self, mode: bool = True):
super().train(mode)
if self.router_config.freeze_backbone:
self.backbone.eval()
return self
def _pool(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
if self.router_config.pooling == "mean":
mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
return (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1)
positions = torch.arange(hidden.shape[1], device=hidden.device).unsqueeze(0)
last_indices = positions.masked_fill(attention_mask == 0, -1).max(dim=1).values
batch_indices = torch.arange(hidden.shape[0], device=hidden.device)
return hidden[batch_indices, last_indices]
def encode_features(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
context = torch.no_grad() if self.router_config.freeze_backbone else nullcontext()
with context:
output = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
pooled = self._pool(output.last_hidden_state, attention_mask)
return self.norm(pooled.float())
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
features = self.encode_features(input_ids, attention_mask)
return self.classifier(self.dropout(features))
def head_state_dict(self) -> dict[str, torch.Tensor]:
state = {}
for prefix, module in (("norm", self.norm), ("classifier", self.classifier)):
for name, value in module.state_dict().items():
state[f"{prefix}.{name}"] = value.detach().cpu().contiguous()
return state
def load_head_state_dict(self, state: dict[str, torch.Tensor]) -> None:
self.norm.load_state_dict(
{key.removeprefix("norm."): value for key, value in state.items() if key.startswith("norm.")}
)
self.classifier.load_state_dict(
{
key.removeprefix("classifier."): value
for key, value in state.items()
if key.startswith("classifier.")
}
)
def save_checkpoint(self, path: str | Path, metadata: dict | None = None) -> None:
destination = Path(path)
destination.mkdir(parents=True, exist_ok=True)
payload = {
"format_version": 1,
"model": self.router_config.model_dump(),
"worker_ids": self.worker_ids,
"metadata": metadata or {},
}
(destination / "router_config.json").write_text(
json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
)
save_file(self.head_state_dict(), destination / "router_head.safetensors")
self.tokenizer.save_pretrained(destination / "tokenizer")
if not self.router_config.freeze_backbone:
self.backbone.save_pretrained(destination / "backbone", safe_serialization=True)
@classmethod
def from_checkpoint(
cls, path: str | Path, device: torch.device | None = None
) -> RouterModel:
source = Path(path)
payload = json.loads((source / "router_config.json").read_text(encoding="utf-8"))
config = ModelConfig.model_validate(payload["model"])
target_device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
backbone_path = source / "backbone"
model = cls(
config,
payload["worker_ids"],
target_device,
backbone_source=backbone_path if backbone_path.exists() else None,
)
tokenizer_path = source / "tokenizer"
if tokenizer_path.exists():
model.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
if model.tokenizer.pad_token_id is None:
model.tokenizer.pad_token = model.tokenizer.eos_token
model.tokenizer.padding_side = "right"
state = load_file(source / "router_head.safetensors", device="cpu")
model.load_head_state_dict(state)
return model
def trainable_parameter_counts(self) -> tuple[int, int]:
trainable = sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
total = sum(parameter.numel() for parameter in self.parameters())
return trainable, total
|