"""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 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.""" 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: 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, ) -> 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 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, ) 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 # 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" 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 load_torch_shard( model_id: str, shard_start: int, shard_end: int, quantization: Quantization = "auto", cache_dir: Path | None = None, force_cpu: bool = False, ) -> TorchModelShard: return TorchModelShard( model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu ) 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, 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: 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 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) )