|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +"""Hugging Face Transformers prefill forward used by the Rust crate.""" |
| 5 | + |
| 6 | +from __future__ import annotations |
| 7 | + |
| 8 | +from typing import Any |
| 9 | + |
| 10 | + |
| 11 | +def _detect_device(torch: Any) -> str: |
| 12 | + if torch.cuda.is_available(): |
| 13 | + return "cuda" |
| 14 | + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): |
| 15 | + return "mps" |
| 16 | + return "cpu" |
| 17 | + |
| 18 | + |
| 19 | +def _resolve_device(torch: Any, override: str | None) -> str: |
| 20 | + device = override.lower().strip() if override is not None else _detect_device(torch) |
| 21 | + if device == "cpu": |
| 22 | + return device |
| 23 | + if device == "mps": |
| 24 | + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): |
| 25 | + return device |
| 26 | + raise ValueError("MPS was requested but is not available") |
| 27 | + if device == "cuda": |
| 28 | + if torch.cuda.is_available(): |
| 29 | + return device |
| 30 | + raise ValueError("CUDA was requested but is not available") |
| 31 | + if device.startswith("cuda:") and device.removeprefix("cuda:").isdigit(): |
| 32 | + index = int(device.removeprefix("cuda:")) |
| 33 | + if torch.cuda.is_available() and index < torch.cuda.device_count(): |
| 34 | + return device |
| 35 | + raise ValueError(f"CUDA device {device} is not available") |
| 36 | + raise ValueError(f"Unsupported device: {device}") |
| 37 | + |
| 38 | + |
| 39 | +class TransformersForward: |
| 40 | + """Lazily load a causal LM and return pooled prefill hidden states.""" |
| 41 | + |
| 42 | + def __init__( |
| 43 | + self, |
| 44 | + model: str, |
| 45 | + *, |
| 46 | + device: str | None = None, |
| 47 | + cache_dir: str | None = None, |
| 48 | + ) -> None: |
| 49 | + self._model_path = model |
| 50 | + self._cache_dir = cache_dir |
| 51 | + self._device_override = device |
| 52 | + self._model = None |
| 53 | + self._tokenizer = None |
| 54 | + self._torch = None |
| 55 | + self.n_layers = 0 |
| 56 | + self.hidden_dim = 0 |
| 57 | + |
| 58 | + def _ensure_loaded(self) -> str: |
| 59 | + if self._model is not None: |
| 60 | + return str(self._model.device) |
| 61 | + |
| 62 | + import torch |
| 63 | + from transformers import AutoModelForCausalLM, AutoTokenizer |
| 64 | + |
| 65 | + self._torch = torch |
| 66 | + device = _resolve_device(torch, self._device_override) |
| 67 | + if device == "cpu": |
| 68 | + dtype = torch.float32 |
| 69 | + elif device == "mps": |
| 70 | + dtype = torch.float16 |
| 71 | + else: |
| 72 | + dtype = ( |
| 73 | + torch.bfloat16 |
| 74 | + if torch.cuda.get_device_capability(device)[0] >= 8 |
| 75 | + else torch.float16 |
| 76 | + ) |
| 77 | + |
| 78 | + self._tokenizer = AutoTokenizer.from_pretrained( |
| 79 | + self._model_path, |
| 80 | + cache_dir=self._cache_dir, |
| 81 | + ) |
| 82 | + if self._tokenizer.pad_token is None: |
| 83 | + self._tokenizer.pad_token = self._tokenizer.eos_token |
| 84 | + |
| 85 | + load_kwargs: dict[str, Any] = { |
| 86 | + "dtype": dtype, |
| 87 | + "cache_dir": self._cache_dir, |
| 88 | + } |
| 89 | + |
| 90 | + self._model = AutoModelForCausalLM.from_pretrained( |
| 91 | + self._model_path, |
| 92 | + **load_kwargs, |
| 93 | + ) |
| 94 | + if device != "cpu": |
| 95 | + self._model.to(device) |
| 96 | + self._model.eval() |
| 97 | + self.n_layers = self._model.config.num_hidden_layers |
| 98 | + self.hidden_dim = self._model.config.hidden_size |
| 99 | + return str(self._model.device) |
| 100 | + |
| 101 | + def extract_batch( |
| 102 | + self, |
| 103 | + prompts: list[str], |
| 104 | + *, |
| 105 | + chat_template_kwargs: dict[str, Any] | None = None, |
| 106 | + extract_layers: list[int] | str = "upper_half", |
| 107 | + pooling_modes: list[str] | None = None, |
| 108 | + batch_size: int = 4, |
| 109 | + max_length: int = 2048, |
| 110 | + ) -> dict[str, Any]: |
| 111 | + """Extract pooled hidden states using the blueprint's direct indexing.""" |
| 112 | + self._ensure_loaded() |
| 113 | + |
| 114 | + if extract_layers == "all": |
| 115 | + layers = list(range(self.n_layers)) |
| 116 | + elif extract_layers == "upper_half": |
| 117 | + layers = list(range(self.n_layers // 2, self.n_layers)) |
| 118 | + elif isinstance(extract_layers, list): |
| 119 | + layers = [int(layer) for layer in extract_layers] |
| 120 | + else: |
| 121 | + raise ValueError(f"Unsupported layer selection: {extract_layers}") |
| 122 | + if not layers: |
| 123 | + raise ValueError("extract_layers resolved to an empty list") |
| 124 | + invalid = [layer for layer in layers if layer < 0 or layer >= self.n_layers] |
| 125 | + if invalid: |
| 126 | + raise ValueError( |
| 127 | + f"Requested indexes {invalid} are outside hidden-state range 0..{self.n_layers - 1}" |
| 128 | + ) |
| 129 | + |
| 130 | + pools = {"last", "mean"} if pooling_modes is None else set(pooling_modes) |
| 131 | + unknown_pools = pools - {"last", "mean"} |
| 132 | + if unknown_pools: |
| 133 | + raise ValueError(f"Unknown pooling modes: {sorted(unknown_pools)}") |
| 134 | + if not pools: |
| 135 | + raise ValueError("At least one pooling mode is required") |
| 136 | + |
| 137 | + template_kwargs = chat_template_kwargs or {} |
| 138 | + conversations = [[{"role": "user", "content": prompt}] for prompt in prompts] |
| 139 | + formatted = self._tokenizer.apply_chat_template( |
| 140 | + conversations, |
| 141 | + tokenize=False, |
| 142 | + add_generation_prompt=True, |
| 143 | + **template_kwargs, |
| 144 | + ) |
| 145 | + all_last = {layer: [] for layer in layers} if "last" in pools else {} |
| 146 | + all_mean = {layer: [] for layer in layers} if "mean" in pools else {} |
| 147 | + |
| 148 | + for batch_start in range(0, len(formatted), batch_size): |
| 149 | + inputs = self._tokenizer( |
| 150 | + formatted[batch_start : batch_start + batch_size], |
| 151 | + return_tensors="pt", |
| 152 | + padding=True, |
| 153 | + truncation=True, |
| 154 | + max_length=max_length, |
| 155 | + ) |
| 156 | + input_ids = inputs["input_ids"].to(self._model.device) |
| 157 | + attention_mask = inputs["attention_mask"].to(self._model.device) |
| 158 | + |
| 159 | + with self._torch.inference_mode(): |
| 160 | + outputs = self._model( |
| 161 | + input_ids=input_ids, |
| 162 | + attention_mask=attention_mask, |
| 163 | + output_hidden_states=True, |
| 164 | + use_cache=False, |
| 165 | + ) |
| 166 | + |
| 167 | + hidden_states = outputs.hidden_states |
| 168 | + token_mask = attention_mask.bool() |
| 169 | + token_count = token_mask.sum(dim=1, keepdim=True) |
| 170 | + positions = self._torch.arange(token_mask.shape[1], device=token_mask.device).expand_as( |
| 171 | + token_mask |
| 172 | + ) |
| 173 | + last_index = positions.masked_fill(~token_mask, -1).max(dim=1).values |
| 174 | + batch_index = self._torch.arange(token_mask.shape[0], device=token_mask.device) |
| 175 | + |
| 176 | + for layer in layers: |
| 177 | + hidden = hidden_states[layer].float() |
| 178 | + if "last" in pools: |
| 179 | + all_last[layer].append(hidden[batch_index, last_index].cpu()) |
| 180 | + if "mean" in pools: |
| 181 | + masked = hidden.masked_fill(~token_mask.unsqueeze(-1), 0) |
| 182 | + all_mean[layer].append((masked.sum(dim=1) / token_count).cpu()) |
| 183 | + |
| 184 | + del outputs, hidden_states, input_ids, attention_mask |
| 185 | + |
| 186 | + return { |
| 187 | + "hidden_last": { |
| 188 | + layer: self._torch.cat(rows).contiguous().numpy().tobytes() |
| 189 | + for layer, rows in all_last.items() |
| 190 | + }, |
| 191 | + "hidden_mean": { |
| 192 | + layer: self._torch.cat(rows).contiguous().numpy().tobytes() |
| 193 | + for layer, rows in all_mean.items() |
| 194 | + }, |
| 195 | + "n_layers": self.n_layers, |
| 196 | + "hidden_dim": self.hidden_dim, |
| 197 | + } |
| 198 | + |
| 199 | + def unload(self) -> None: |
| 200 | + self._model = None |
| 201 | + self._tokenizer = None |
| 202 | + self.n_layers = 0 |
| 203 | + self.hidden_dim = 0 |
| 204 | + if self._torch is not None: |
| 205 | + if self._torch.cuda.is_available(): |
| 206 | + self._torch.cuda.empty_cache() |
| 207 | + if hasattr(self._torch, "mps") and self._torch.backends.mps.is_available(): |
| 208 | + self._torch.mps.empty_cache() |
0 commit comments