477 lines
18 KiB
Python
477 lines
18 KiB
Python
"""HuggingFace/PyTorch shard backend for real node inference."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from dataclasses import dataclass
|
|
from typing import Any, Literal
|
|
|
|
Quantization = Literal["bfloat16", "int8", "nf4"]
|
|
|
|
|
|
class ModelBackendError(RuntimeError):
|
|
"""Base class for real model backend startup and execution failures."""
|
|
|
|
|
|
class MissingModelDependencyError(ModelBackendError):
|
|
"""Raised when optional model dependencies are not installed."""
|
|
|
|
|
|
class InsufficientVRAMError(ModelBackendError):
|
|
"""Raised when a requested shard cannot fit in available CUDA memory."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TensorPayload:
|
|
body: bytes
|
|
shape: list[int]
|
|
attention_mask_header: str | None
|
|
position_ids_header: str | None
|
|
|
|
|
|
def validate_quantization(value: str) -> Quantization:
|
|
if value not in {"bfloat16", "int8", "nf4"}:
|
|
raise ValueError("quantization must be one of: bfloat16, int8, nf4")
|
|
return value # type: ignore[return-value]
|
|
|
|
|
|
def build_quantization_config(quantization: Quantization) -> Any | None:
|
|
"""Return a transformers BitsAndBytesConfig for quantized weights."""
|
|
if quantization == "bfloat16":
|
|
return None
|
|
try:
|
|
import torch
|
|
from transformers import BitsAndBytesConfig
|
|
except ModuleNotFoundError as exc:
|
|
raise MissingModelDependencyError(
|
|
"transformers and torch are required for int8/nf4 quantization"
|
|
) from exc
|
|
|
|
if quantization == "int8":
|
|
return BitsAndBytesConfig(load_in_8bit=True)
|
|
return BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
)
|
|
|
|
|
|
class TorchModelShard:
|
|
"""Executable subset of a HuggingFace causal language model."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_id: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
quantization: Quantization = "bfloat16",
|
|
) -> None:
|
|
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
|
|
raise ValueError("shard_start must be <= shard_end and non-negative")
|
|
self.model_id = model_id
|
|
self.shard_start = shard_start
|
|
self.shard_end = shard_end
|
|
self.quantization = quantization
|
|
|
|
try:
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
except ModuleNotFoundError as exc:
|
|
raise MissingModelDependencyError(
|
|
"real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
|
|
) from exc
|
|
|
|
self.torch = torch
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
quant_config = build_quantization_config(quantization)
|
|
try:
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_id,
|
|
quantization_config=quant_config,
|
|
device_map="auto" if quant_config is not None else None,
|
|
torch_dtype=torch.bfloat16,
|
|
low_cpu_mem_usage=True,
|
|
use_safetensors=True,
|
|
)
|
|
if quant_config is None:
|
|
self.model.to(self.device)
|
|
except Exception as exc:
|
|
if _looks_like_oom(exc):
|
|
raise InsufficientVRAMError(
|
|
f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
|
|
f"with {quantization} quantization; choose a smaller shard or lower quantization"
|
|
) from exc
|
|
raise
|
|
|
|
self.model.eval()
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
self.layers = _model_layers(self.model)
|
|
self.total_layers = len(self.layers)
|
|
# shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
|
|
if shard_end >= self.total_layers:
|
|
raise ValueError(
|
|
f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}"
|
|
)
|
|
self.is_head = shard_start == 0
|
|
self.is_tail = shard_end >= self.total_layers - 1
|
|
self.hidden_size = int(
|
|
getattr(self.model.config, "hidden_size", 0)
|
|
or getattr(self.model.config, "n_embd", 0)
|
|
)
|
|
self._embed_tokens = _embed_tokens(self.model) if self.is_head else None
|
|
self._position_embeddings = _position_embeddings(self.model)
|
|
self._norm = _final_norm(self.model) if self.is_tail else None
|
|
self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None
|
|
|
|
def encode_prompt(self, prompt: str) -> TensorPayload:
|
|
if not self.is_head or self._embed_tokens is None:
|
|
raise ModelBackendError("text prompts can only be accepted by the head shard")
|
|
encoded = self.tokenizer(prompt, return_tensors="pt")
|
|
input_ids = encoded["input_ids"].to(self.device)
|
|
attention_mask = encoded.get("attention_mask")
|
|
if attention_mask is None:
|
|
attention_mask = self.torch.ones_like(input_ids)
|
|
attention_mask = attention_mask.to(self.device)
|
|
position_ids = _position_ids(attention_mask, self.torch)
|
|
hidden_states = self._embed_tokens(input_ids)
|
|
if self._position_embeddings is not None:
|
|
hidden_states = hidden_states + self._position_embeddings(position_ids)
|
|
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
|
|
return self._payload(hidden_states, attention_mask, position_ids)
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body: bytes,
|
|
shape: list[int],
|
|
attention_mask_header: str | None,
|
|
position_ids_header: str | None,
|
|
) -> TensorPayload | str:
|
|
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
|
self.device
|
|
)
|
|
attention_mask = _tensor_from_int64_header(
|
|
attention_mask_header, self.torch, self.device
|
|
)
|
|
position_ids = _tensor_from_int64_header(
|
|
position_ids_header, self.torch, self.device
|
|
)
|
|
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
|
|
if self.is_tail:
|
|
return self.decode_tail(hidden_states)
|
|
return self._payload(hidden_states, attention_mask, position_ids)
|
|
|
|
def decode_tail(self, hidden_states: Any) -> str:
|
|
if self._norm is not None:
|
|
hidden_states = self._norm(hidden_states)
|
|
if self._lm_head is None:
|
|
raise ModelBackendError("tail shard has no lm_head")
|
|
logits = self._lm_head(hidden_states)
|
|
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
|
|
return self.tokenizer.decode([token_id], skip_special_tokens=True)
|
|
|
|
def generate_text(
|
|
self,
|
|
messages: list[dict],
|
|
max_new_tokens: int = 256,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
) -> str:
|
|
"""Autoregressive generation using HF generate() — single-node (head+tail) mode."""
|
|
if not self.is_head or not self.is_tail:
|
|
raise ModelBackendError("local generation requires a full head+tail shard")
|
|
encoded = self._encode_messages(messages)
|
|
input_ids = encoded["input_ids"].to(self.device)
|
|
attention_mask = encoded.get("attention_mask")
|
|
if attention_mask is not None:
|
|
attention_mask = attention_mask.to(self.device)
|
|
pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
|
|
do_sample = temperature != 1.0 or top_p != 1.0
|
|
with self.torch.inference_mode():
|
|
generated = self.model.generate(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
max_new_tokens=max(1, int(max_new_tokens)),
|
|
do_sample=do_sample,
|
|
temperature=temperature if do_sample else None,
|
|
top_p=top_p if do_sample else None,
|
|
pad_token_id=pad_token_id,
|
|
)
|
|
new_tokens = generated[0, input_ids.shape[-1]:]
|
|
return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
|
|
|
|
def generate_text_streaming(
|
|
self,
|
|
messages: list[dict],
|
|
max_new_tokens: int = 256,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
):
|
|
"""Yield decoded token strings one at a time using HF TextIteratorStreamer."""
|
|
if not self.is_head or not self.is_tail:
|
|
raise ModelBackendError("streaming generation requires a full head+tail shard")
|
|
import threading
|
|
try:
|
|
from transformers import TextIteratorStreamer # type: ignore[import]
|
|
except ImportError:
|
|
yield self.generate_text(messages, max_new_tokens, temperature, top_p)
|
|
return
|
|
|
|
encoded = self._encode_messages(messages)
|
|
input_ids = encoded["input_ids"].to(self.device)
|
|
attention_mask = encoded.get("attention_mask")
|
|
if attention_mask is not None:
|
|
attention_mask = attention_mask.to(self.device)
|
|
pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
|
|
do_sample = temperature != 1.0 or top_p != 1.0
|
|
|
|
streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
|
|
gen_kwargs = dict(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
max_new_tokens=max(1, int(max_new_tokens)),
|
|
do_sample=do_sample,
|
|
temperature=temperature if do_sample else None,
|
|
top_p=top_p if do_sample else None,
|
|
pad_token_id=pad_token_id,
|
|
streamer=streamer,
|
|
)
|
|
t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True)
|
|
t.start()
|
|
for token_text in streamer:
|
|
yield token_text
|
|
t.join()
|
|
|
|
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
|
"""Return tokenizer-backed prompt token count for OpenAI usage metadata."""
|
|
encoded = self._encode_messages(messages)
|
|
input_ids = encoded["input_ids"]
|
|
return int(input_ids.shape[-1])
|
|
|
|
def count_text_tokens(self, text: str) -> int:
|
|
"""Return tokenizer-backed completion token count for OpenAI usage metadata."""
|
|
try:
|
|
encoded = self.tokenizer(
|
|
text,
|
|
return_tensors="pt",
|
|
add_special_tokens=False,
|
|
)
|
|
except TypeError:
|
|
encoded = self.tokenizer(text, return_tensors="pt")
|
|
return int(encoded["input_ids"].shape[-1])
|
|
|
|
def _encode_messages(self, messages: list[dict]) -> dict:
|
|
"""Format messages with chat template (if available) and tokenize."""
|
|
if hasattr(self.tokenizer, "apply_chat_template"):
|
|
try:
|
|
prompt_str = self.tokenizer.apply_chat_template(
|
|
messages,
|
|
add_generation_prompt=True,
|
|
tokenize=False,
|
|
)
|
|
return dict(self.tokenizer(prompt_str, return_tensors="pt"))
|
|
except Exception:
|
|
pass
|
|
prompt = " ".join(
|
|
str(m.get("content", ""))
|
|
for m in messages
|
|
if isinstance(m, dict) and m.get("role") == "user"
|
|
)
|
|
return dict(self.tokenizer(prompt, return_tensors="pt"))
|
|
|
|
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
|
position_embeddings = _rotary_position_embeddings(
|
|
self.model,
|
|
hidden_states,
|
|
position_ids,
|
|
)
|
|
layer_attention_mask = _decoder_attention_mask(
|
|
attention_mask,
|
|
hidden_states,
|
|
self.torch,
|
|
)
|
|
with self.torch.inference_mode():
|
|
for layer in self.layers[self.shard_start:self.shard_end + 1]:
|
|
hidden_states = _call_layer(
|
|
layer,
|
|
hidden_states,
|
|
layer_attention_mask,
|
|
position_ids,
|
|
position_embeddings,
|
|
)
|
|
return hidden_states.to(self.torch.bfloat16)
|
|
|
|
def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
|
|
hidden_states = hidden_states.to(self.torch.bfloat16).contiguous()
|
|
return TensorPayload(
|
|
body=_tensor_to_bytes(hidden_states),
|
|
shape=list(hidden_states.shape),
|
|
attention_mask_header=_int_tensor_header(attention_mask)
|
|
if attention_mask is not None
|
|
else None,
|
|
position_ids_header=_int_tensor_header(position_ids)
|
|
if position_ids is not None
|
|
else None,
|
|
)
|
|
|
|
|
|
def load_torch_shard(
|
|
model_id: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
quantization: Quantization = "bfloat16",
|
|
) -> TorchModelShard:
|
|
return TorchModelShard(model_id, shard_start, shard_end, quantization)
|
|
|
|
|
|
def _model_layers(model: Any) -> Any:
|
|
if hasattr(model, "model") and hasattr(model.model, "layers"):
|
|
return model.model.layers
|
|
if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
|
|
return model.transformer.h
|
|
raise ModelBackendError(
|
|
"unsupported HuggingFace model architecture: no transformer layers found"
|
|
)
|
|
|
|
|
|
def _embed_tokens(model: Any) -> Any:
|
|
if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
|
|
return model.model.embed_tokens
|
|
if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
|
|
return model.transformer.wte
|
|
raise ModelBackendError(
|
|
"unsupported HuggingFace model architecture: no token embeddings found"
|
|
)
|
|
|
|
|
|
def _position_embeddings(model: Any) -> Any | None:
|
|
if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
|
|
return model.transformer.wpe
|
|
return None
|
|
|
|
|
|
def _final_norm(model: Any) -> Any | None:
|
|
if hasattr(model, "model") and hasattr(model.model, "norm"):
|
|
return model.model.norm
|
|
if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
|
|
return model.transformer.ln_f
|
|
return None
|
|
|
|
|
|
def _position_ids(attention_mask: Any, torch: Any) -> Any:
|
|
position_ids = attention_mask.long().cumsum(-1) - 1
|
|
return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)
|
|
|
|
|
|
def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any:
|
|
"""Build a causal additive mask for decoder layers called outside model.forward."""
|
|
if attention_mask is None:
|
|
return None
|
|
if len(getattr(attention_mask, "shape", ())) != 2:
|
|
return attention_mask
|
|
batch_size, seq_len = attention_mask.shape
|
|
if seq_len <= 1:
|
|
return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype)
|
|
|
|
min_value = torch.finfo(hidden_states.dtype).min
|
|
causal = torch.full(
|
|
(seq_len, seq_len),
|
|
min_value,
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
causal = torch.triu(causal, diagonal=1)
|
|
causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()
|
|
|
|
padding = attention_mask.to(device=hidden_states.device)
|
|
if not bool(padding.all()):
|
|
causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value)
|
|
return causal
|
|
|
|
|
|
def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None:
|
|
"""Return model-level rotary embeddings required by newer HF decoder layers."""
|
|
if position_ids is None:
|
|
return None
|
|
rotary = None
|
|
if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
|
|
rotary = model.model.rotary_emb
|
|
elif hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
|
|
rotary = model.transformer.rotary_emb
|
|
if rotary is None:
|
|
return None
|
|
return rotary(hidden_states, position_ids)
|
|
|
|
|
|
def _call_layer(
|
|
layer: Any,
|
|
hidden_states: Any,
|
|
attention_mask: Any,
|
|
position_ids: Any,
|
|
position_embeddings: Any | None = None,
|
|
) -> Any:
|
|
attempts = (
|
|
{
|
|
"attention_mask": attention_mask,
|
|
"position_ids": position_ids,
|
|
"position_embeddings": position_embeddings,
|
|
"use_cache": False,
|
|
},
|
|
{
|
|
"attention_mask": attention_mask,
|
|
"position_ids": position_ids,
|
|
"use_cache": False,
|
|
},
|
|
{"attention_mask": attention_mask, "use_cache": False},
|
|
{"use_cache": False},
|
|
{},
|
|
)
|
|
last_exc: Exception | None = None
|
|
for kwargs in attempts:
|
|
filtered = {key: value for key, value in kwargs.items() if value is not None}
|
|
try:
|
|
output = layer(hidden_states, **filtered)
|
|
return output[0] if isinstance(output, tuple) else output
|
|
except TypeError as exc:
|
|
last_exc = exc
|
|
if last_exc is not None:
|
|
raise last_exc
|
|
return layer(hidden_states)[0]
|
|
|
|
|
|
def _tensor_to_bytes(tensor: Any) -> bytes:
|
|
import torch
|
|
|
|
return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()
|
|
|
|
|
|
def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
|
|
tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16)
|
|
return tensor.reshape(shape)
|
|
|
|
|
|
def _int_tensor_header(tensor: Any) -> str:
|
|
data = tensor.detach().cpu().long().contiguous()
|
|
raw = data.numpy().tobytes()
|
|
shape = ",".join(str(dim) for dim in data.shape)
|
|
encoded = base64.b64encode(raw).decode("ascii")
|
|
return f"{shape}:{encoded}"
|
|
|
|
|
|
def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None:
|
|
if not value:
|
|
return None
|
|
shape_text, encoded = value.split(":", 1)
|
|
shape = [int(part) for part in shape_text.split(",") if part]
|
|
raw = base64.b64decode(encoded.encode("ascii"))
|
|
return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device)
|
|
|
|
|
|
def _looks_like_oom(exc: BaseException) -> bool:
|
|
current: BaseException | None = exc
|
|
while current is not None:
|
|
text = str(current).lower()
|
|
if "out of memory" in text or "cuda error: out of memory" in text:
|
|
return True
|
|
current = current.__cause__ or current.__context__
|
|
return False
|