| 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 |
|
|
