560de08edd
Adds a committed .gitattributes so Windows and Linux checkouts converge on LF for all text files, overriding each developer's local core.autocrlf. Renormalizes existing blobs (server.py, dashboard.html, etc.) that had CRLF baked in, clearing the repo-wide phantom "modified" churn. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
828 lines
30 KiB
Python
828 lines
30 KiB
Python
"""HuggingFace/PyTorch shard backend for real node inference."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from dataclasses import dataclass
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Literal
|
|
|
|
Quantization = Literal["auto", "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."""
|
|
|
|
|
|
class PartialModelLoadUnsupported(ModelBackendError):
|
|
"""Raised when a shard cannot be materialized from a local snapshot subset."""
|
|
|
|
|
|
@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 {"auto", "bfloat16", "int8", "nf4"}:
|
|
raise ValueError("quantization must be one of: auto, 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 in {"auto", "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 = "auto",
|
|
cache_dir: Path | None = None,
|
|
) -> 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 AutoConfig, 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")
|
|
load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
|
|
quant_config, dtype, uses_quantized_weights = _model_load_plan(
|
|
AutoConfig,
|
|
load_source,
|
|
quantization,
|
|
torch,
|
|
None if load_source != model_id else cache_dir,
|
|
)
|
|
try:
|
|
total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source)
|
|
if _should_partial_materialize_shard(
|
|
load_source,
|
|
shard_start,
|
|
shard_end,
|
|
total_layers_hint=total_layers_hint,
|
|
uses_quantized_weights=uses_quantized_weights,
|
|
):
|
|
self.model = _load_partial_model_from_snapshot(
|
|
AutoConfig,
|
|
AutoModelForCausalLM,
|
|
torch,
|
|
load_source,
|
|
shard_start,
|
|
shard_end,
|
|
dtype,
|
|
self.device,
|
|
)
|
|
else:
|
|
load_kwargs = {
|
|
"device_map": "auto" if uses_quantized_weights else None,
|
|
"dtype": dtype,
|
|
"low_cpu_mem_usage": True,
|
|
"cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
|
|
}
|
|
if quant_config is not None:
|
|
load_kwargs["quantization_config"] = quant_config
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
load_source,
|
|
**load_kwargs,
|
|
)
|
|
if not uses_quantized_weights:
|
|
self.model.to(self.device)
|
|
except Exception as exc:
|
|
if _looks_like_oom(exc):
|
|
memory_kind = "VRAM" if self.device.type == "cuda" else "RAM"
|
|
raise InsufficientVRAMError(
|
|
f"insufficient {memory_kind} 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(
|
|
load_source,
|
|
cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None,
|
|
)
|
|
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,
|
|
start_layer: int | None = 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, start_layer=start_layer
|
|
)
|
|
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 = 5120,
|
|
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 = 5000,
|
|
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,
|
|
start_layer: int | None = None,
|
|
) -> Any:
|
|
# start_layer overrides shard_start for overlapping-shard routing
|
|
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
|
|
# indexing outside the loaded weights.
|
|
effective_start = (
|
|
max(self.shard_start, start_layer)
|
|
if start_layer is not None
|
|
else self.shard_start
|
|
)
|
|
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[effective_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 = "auto",
|
|
cache_dir: Path | None = None,
|
|
) -> TorchModelShard:
|
|
return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
|
|
|
|
|
|
def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None:
|
|
snapshot_dir = Path(load_source)
|
|
if not (snapshot_dir / "config.json").exists():
|
|
return None
|
|
from .model_catalog import layers_from_config
|
|
|
|
try:
|
|
cfg = auto_config.from_pretrained(str(snapshot_dir))
|
|
except Exception:
|
|
return None
|
|
return layers_from_config(cfg)
|
|
|
|
|
|
def _should_partial_materialize_shard(
|
|
load_source: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
*,
|
|
total_layers_hint: int | None,
|
|
uses_quantized_weights: bool,
|
|
) -> bool:
|
|
if uses_quantized_weights:
|
|
return False
|
|
snapshot_dir = Path(load_source)
|
|
if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists():
|
|
return False
|
|
if not (snapshot_dir / "model.safetensors.index.json").exists():
|
|
return False
|
|
if total_layers_hint is None:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _load_partial_model_from_snapshot(
|
|
auto_config: Any,
|
|
auto_model_for_causal_lm: Any,
|
|
torch: Any,
|
|
load_source: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
dtype: Any,
|
|
device: Any,
|
|
*,
|
|
init_empty_weights_fn: Any | None = None,
|
|
set_tensor_fn: Any | None = None,
|
|
safe_open_fn: Any | None = None,
|
|
) -> Any:
|
|
from .model_catalog import layers_from_config
|
|
from .safetensors_selection import (
|
|
INDEX_FILENAME,
|
|
select_tensor_names_for_layers_from_index,
|
|
)
|
|
|
|
if init_empty_weights_fn is None:
|
|
from accelerate import init_empty_weights as init_empty_weights_fn
|
|
if set_tensor_fn is None:
|
|
from accelerate.utils import set_module_tensor_to_device as set_tensor_fn
|
|
if safe_open_fn is None:
|
|
from safetensors import safe_open as safe_open_fn
|
|
|
|
snapshot_dir = Path(load_source)
|
|
cfg = auto_config.from_pretrained(str(snapshot_dir))
|
|
total_layers = layers_from_config(cfg)
|
|
if total_layers is None:
|
|
raise PartialModelLoadUnsupported(
|
|
f"could not determine num_hidden_layers for local snapshot {snapshot_dir}"
|
|
)
|
|
if shard_end >= total_layers:
|
|
raise ValueError(
|
|
f"shard_end {shard_end} exceeds last layer index {total_layers - 1}"
|
|
)
|
|
|
|
index_path = snapshot_dir / INDEX_FILENAME
|
|
try:
|
|
index = json.loads(index_path.read_text(encoding="utf-8"))
|
|
except FileNotFoundError as exc:
|
|
raise PartialModelLoadUnsupported(
|
|
f"missing SafeTensors index for partial load: {index_path}"
|
|
) from exc
|
|
weight_map = index.get("weight_map")
|
|
if not isinstance(weight_map, dict):
|
|
raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object")
|
|
|
|
tensor_names = select_tensor_names_for_layers_from_index(
|
|
weight_map,
|
|
shard_start,
|
|
shard_end,
|
|
total_layers=total_layers,
|
|
)
|
|
if not tensor_names:
|
|
raise PartialModelLoadUnsupported(
|
|
f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}"
|
|
)
|
|
|
|
with init_empty_weights_fn():
|
|
model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype)
|
|
tie_weights = getattr(model, "tie_weights", None)
|
|
if callable(tie_weights):
|
|
tie_weights()
|
|
|
|
# Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and
|
|
# multi-token-prediction tensors the text-only CausalLM never builds;
|
|
# transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected,
|
|
# so the manual loader must skip them too.
|
|
expected_keys = _model_state_dict_keys(model)
|
|
tensors_by_file: dict[str, list[str]] = {}
|
|
skipped: list[str] = []
|
|
for tensor_name in sorted(tensor_names):
|
|
rel_file = weight_map.get(tensor_name)
|
|
if not isinstance(rel_file, str):
|
|
continue
|
|
if (
|
|
expected_keys is not None
|
|
and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys
|
|
):
|
|
skipped.append(tensor_name)
|
|
continue
|
|
tensors_by_file.setdefault(rel_file, []).append(tensor_name)
|
|
if skipped:
|
|
preview = ", ".join(skipped[:3])
|
|
print(
|
|
f" Skipping {len(skipped)} checkpoint tensors absent from the causal LM "
|
|
f"(e.g. {preview})",
|
|
flush=True,
|
|
)
|
|
if not tensors_by_file:
|
|
raise PartialModelLoadUnsupported(
|
|
f"no checkpoint tensors for layers {shard_start}-{shard_end} match the "
|
|
f"causal LM built from {snapshot_dir}"
|
|
)
|
|
|
|
for rel_file, names in tensors_by_file.items():
|
|
checkpoint_file = snapshot_dir / rel_file
|
|
if not checkpoint_file.exists():
|
|
raise PartialModelLoadUnsupported(
|
|
f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}"
|
|
)
|
|
with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle:
|
|
for tensor_name in names:
|
|
set_tensor_fn(
|
|
model,
|
|
_checkpoint_tensor_name_for_model(model, tensor_name),
|
|
device,
|
|
value=handle.get_tensor(tensor_name),
|
|
dtype=dtype,
|
|
)
|
|
|
|
for module in _active_modules_for_shard(model, shard_start, shard_end):
|
|
if hasattr(module, "to"):
|
|
module.to(device)
|
|
return model
|
|
|
|
|
|
def _model_load_plan(
|
|
auto_config: Any,
|
|
model_id: str,
|
|
quantization: Quantization,
|
|
torch: Any,
|
|
cache_dir: Path | None = None,
|
|
) -> tuple[Any | None, Any, bool]:
|
|
"""Return (explicit quant config, dtype, uses quantized weights)."""
|
|
if quantization != "auto":
|
|
quant_config = build_quantization_config(quantization)
|
|
return quant_config, torch.bfloat16, quant_config is not None
|
|
|
|
cfg = auto_config.from_pretrained(
|
|
model_id,
|
|
cache_dir=str(cache_dir) if cache_dir is not None else None,
|
|
)
|
|
if _native_quantization_config(cfg) is not None:
|
|
return None, _native_torch_dtype(cfg, torch), True
|
|
return None, _native_torch_dtype(cfg, torch), False
|
|
|
|
|
|
def _config_candidates(cfg: Any) -> list[Any]:
|
|
candidates = [cfg]
|
|
get_text_config = getattr(cfg, "get_text_config", None)
|
|
if callable(get_text_config):
|
|
try:
|
|
candidates.append(get_text_config())
|
|
except Exception:
|
|
pass
|
|
text_config = getattr(cfg, "text_config", None)
|
|
if text_config is not None:
|
|
candidates.append(text_config)
|
|
return candidates
|
|
|
|
|
|
def _native_quantization_config(cfg: Any) -> Any | None:
|
|
for candidate in _config_candidates(cfg):
|
|
quant_config = getattr(candidate, "quantization_config", None)
|
|
if quant_config:
|
|
return quant_config
|
|
return None
|
|
|
|
|
|
def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
|
|
for candidate in _config_candidates(cfg):
|
|
for attr in ("dtype", "torch_dtype"):
|
|
dtype = getattr(candidate, attr, None)
|
|
if dtype is None:
|
|
continue
|
|
if isinstance(dtype, str):
|
|
dtype_name = dtype.removeprefix("torch.")
|
|
dtype_value = getattr(torch, dtype_name, None)
|
|
if dtype_value is not None:
|
|
return dtype_value
|
|
else:
|
|
return dtype
|
|
return torch.bfloat16
|
|
|
|
|
|
def _causal_lm_config(cfg: Any) -> Any:
|
|
"""Use the text decoder config for composite VLM/MoE presets."""
|
|
get_text_config = getattr(cfg, "get_text_config", None)
|
|
if callable(get_text_config):
|
|
try:
|
|
return get_text_config()
|
|
except Exception:
|
|
pass
|
|
text_config = getattr(cfg, "text_config", None)
|
|
if text_config is not None:
|
|
return text_config
|
|
return cfg
|
|
|
|
|
|
def _model_state_dict_keys(model: Any) -> set[str] | None:
|
|
"""Expected parameter/buffer names, or None when the model can't report them."""
|
|
state_dict = getattr(model, "state_dict", None)
|
|
if not callable(state_dict):
|
|
return None
|
|
try:
|
|
return set(state_dict().keys())
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
|
|
"""Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
|
|
inner = getattr(model, "model", None)
|
|
if inner is not None and hasattr(inner, "language_model"):
|
|
return tensor_name
|
|
if ".language_model." in tensor_name:
|
|
return tensor_name.replace(".language_model.", ".")
|
|
return tensor_name
|
|
|
|
|
|
def _transformer_backbone(model: Any) -> Any:
|
|
if hasattr(model, "model"):
|
|
inner = model.model
|
|
language_model = getattr(inner, "language_model", None)
|
|
if language_model is not None:
|
|
return language_model
|
|
return inner
|
|
if hasattr(model, "transformer"):
|
|
return model.transformer
|
|
raise ModelBackendError(
|
|
"unsupported HuggingFace model architecture: no transformer backbone found"
|
|
)
|
|
|
|
|
|
def _model_layers(model: Any) -> Any:
|
|
backbone = _transformer_backbone(model)
|
|
for attr in ("layers", "h", "blocks"):
|
|
layers = getattr(backbone, attr, None)
|
|
if layers is not None:
|
|
return layers
|
|
raise ModelBackendError(
|
|
"unsupported HuggingFace model architecture: no transformer layers found"
|
|
)
|
|
|
|
|
|
def _embed_tokens(model: Any) -> Any:
|
|
backbone = _transformer_backbone(model)
|
|
for attr in ("embed_tokens", "wte"):
|
|
embed = getattr(backbone, attr, None)
|
|
if embed is not None:
|
|
return embed
|
|
raise ModelBackendError(
|
|
"unsupported HuggingFace model architecture: no token embeddings found"
|
|
)
|
|
|
|
|
|
def _position_embeddings(model: Any) -> Any | None:
|
|
backbone = _transformer_backbone(model)
|
|
return getattr(backbone, "wpe", None)
|
|
|
|
|
|
def _rotary_embedding_module(model: Any) -> Any | None:
|
|
backbone = _transformer_backbone(model)
|
|
return getattr(backbone, "rotary_emb", None)
|
|
|
|
|
|
def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
|
|
active: list[Any] = []
|
|
|
|
def add(module: Any | None) -> None:
|
|
if module is None:
|
|
return
|
|
if any(existing is module for existing in active):
|
|
return
|
|
active.append(module)
|
|
|
|
if shard_start == 0:
|
|
add(_embed_tokens(model))
|
|
add(_position_embeddings(model))
|
|
add(_rotary_embedding_module(model))
|
|
for layer in _model_layers(model)[shard_start:shard_end + 1]:
|
|
add(layer)
|
|
total_layers = len(_model_layers(model))
|
|
if shard_end >= total_layers - 1:
|
|
add(_final_norm(model))
|
|
add(getattr(model, "lm_head", None))
|
|
return active
|
|
|
|
|
|
def _final_norm(model: Any) -> Any | None:
|
|
backbone = _transformer_backbone(model)
|
|
for attr in ("norm", "ln_f", "final_layer_norm"):
|
|
norm = getattr(backbone, attr, None)
|
|
if norm is not None:
|
|
return norm
|
|
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 = _rotary_embedding_module(model)
|
|
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
|
|
or "paging file is too small" in text
|
|
or "os error 1455" in text
|
|
):
|
|
return True
|
|
current = current.__cause__ or current.__context__
|
|
return False
|