1232 lines
46 KiB
Python
1232 lines
46 KiB
Python
"""HuggingFace/PyTorch shard backend for real node inference."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from collections import OrderedDict
|
|
from dataclasses import dataclass
|
|
import json
|
|
import os
|
|
import threading
|
|
import time
|
|
import warnings
|
|
from pathlib import Path
|
|
from typing import Any, Literal, Mapping
|
|
|
|
Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
|
|
|
|
# Recipe params this backend knows how to apply (see meshnet_node.recipe_manifest).
|
|
# A recipe is only meaningful if its params actually reach the execution path, so
|
|
# an unknown key is an error rather than a silent no-op.
|
|
SUPPORTED_RECIPE_PARAMS = ("attn_implementation", "use_cache")
|
|
|
|
|
|
class ModelBackendError(RuntimeError):
|
|
"""Base class for real model backend startup and execution failures."""
|
|
|
|
|
|
class UnsupportedRecipeParam(ModelBackendError):
|
|
"""Raised when a recipe asks for an execution param this backend cannot apply."""
|
|
|
|
|
|
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."""
|
|
|
|
|
|
class KVCacheMiss(ModelBackendError):
|
|
"""Raised when a decode step references session state this node no longer holds.
|
|
|
|
The head recovers by re-prefilling the full sequence (the stateless path),
|
|
so eviction or a node restart degrades throughput instead of corrupting output.
|
|
"""
|
|
|
|
|
|
def _torch_cuda_is_executable(torch_module: Any) -> bool:
|
|
"""Return True only when this process can actually execute a CUDA/HIP op.
|
|
|
|
On ROCm, ``torch.cuda.is_available()`` can be true for an AMD GPU even when
|
|
the installed PyTorch wheel has no runnable kernels for that GPU target.
|
|
Loading weights onto such a device can segfault in native code, so the model
|
|
backend must use the same executable-device check as startup hardware
|
|
detection rather than trusting inventory alone.
|
|
"""
|
|
try:
|
|
if not torch_module.cuda.is_available():
|
|
return False
|
|
probe = torch_module.empty((1,), device="cuda")
|
|
probe += 1
|
|
torch_module.cuda.synchronize()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TensorPayload:
|
|
"""An immutable, request-owned binary activation payload.
|
|
|
|
``body`` is always the exact bfloat16 wire body. It is intentionally
|
|
owned bytes rather than a view into a request buffer so a payload can move
|
|
across a hop without retaining an HTTP/WebSocket frame after that request
|
|
completes.
|
|
"""
|
|
|
|
body: bytes
|
|
shape: list[int]
|
|
attention_mask_header: str | None
|
|
position_ids_header: str | None
|
|
# Number of tokens already cached before this payload's tokens (decode steps).
|
|
past_len: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TailTokenResult:
|
|
"""Tail-shard decode result: decoded text plus the raw token id.
|
|
|
|
The token id lets the head feed the next decode step (and detect EOS)
|
|
without re-tokenizing text, which is not guaranteed to round-trip.
|
|
"""
|
|
|
|
text: str
|
|
token_id: int
|
|
|
|
|
|
@dataclass
|
|
class SessionCacheEntry:
|
|
"""Per-session cached state for one shard's layer range.
|
|
|
|
`cache` is whatever `use_cache=True` produces for these layers — a
|
|
transformers Cache holding K/V tensors for standard attention, or
|
|
recurrent conv/delta state for hybrid linear-attention layers. The store
|
|
treats it as opaque.
|
|
"""
|
|
|
|
cache: Any
|
|
seq_len: int
|
|
effective_start: int
|
|
last_used: float
|
|
|
|
|
|
class SessionCacheStore:
|
|
"""TTL + LRU bounded map of session_id → SessionCacheEntry.
|
|
|
|
Each node caches state only for its own layer range; no node ever holds
|
|
another node's cache. Stale or mismatched entries raise KVCacheMiss so the
|
|
head falls back to a full re-prefill instead of producing corrupt output.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
max_sessions: int = 8,
|
|
ttl_seconds: float = 600.0,
|
|
clock: Any = None,
|
|
) -> None:
|
|
self.max_sessions = max(1, int(max_sessions))
|
|
self.ttl_seconds = float(ttl_seconds)
|
|
self._clock = clock or time.monotonic
|
|
self._entries: OrderedDict[str, SessionCacheEntry] = OrderedDict()
|
|
self._lock = threading.Lock()
|
|
|
|
def __len__(self) -> int:
|
|
with self._lock:
|
|
return len(self._entries)
|
|
|
|
def store(self, session_id: str, cache: Any, seq_len: int, effective_start: int) -> SessionCacheEntry:
|
|
now = self._clock()
|
|
with self._lock:
|
|
self._entries.pop(session_id, None)
|
|
entry = SessionCacheEntry(cache, seq_len, effective_start, now)
|
|
self._entries[session_id] = entry
|
|
self._evict_locked(now)
|
|
return entry
|
|
|
|
def lookup(
|
|
self,
|
|
session_id: str,
|
|
*,
|
|
expected_seq_len: int | None = None,
|
|
effective_start: int | None = None,
|
|
) -> SessionCacheEntry:
|
|
now = self._clock()
|
|
with self._lock:
|
|
self._evict_locked(now)
|
|
entry = self._entries.get(session_id)
|
|
if entry is None:
|
|
raise KVCacheMiss(f"no cached state for session {session_id[:8]}")
|
|
if expected_seq_len is not None and entry.seq_len != expected_seq_len:
|
|
del self._entries[session_id]
|
|
raise KVCacheMiss(
|
|
f"session {session_id[:8]} cache holds {entry.seq_len} tokens, "
|
|
f"expected {expected_seq_len}"
|
|
)
|
|
if effective_start is not None and entry.effective_start != effective_start:
|
|
del self._entries[session_id]
|
|
raise KVCacheMiss(
|
|
f"session {session_id[:8]} cached with start_layer "
|
|
f"{entry.effective_start}, requested {effective_start}"
|
|
)
|
|
entry.last_used = now
|
|
self._entries.move_to_end(session_id)
|
|
return entry
|
|
|
|
def drop(self, session_id: str) -> None:
|
|
with self._lock:
|
|
self._entries.pop(session_id, None)
|
|
|
|
def _evict_locked(self, now: float) -> None:
|
|
if self.ttl_seconds > 0:
|
|
expired = [
|
|
sid for sid, entry in self._entries.items()
|
|
if now - entry.last_used > self.ttl_seconds
|
|
]
|
|
for sid in expired:
|
|
del self._entries[sid]
|
|
while len(self._entries) > self.max_sessions:
|
|
self._entries.popitem(last=False)
|
|
|
|
|
|
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,
|
|
force_cpu: bool = False,
|
|
recipe_params: Mapping[str, Any] | 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
|
|
self.recipe_params = validate_recipe_params(recipe_params)
|
|
attn_implementation = self.recipe_params.get("attn_implementation")
|
|
|
|
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
|
|
if force_cpu:
|
|
self.device = torch.device("cpu")
|
|
else:
|
|
self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) 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,
|
|
attn_implementation=attn_implementation,
|
|
)
|
|
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
|
|
if attn_implementation is not None:
|
|
load_kwargs["attn_implementation"] = attn_implementation
|
|
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
|
|
# Per-session KV/recurrent-state cache for this shard's layer range.
|
|
# Hybrid/linear-attention models such as Qwen3.6 can dispatch Triton
|
|
# recurrent-cache kernels when use_cache=True. Those kernels cannot
|
|
# consume CPU tensors ("Pointer argument cannot be accessed from Triton"),
|
|
# so CPU shards intentionally stay on the stateless prefill path.
|
|
self.supports_kv_cache = self.device.type != "cpu"
|
|
if self.recipe_params.get("use_cache") is False:
|
|
self.supports_kv_cache = False
|
|
self.kv_sessions = SessionCacheStore(
|
|
max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
|
|
ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
|
|
)
|
|
|
|
def encode_prompt(self, prompt: str, session_id: str | None = None) -> 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_session(
|
|
hidden_states, attention_mask, position_ids,
|
|
session_id=session_id, cache_mode="prefill" if session_id else None,
|
|
)
|
|
return self._payload(hidden_states, attention_mask, position_ids)
|
|
|
|
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
|
|
"""Decode step: embed one new token against this head's cached session.
|
|
|
|
Raises KVCacheMiss if the session was evicted — callers fall back to a
|
|
full re-prefill via encode_prompt.
|
|
"""
|
|
if not self.is_head or self._embed_tokens is None:
|
|
raise ModelBackendError("decode steps can only start at the head shard")
|
|
if not self.supports_kv_cache:
|
|
raise KVCacheMiss("kv cache disabled on this backend")
|
|
entry = self.kv_sessions.lookup(
|
|
session_id, effective_start=self._effective_start(None)
|
|
)
|
|
past_len = entry.seq_len
|
|
input_ids = self.torch.tensor([[int(token_id)]], dtype=self.torch.long, device=self.device)
|
|
position_ids = self.torch.tensor([[past_len]], dtype=self.torch.long, device=self.device)
|
|
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, None, position_ids,
|
|
cache=entry.cache, past_len=past_len,
|
|
)
|
|
entry.seq_len = past_len + 1
|
|
return TensorPayload(
|
|
body=_tensor_to_bytes(hidden_states.to(self.torch.bfloat16).contiguous()),
|
|
shape=list(hidden_states.shape),
|
|
attention_mask_header=None,
|
|
position_ids_header=_int_tensor_header(position_ids),
|
|
past_len=past_len,
|
|
)
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body: bytes,
|
|
shape: list[int],
|
|
attention_mask_header: str | None,
|
|
position_ids_header: str | None,
|
|
start_layer: int | None = None,
|
|
session_id: str | None = None,
|
|
cache_mode: str | None = None,
|
|
past_len: int | None = None,
|
|
) -> TensorPayload | TailTokenResult | 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_session(
|
|
hidden_states, attention_mask, position_ids, start_layer=start_layer,
|
|
session_id=session_id, cache_mode=cache_mode, past_len=past_len,
|
|
)
|
|
if self.is_tail:
|
|
return self.decode_tail_token(hidden_states)
|
|
return self._payload(hidden_states, attention_mask, position_ids)
|
|
|
|
def decode_tail(self, hidden_states: Any) -> str:
|
|
return self.decode_tail_token(hidden_states).text
|
|
|
|
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
|
|
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 TailTokenResult(
|
|
text=self.tokenizer.decode([token_id], skip_special_tokens=True),
|
|
token_id=token_id,
|
|
)
|
|
|
|
def eos_token_ids(self) -> list[int]:
|
|
"""All token ids that should terminate generation (tokenizer + generation config)."""
|
|
ids: set[int] = set()
|
|
tok_eos = getattr(self.tokenizer, "eos_token_id", None)
|
|
gen_config = getattr(self.model, "generation_config", None)
|
|
gen_eos = getattr(gen_config, "eos_token_id", None) if gen_config is not None else None
|
|
for value in (tok_eos, gen_eos):
|
|
if value is None:
|
|
continue
|
|
if isinstance(value, (list, tuple)):
|
|
ids.update(int(v) for v in value)
|
|
else:
|
|
ids.add(int(value))
|
|
return sorted(ids)
|
|
|
|
def release_session(self, session_id: str) -> None:
|
|
self.kv_sessions.drop(session_id)
|
|
|
|
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 _effective_start(self, start_layer: int | None) -> int:
|
|
# 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.
|
|
return (
|
|
max(self.shard_start, start_layer)
|
|
if start_layer is not None
|
|
else self.shard_start
|
|
)
|
|
|
|
def _new_session_cache(self) -> Any | None:
|
|
"""Build the model-appropriate cache object for one session.
|
|
|
|
DynamicCache(config=...) lets transformers pick the right per-layer
|
|
state (K/V for standard attention, conv/recurrent state for hybrid
|
|
linear-attention layers) — the same construction the model's own
|
|
forward() uses when use_cache=True.
|
|
"""
|
|
try:
|
|
from transformers import DynamicCache
|
|
except ImportError:
|
|
return None
|
|
try:
|
|
return DynamicCache(config=self.model.config)
|
|
except TypeError:
|
|
return DynamicCache()
|
|
|
|
def _run_layers_session(
|
|
self,
|
|
hidden_states: Any,
|
|
attention_mask: Any,
|
|
position_ids: Any,
|
|
start_layer: int | None = None,
|
|
session_id: str | None = None,
|
|
cache_mode: str | None = None,
|
|
past_len: int | None = None,
|
|
) -> Any:
|
|
"""Run this shard's layers, keying cached state by session when requested.
|
|
|
|
cache_mode "prefill" creates fresh session state; "decode" requires an
|
|
existing entry (KVCacheMiss otherwise). None runs fully stateless —
|
|
today's behavior, kept as the recovery path.
|
|
"""
|
|
effective_start = self._effective_start(start_layer)
|
|
if not (session_id and cache_mode and self.supports_kv_cache):
|
|
if cache_mode == "decode":
|
|
# A decode payload is one token — running it stateless would
|
|
# silently produce garbage. Force the head to re-prefill.
|
|
raise KVCacheMiss("kv cache disabled on this backend")
|
|
return self._run_layers(
|
|
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
|
)
|
|
if cache_mode == "decode":
|
|
entry = self.kv_sessions.lookup(
|
|
session_id,
|
|
expected_seq_len=past_len,
|
|
effective_start=effective_start,
|
|
)
|
|
seq_len = int(hidden_states.shape[1])
|
|
# Decode attends over cache + new token; no padding, so no mask needed.
|
|
hidden_states = self._run_layers(
|
|
hidden_states, None, position_ids,
|
|
start_layer=start_layer, cache=entry.cache, past_len=entry.seq_len,
|
|
)
|
|
entry.seq_len += seq_len
|
|
return hidden_states
|
|
# Prefill: fresh cache for this session (replaces any stale entry).
|
|
cache = self._new_session_cache()
|
|
if cache is None:
|
|
return self._run_layers(
|
|
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
|
)
|
|
try:
|
|
result = self._run_layers(
|
|
hidden_states, attention_mask, position_ids,
|
|
start_layer=start_layer, cache=cache, past_len=0,
|
|
)
|
|
except Exception as exc:
|
|
if not _cache_unsupported_for_shard(exc):
|
|
raise
|
|
# Layers reject cache kwargs (exotic architecture) — disable caching
|
|
# for this backend and stay on the stateless path. Some hybrid
|
|
# CPU paths also accept cache kwargs but fail at runtime inside
|
|
# Triton-only kernels; treat those as cache-unsupported too.
|
|
self.supports_kv_cache = False
|
|
print(f" [node] kv cache unsupported by {self.model_id}: {exc}", flush=True)
|
|
return self._run_layers(
|
|
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
|
)
|
|
self.kv_sessions.store(
|
|
session_id, cache,
|
|
seq_len=int(hidden_states.shape[1]),
|
|
effective_start=effective_start,
|
|
)
|
|
return result
|
|
|
|
def _run_layers(
|
|
self,
|
|
hidden_states: Any,
|
|
attention_mask: Any,
|
|
position_ids: Any,
|
|
start_layer: int | None = None,
|
|
cache: Any = None,
|
|
past_len: int = 0,
|
|
) -> Any:
|
|
effective_start = self._effective_start(start_layer)
|
|
position_embeddings = _rotary_position_embeddings(
|
|
self.model,
|
|
hidden_states,
|
|
position_ids,
|
|
)
|
|
layer_attention_mask = _decoder_attention_mask(
|
|
attention_mask,
|
|
hidden_states,
|
|
self.torch,
|
|
)
|
|
cache_position = None
|
|
if cache is not None:
|
|
seq_len = int(hidden_states.shape[1])
|
|
cache_position = self.torch.arange(
|
|
past_len, past_len + seq_len, device=hidden_states.device
|
|
)
|
|
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,
|
|
cache=cache,
|
|
cache_position=cache_position,
|
|
)
|
|
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 validate_recipe_params(params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
"""Return recipe params this backend can honour, or raise naming the bad key."""
|
|
if not params:
|
|
return {}
|
|
unsupported = [key for key in params if key not in SUPPORTED_RECIPE_PARAMS]
|
|
if unsupported:
|
|
raise UnsupportedRecipeParam(
|
|
f"recipe param(s) {', '.join(sorted(unsupported))} are not supported by this "
|
|
f"backend; it applies: {', '.join(SUPPORTED_RECIPE_PARAMS)}"
|
|
)
|
|
return dict(params)
|
|
|
|
|
|
def load_torch_shard(
|
|
model_id: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
quantization: Quantization = "auto",
|
|
cache_dir: Path | None = None,
|
|
force_cpu: bool = False,
|
|
recipe_params: Mapping[str, Any] | None = None,
|
|
) -> TorchModelShard:
|
|
return TorchModelShard(
|
|
model_id,
|
|
shard_start,
|
|
shard_end,
|
|
quantization,
|
|
cache_dir,
|
|
force_cpu=force_cpu,
|
|
recipe_params=recipe_params,
|
|
)
|
|
|
|
|
|
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,
|
|
attn_implementation: str | 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))
|
|
if attn_implementation is not None:
|
|
# The partial path instantiates from the config, so the attention choice
|
|
# has to be set on it rather than passed to from_pretrained.
|
|
cfg._attn_implementation = attn_implementation
|
|
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,
|
|
cache: Any = None,
|
|
cache_position: Any = 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}
|
|
if cache is not None:
|
|
# transformers 5.x layers take a Cache via past_key_values and
|
|
# mutate it in place; cache_position is required by sliding-window
|
|
# and hybrid recurrent layers.
|
|
filtered["past_key_values"] = cache
|
|
filtered["use_cache"] = True
|
|
if cache_position is not None:
|
|
filtered["cache_position"] = cache_position
|
|
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:
|
|
# ``frombuffer`` views the immutable request-owned bytes for this forward
|
|
# only. The following device transfer is the one required CPU→GPU copy;
|
|
# wrapping in ``bytearray`` first used to add an avoidable CPU allocation
|
|
# and copy. Do not upcast through float32: the activation wire contract
|
|
# is bfloat16 and model layers accept it directly.
|
|
# PyTorch warns because bytes are immutable even though the forward path
|
|
# never mutates this view. Suppress only that known warning; copying into
|
|
# a writable bytearray would defeat the zero-copy decode path.
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
message="The given buffer is not writable.*",
|
|
category=UserWarning,
|
|
)
|
|
tensor = torch.frombuffer(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
|
|
|
|
|
|
def _cache_unsupported_for_shard(exc: BaseException) -> bool:
|
|
"""True when a layer failure means session cache is unsupported, not fatal."""
|
|
text = str(exc).lower()
|
|
return (
|
|
isinstance(exc, TypeError)
|
|
or "pointer argument cannot be accessed from triton" in text
|
|
or ("triton" in text and "cpu tensor" in text)
|
|
)
|