""" chunked_model.py — Memory-efficient layer-by-layer inference engine. Drop-in replacement for airllm. Zero airllm dependency. How it works ──────────── • Model weights are split into small shard files (~chunk_mb MB each) once on first load, then reused on every subsequent run. • During inference each layer's weights are loaded from its shard, the layer forward pass is executed, then the weights are released from RAM. • An in-memory layer cache avoids re-reading the same shard twice in one generation call (huge speedup for multi-token generation). • KV cache: after the prefill pass only one new token is forwarded per step, so the per-token cost is one layer-by-layer pass over 1 token — much faster than airllm's O(n²) approach. Supported architectures ─────────────────────── Qwen2ForCausalLM · Qwen3ForCausalLM · Qwen3_5ForConditionalGeneration LlamaForCausalLM · MistralForCausalLM · MixtralForCausalLM GemmaForCausalLM · Gemma2ForCausalLM · Phi3ForCausalLM Usage ───── from chunked_model import ChunkedModel model = ChunkedModel("./model", chunk_mb=75) out = model.generate(input_ids, max_new_tokens=200, temperature=0.7) """ from __future__ import annotations import gc, json, os from pathlib import Path from typing import Dict, List, Optional, Tuple import torch import torch.nn.functional as F from safetensors import safe_open from safetensors.torch import save_file from transformers import AutoConfig # ───────────────────────────────────────────────────────────────────────────── # Math primitives # ───────────────────────────────────────────────────────────────────────────── def _rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: # Compute norm in float32 for stability, then cast back to original dtype x_f32 = x.float() variance = x_f32.pow(2).mean(-1, keepdim=True) normed = x_f32 * torch.rsqrt(variance + eps) return (w.float() * normed).to(x.dtype) def _rotate_half(x: torch.Tensor) -> torch.Tensor: h = x.shape[-1] // 2 return torch.cat([-x[..., h:], x[..., :h]], dim=-1) def _apply_rope( q: torch.Tensor, k: torch.Tensor, position_ids: torch.Tensor, head_dim: int, rope_theta: float = 1_000_000.0, partial_factor: float = 1.0, ) -> Tuple[torch.Tensor, torch.Tensor]: """Apply Rotary Position Embedding to query and key tensors. partial_factor < 1.0: only apply RoPE to the first int(head_dim*partial_factor) dimensions (used by Qwen3.5 which uses partial_rotary_factor=0.25). """ device = q.device rot_dim = int(head_dim * partial_factor) if rot_dim % 2 != 0: rot_dim -= 1 inv_freq = 1.0 / ( rope_theta ** ( torch.arange(0, rot_dim, 2, device=device, dtype=torch.float32) / rot_dim ) ) freqs = torch.einsum("bi,j->bij", position_ids.float(), inv_freq) emb = torch.cat([freqs, freqs], dim=-1) # (B, T, rot_dim) cos = emb.cos().unsqueeze(1) # (B, 1, T, rot_dim) sin = emb.sin().unsqueeze(1) # Cast cos/sin to q's dtype so RoPE doesn't upcast bfloat16 tensors cos = cos.to(q.dtype) sin = sin.to(q.dtype) if partial_factor < 1.0: q_rot, q_pass = q[..., :rot_dim], q[..., rot_dim:] k_rot, k_pass = k[..., :rot_dim], k[..., rot_dim:] q_rot = q_rot * cos + _rotate_half(q_rot) * sin k_rot = k_rot * cos + _rotate_half(k_rot) * sin q = torch.cat([q_rot, q_pass], dim=-1) k = torch.cat([k_rot, k_pass], dim=-1) else: q = q * cos + _rotate_half(q) * sin k = k * cos + _rotate_half(k) * sin return q, k def _swiglu( x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor, chunk_rows: int = 0, ) -> torch.Tensor: """SwiGLU feed-forward: down( silu(gate(x)) ⊙ up(x) ). chunk_rows > 0: compute the intermediate dimension in slices so that only `chunk_rows` worth of gate/up/down are live at once. Use this for very large models where the MLP weights exceed the chunk budget. """ interm = gate_w.shape[0] if chunk_rows <= 0 or chunk_rows >= interm: return F.linear( F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w, ) # Chunked path — accumulate into output buffer row-slice by row-slice out = torch.zeros(*x.shape[:-1], down_w.shape[0], dtype=x.dtype, device=x.device) for start in range(0, interm, chunk_rows): end = min(start + chunk_rows, interm) gate = F.silu(F.linear(x, gate_w[start:end])) up = F.linear(x, up_w[start:end]) out += F.linear(gate * up, down_w[:, start:end]) del gate, up return out # ───────────────────────────────────────────────────────────────────────────── # Shard manager # ───────────────────────────────────────────────────────────────────────────── class _ShardManager: """ One-time split of model safetensors into fixed-size shard files. Builds index.json mapping every weight key to its shard filename. On subsequent runs the existing shards + index are reused as-is. """ _INDEX = "index.json" def __init__(self, model_path: Path, chunk_mb: int, dtype: torch.dtype): self.model_path = model_path self.chunk_bytes = chunk_mb * 1024 * 1024 self.dtype = dtype self.shard_dir = model_path / f"_chunks_{chunk_mb}mb" self.index: Dict[str, str] = {} # ── Public API ──────────────────────────────────────────────────────────── def prepare(self, num_layers: int) -> None: idx = self.shard_dir / self._INDEX if idx.exists(): with open(idx) as f: self.index = json.load(f) n_shards = len(set(self.index.values())) print(f"[chunked] Reusing {n_shards} shards ({self.shard_dir.name})", flush=True) else: self._build() def set_layer_prefix(self, prefix: str) -> None: """Configure the key prefix used by load_layer (e.g. 'model.language_model.').""" self._layer_prefix = prefix def load_layer(self, i: int) -> Dict[str, torch.Tensor]: pfx = getattr(self, "_layer_prefix", "model.") key = f"{pfx}layers.{i}." return self._load_iter(k for k in self.index if k.startswith(key)) def load_keys(self, *keys: str) -> Dict[str, torch.Tensor]: return self._load_iter(k for k in keys if k in self.index) # ── Internal ────────────────────────────────────────────────────────────── def _load_iter(self, keys) -> Dict[str, torch.Tensor]: shard_map: Dict[str, List[str]] = {} for k in keys: sf = self.index.get(k) if sf: shard_map.setdefault(sf, []).append(k) out: Dict[str, torch.Tensor] = {} for sf, ks in shard_map.items(): with safe_open(str(self.shard_dir / sf), framework="pt", device="cpu") as f: for k in ks: # Cast to target dtype at load time (shards keep original dtype) out[k] = f.get_tensor(k).to(self.dtype) return out # Key prefixes to SKIP during shard building (vision encoder, MTP, etc.) _SKIP_PREFIXES = ("model.visual", "mtp.", "model.embed.visual") def _read_source(self) -> Dict[str, torch.Tensor]: """Read all text-model weights from safetensors, keeping original dtype. Vision-encoder and MTP weights are skipped — they are not needed for text-only inference and would double RAM usage during shard building. dtype conversion (to float32) happens at inference time, not here. """ weights: Dict[str, torch.Tensor] = {} for st in sorted(self.model_path.glob("*.safetensors")): with safe_open(str(st), framework="pt", device="cpu") as f: for k in f.keys(): if any(k.startswith(pfx) for pfx in self._SKIP_PREFIXES): continue weights[k] = f.get_tensor(k) # keep original dtype (bf16/fp16) if not weights: raise RuntimeError(f"No .safetensors files found in {self.model_path}") return weights @staticmethod def _group_of(key: str) -> str: # Support both model.layers.{i}. and model.language_model.layers.{i}. if ".layers." in key and not key.startswith("model.visual") \ and not key.startswith("mtp"): parts = key.split(".layers.") if len(parts) >= 2: idx = parts[1].split(".")[0] if idx.isdigit(): return "layer_" + idx.zfill(5) if "embed_tokens" in key: return "00_embed" return "zz_head" def _build(self) -> None: chunk_mb = self.chunk_bytes // (1024 * 1024) self.shard_dir.mkdir(parents=True, exist_ok=True) print(f"[chunked] Building shards (~{chunk_mb} MB each) — one-time setup ...", flush=True) weights = self._read_source() # Group weights by logical layer groups: Dict[str, Dict[str, torch.Tensor]] = {} for k, t in weights.items(): g = self._group_of(k) groups.setdefault(g, {})[k] = t index: Dict[str, str] = {} shard_n = 0 def _flush(bucket: Dict[str, torch.Tensor]) -> str: nonlocal shard_n fname = f"s{shard_n:05d}.safetensors" save_file(bucket, str(self.shard_dir / fname)) shard_n += 1 return fname for g_name in sorted(groups): group = groups[g_name] g_bytes = sum(t.numel() * t.element_size() for t in group.values()) if g_bytes <= self.chunk_bytes: fname = _flush(group) for k in group: index[k] = fname else: bucket: Dict[str, torch.Tensor] = {} bucket_bytes = 0 for k, t in group.items(): tb = t.numel() * t.element_size() if tb > self.chunk_bytes and not bucket: # Single oversized tensor — save alone; chunked inference # handles it at runtime via _swiglu(chunk_rows=...) fname = _flush({k: t}) index[k] = fname continue if bucket and bucket_bytes + tb > self.chunk_bytes: fname = _flush(bucket) for bk in bucket: index[bk] = fname bucket, bucket_bytes = {}, 0 bucket[k] = t bucket_bytes += tb if bucket: fname = _flush(bucket) for bk in bucket: index[bk] = fname with open(self.shard_dir / self._INDEX, "w") as f: json.dump(index, f, indent=2) self.index = index total_mb = sum( (self.shard_dir / sf).stat().st_size for sf in set(index.values()) ) // (1024 * 1024) print(f"[chunked] {shard_n} shards created ({total_mb} MB) in {self.shard_dir.name}", flush=True) # ───────────────────────────────────────────────────────────────────────────── # In-memory layer cache # ───────────────────────────────────────────────────────────────────────────── class _LayerCache: """ Caches layer weight dicts in memory to avoid re-reading shards on every decode step. max_layers=None means unlimited (keep everything). For small models (≤ ~4 GB weights) this is effectively a full warm cache after the first generate() call. """ def __init__(self, max_layers: Optional[int] = None): self._cache: Dict[str, Dict[str, torch.Tensor]] = {} self._order: List[str] = [] # LRU insertion order self.max = max_layers def get(self, key: str) -> Optional[Dict[str, torch.Tensor]]: if key in self._cache: self._order.remove(key) self._order.append(key) return self._cache[key] return None def put(self, key: str, weights: Dict[str, torch.Tensor]) -> None: if self.max is not None and len(self._cache) >= self.max and key not in self._cache: evict = self._order.pop(0) del self._cache[evict] self._cache[key] = weights if key in self._order: self._order.remove(key) self._order.append(key) def clear(self) -> None: self._cache.clear() self._order.clear() gc.collect() def __len__(self) -> int: return len(self._cache) # ───────────────────────────────────────────────────────────────────────────── # ChunkedModel — the inference engine # ───────────────────────────────────────────────────────────────────────────── class ChunkedModel: """ Memory-efficient transformer inference — loads one shard at a time. Parameters ---------- model_path : str | Path Directory containing config.json + model.safetensors. chunk_mb : int Target shard size in MB (default 75). Smaller = less peak RAM, more disk I/O on cold start. dtype : torch.dtype float32 (default, safest) or float16 (faster, needs modern CPU). cache_layers : int | None How many layer weight dicts to keep in RAM between steps. None (default) = keep everything — fast generation, uses more RAM. Set to e.g. 4 for 70B+ models to cap working-set memory. mlp_chunk_rows : int > 0 to split MLP projections into row-slices at runtime. Auto-computed when 0 (default). """ SUPPORTED = { "Qwen2ForCausalLM", "Qwen3ForCausalLM", "Qwen3_5ForConditionalGeneration", "LlamaForCausalLM", "MistralForCausalLM", "MixtralForCausalLM", "GemmaForCausalLM", "Gemma2ForCausalLM", "Phi3ForCausalLM", } def __init__( self, model_path: str, chunk_mb: int = 75, dtype: Optional[torch.dtype] = None, # None → auto-detect from config cache_layers: Optional[int] = None, mlp_chunk_rows: int = 0, ): self.model_path = Path(model_path) if not self.model_path.is_dir(): raise RuntimeError( f"Model directory not found: {self.model_path}\n" "Run 'bash install.sh' to download the model." ) print(f"[chunked] Loading config from {self.model_path.name}", flush=True) cfg = AutoConfig.from_pretrained(str(self.model_path), local_files_only=True) self.cfg = cfg arch = (cfg.architectures or ["Unknown"])[0] if arch not in self.SUPPORTED: print(f"[chunked] WARNING: {arch} not in supported list, attempting generic path") # Auto-detect dtype from model config — bfloat16 halves RAM vs float32 if dtype is None: raw = getattr(cfg, "torch_dtype", None) or getattr(cfg, "dtype", None) raw = str(raw) if raw is not None else "" if "bfloat16" in raw or raw is torch.bfloat16: dtype = torch.bfloat16 elif "float16" in raw or raw is torch.float16: dtype = torch.float16 else: dtype = torch.float32 self.dtype = dtype # ── Config extraction: try top-level then text_config fallback ────────── # Qwen3.5 / multimodal models nest text params under cfg.text_config def _get(attr, default=None): v = getattr(cfg, attr, None) if v is None and hasattr(cfg, "text_config"): v = getattr(cfg.text_config, attr, None) return v if v is not None else default self.arch = arch self.L = _get("num_hidden_layers") self.H = _get("num_attention_heads") self.KVH = _get("num_key_value_heads", self.H) self.D = _get("hidden_size") self.Dh = _get("head_dim", self.D // self.H) self.interm = _get("intermediate_size") self.vocab = _get("vocab_size") self.eps = float(_get("rms_norm_eps", 1e-6)) self.tied = bool(_get("tie_word_embeddings", False)) # rope_theta may be nested in rope_parameters dict (Qwen3.5) rope_params = _get("rope_parameters") or {} self.theta = float( rope_params.get("rope_theta", None) or _get("rope_theta", None) or 1_000_000.0 ) # partial RoPE: Qwen3.5 uses partial_rotary_factor=0.25 self.rope_partial = float(_get("partial_rotary_factor", 1.0)) # Qwen3 / Qwen3.5: per-head RMSNorm on Q and K self.qk_norm = "Qwen3" in arch # Hybrid layer types: Qwen3.5 has linear_attention + full_attention raw_layer_types = _get("layer_types", []) if raw_layer_types: self.layer_types = list(raw_layer_types) else: self.layer_types = ["full_attention"] * self.L # Linear attention dims (Qwen3.5 SSM layers) self.lin_kh = _get("linear_num_key_heads", self.H) # K heads self.lin_vh = _get("linear_num_value_heads", self.H) # V heads self.lin_kdh = _get("linear_key_head_dim", self.Dh) # K head dim # V head dim = out_proj_input / lin_vh (auto from weights) self.lin_qdim = self.lin_kh * self.lin_kdh # Q total dim (=K total dim) self.lin_kdim = self.lin_kh * self.lin_kdh # K total dim # V total dim = out_proj rows = hidden_size … use out_proj shape at runtime # ── Weight key prefix detection ───────────────────────────────────────── # Standard models use "model." prefix; Qwen3.5 uses "model.language_model." self.prefix = "model.language_model." \ if arch == "Qwen3_5ForConditionalGeneration" else "model." print( f"[chunked] {arch} | {self.L} layers | " f"hidden={self.D} | " f"heads={self.H}(kv={self.KVH}) | " f"prefix={self.prefix.rstrip('.')}", flush=True, ) has_linear = any(lt == "linear_attention" for lt in self.layer_types) if has_linear: n_lin = sum(1 for lt in self.layer_types if lt == "linear_attention") n_full = self.L - n_lin print(f"[chunked] Hybrid layers: {n_full} full_attention + {n_lin} linear_attention (SSM approx)", flush=True) # MLP runtime chunking (for huge layer weights) if mlp_chunk_rows > 0: self.mlp_chunk = mlp_chunk_rows else: row_b = self.D * _dtype_bytes(dtype) rpc = max(256, (chunk_mb * 1024 * 1024 // 4) // row_b) mlp_b = 3 * self.interm * self.D * _dtype_bytes(dtype) self.mlp_chunk = rpc if mlp_b > chunk_mb * 1024 * 1024 else 0 # Shard storage self._sm = _ShardManager(self.model_path, chunk_mb, dtype) self._sm.set_layer_prefix(self.prefix) self._sm.prepare(self.L) # In-memory cache (avoids re-reading shards during decode steps) # Default: unlimited — warm on first generate(), stays warm. self._cache = _LayerCache(max_layers=cache_layers) # Always-pinned: embed and head weights (small overhead, huge speedup) self._embed_w: Optional[torch.Tensor] = None self._norm_w: Optional[torch.Tensor] = None self._head_w: Optional[torch.Tensor] = None if self.mlp_chunk > 0: print(f"[chunked] MLP sub-chunking active: {self.mlp_chunk} rows/pass", flush=True) print( f"[chunked] Ready — chunk={chunk_mb} MB " f"rope_theta={self.theta:.0f} " f"rope_partial={self.rope_partial} " f"tied={self.tied} " f"layer_cache={'unlimited' if cache_layers is None else cache_layers}", flush=True, ) # ── Cached weight loaders ───────────────────────────────────────────────── def _embed(self) -> torch.Tensor: if self._embed_w is None: key = self.prefix + "embed_tokens.weight" w = self._sm.load_keys(key) self._embed_w = w[key] return self._embed_w def _norm(self) -> torch.Tensor: if self._norm_w is None: key = self.prefix + "norm.weight" w = self._sm.load_keys(key) self._norm_w = w[key] return self._norm_w def _lm_head(self) -> torch.Tensor: if self._head_w is None: lm_key = self.prefix + "lm_head.weight" if self.tied or lm_key not in self._sm.index: self._head_w = self._embed() else: w = self._sm.load_keys(lm_key) self._head_w = w[lm_key] return self._head_w def _layer(self, i: int) -> Dict[str, torch.Tensor]: key = f"layer_{i}" cached = self._cache.get(key) if cached is not None: return cached w = self._sm.load_layer(i) self._cache.put(key, w) return w # ── Linear-attention (SSM approximation) forward ────────────────────────── def _linear_attn_step( self, hidden: torch.Tensor, w: dict, p: str, ) -> torch.Tensor: """ Approximate forward for Qwen3.5 linear_attention (Mamba/SSM) layers. The true computation is a state-space recurrence; here we approximate it as standard scaled-dot-product attention over the same QKV projections. Quality is slightly lower than the full SSM but produces coherent text. All dims are auto-detected from actual weight shapes (no config needed): in_proj_qkv : [qkv_total, D] → Q + K + V (V = out_proj input dim) in_proj_z : [gate_dim, D] → gate norm.weight : [vdh] → value head dim (e.g. 128) out_proj : [D, v_dim] """ B, T, D = hidden.shape out_proj_w = w[p + "linear_attn.out_proj.weight"] # [D, v_dim] v_dim = out_proj_w.shape[1] # e.g. 4096 qkv = F.linear(hidden, w[p + "linear_attn.in_proj_qkv.weight"]) # [B,T,qkv_total] z = F.linear(hidden, w[p + "linear_attn.in_proj_z.weight"]) # [B,T,gate_dim] # Auto-derive Q and K dims: total - V, split equally between Q and K qkv_total = qkv.shape[-1] # e.g. 8192 q_dim = k_dim = (qkv_total - v_dim) // 2 # e.g. 2048 each q = qkv[..., :q_dim] # [B,T,q_dim] k = qkv[..., q_dim:q_dim + k_dim] # [B,T,k_dim] v = qkv[..., q_dim + k_dim:] # [B,T,v_dim] # Head count from norm.weight (norm applied per value head) norm_key = p + "linear_attn.norm.weight" nw = w.get(norm_key) if nw is not None: vdh = int(nw.shape[0]) # value head dim (e.g. 128) else: vdh = max(1, v_dim // 32) # fallback kh = v_dim // vdh # number of heads (e.g. 32) kdh = q_dim // kh # query/key head dim (e.g. 64) q = q.view(B, T, kh, kdh).transpose(1, 2) # [B, kh, T, kdh] k = k.view(B, T, kh, kdh).transpose(1, 2) # [B, kh, T, kdh] v = v.view(B, T, kh, vdh).transpose(1, 2) # [B, kh, T, vdh] is_causal = T > 1 attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal) attn = attn.transpose(1, 2).contiguous().view(B, T, v_dim) # [B,T,v_dim] # Per-value-head RMS norm (norm.weight tiled across all heads) if nw is not None: nw_tiled = nw.repeat(kh) if kh > 1 else nw # [vdh * kh] = [v_dim] attn = _rms_norm(attn, nw_tiled.to(attn.dtype), self.eps) # Gating: element-wise silu gate gated = attn * F.silu(z) # [B,T,v_dim] return F.linear(gated, out_proj_w) # [B,T,D] # ── Forward pass ────────────────────────────────────────────────────────── def _forward( self, input_ids: torch.Tensor, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, ): B, T = input_ids.shape past_n = past_kv[0][0].shape[2] if past_kv is not None else 0 pos_ids = torch.arange(past_n, past_n + T, device=input_ids.device).unsqueeze(0) # Embedding hidden = F.embedding(input_ids, self._embed().to(self.dtype)) new_kv: List[Tuple[torch.Tensor, torch.Tensor]] = [] for i in range(self.L): w = self._layer(i) p = f"{self.prefix}layers.{i}." ltype = self.layer_types[i] if i < len(self.layer_types) else "full_attention" # ── Pre-attention norm ────────────────────────────────────────── res = hidden hidden = _rms_norm(hidden, w[p + "input_layernorm.weight"], self.eps) # ── Attention block (full or linear) ──────────────────────────── if ltype == "linear_attention": # SSM approximation — no KV cache for these layers attn_out = self._linear_attn_step(hidden, w, p) new_kv.append(( torch.zeros(B, self.KVH, 0, self.Dh), torch.zeros(B, self.KVH, 0, self.Dh), )) else: # Full self-attention (Qwen3.5: q_proj outputs Q+gate combined) q_raw = F.linear(hidden, w[p + "self_attn.q_proj.weight"]) k = F.linear(hidden, w[p + "self_attn.k_proj.weight"]) v = F.linear(hidden, w[p + "self_attn.v_proj.weight"]) # Detect gated Q: q_proj output is 2× expected → split Q and gate expected_q = self.H * self.Dh # e.g. 16*256=4096 q_attn_gate = q_raw.shape[-1] == expected_q * 2 # Qwen3.5 gated if q_attn_gate: q, q_gate = q_raw.chunk(2, dim=-1) # each [B,T,H*Dh] else: q, q_gate = q_raw, None q = q.view(B, T, self.H, self.Dh).transpose(1, 2) k = k.view(B, T, self.KVH, self.Dh).transpose(1, 2) v = v.view(B, T, self.KVH, self.Dh).transpose(1, 2) # Per-head RMSNorm on Q and K (Qwen3 / Qwen3.5) if self.qk_norm: q = _rms_norm(q, w[p + "self_attn.q_norm.weight"], self.eps) k = _rms_norm(k, w[p + "self_attn.k_norm.weight"], self.eps) q, k = _apply_rope(q, k, pos_ids, self.Dh, self.theta, self.rope_partial) if past_kv is not None: pk, pv = past_kv[i] if pk.shape[2] > 0: # skip empty SSM placeholders k = torch.cat([pk, k], dim=2) v = torch.cat([pv, v], dim=2) new_kv.append((k.detach(), v.detach())) # GQA: broadcast KV heads to match Q heads if self.KVH != self.H: k = k.repeat_interleave(self.H // self.KVH, dim=1) v = v.repeat_interleave(self.H // self.KVH, dim=1) is_causal = T > 1 and past_kv is None attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal) attn = attn.transpose(1, 2).contiguous().view(B, T, self.H * self.Dh) # Output gate (Qwen3.5 attn_output_gate=True) if q_attn_gate: attn = attn * F.silu(q_gate) attn_out = F.linear(attn, w[p + "self_attn.o_proj.weight"]) hidden = res + attn_out # ── Post-attention norm + FFN ─────────────────────────────────── res = hidden hidden = _rms_norm(hidden, w[p + "post_attention_layernorm.weight"], self.eps) hidden = res + _swiglu( hidden, w[p + "mlp.gate_proj.weight"], w[p + "mlp.up_proj.weight"], w[p + "mlp.down_proj.weight"], self.mlp_chunk, ) # Final norm + LM head hidden = _rms_norm(hidden, self._norm().to(self.dtype), self.eps) logits = F.linear(hidden, self._lm_head().to(self.dtype)) return logits, new_kv # ── Text generation ─────────────────────────────────────────────────────── def _sample_next( self, logits: torch.Tensor, generated: torch.Tensor, temperature: float, top_p: float, top_k: int, repetition_penalty: float, do_sample: bool, ) -> torch.Tensor: """Sample the next token from logits.""" next_logits = logits[:, -1, :].float() if repetition_penalty != 1.0: for tok in generated[0].tolist(): v = next_logits[0, tok] next_logits[0, tok] = v / repetition_penalty if v > 0 \ else v * repetition_penalty if not do_sample: return next_logits.argmax(dim=-1, keepdim=True) if temperature > 0: next_logits = next_logits / max(temperature, 1e-6) if top_k > 0: vals, _ = torch.topk(next_logits, top_k) next_logits[next_logits < vals[:, -1:]] = float("-inf") if 0.0 < top_p < 1.0: srt_l, srt_i = torch.sort(next_logits, descending=True) cum = torch.cumsum(F.softmax(srt_l, dim=-1), dim=-1) srt_l[cum - F.softmax(srt_l, dim=-1) > top_p] = float("-inf") next_logits = torch.full_like(next_logits, float("-inf")).scatter_( 1, srt_i, srt_l ) probs = F.softmax(next_logits, dim=-1) return torch.multinomial(probs, num_samples=1) def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 200, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 0, repetition_penalty: float = 1.0, eos_token_id: Optional[int] = None, pad_token_id: int = 0, use_cache: bool = True, do_sample: bool = True, **_, ) -> torch.Tensor: """ Generate tokens with optional KV cache. With use_cache=True (default) the first forward pass processes the entire prompt, and every subsequent step processes only the one new token — dramatically faster than full-context re-computation. The layer weight cache means shard files are read from disk only once per generate() call; later steps use the in-memory copies. """ generated = input_ids.clone() past_kv = None cur_ids = input_ids with torch.inference_mode(): for _ in range(max_new_tokens): logits, new_kv = self._forward(cur_ids, past_kv if use_cache else None) if use_cache: past_kv = new_kv next_token = self._sample_next( logits, generated, temperature, top_p, top_k, repetition_penalty, do_sample, ) generated = torch.cat([generated, next_token], dim=1) cur_ids = next_token if eos_token_id is not None and (next_token == eos_token_id).all(): break return generated def generate_stream( self, input_ids: torch.Tensor, max_new_tokens: int = 200, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 0, repetition_penalty: float = 1.0, eos_token_id: Optional[int] = None, pad_token_id: int = 0, use_cache: bool = True, do_sample: bool = True, **_, ): """ Token-streaming variant — yields each new token ID as a 1-D tensor the moment it is sampled, so callers can decode and stream to clients without waiting for the full generation to finish. """ generated = input_ids.clone() past_kv = None cur_ids = input_ids with torch.inference_mode(): for _ in range(max_new_tokens): logits, new_kv = self._forward(cur_ids, past_kv if use_cache else None) if use_cache: past_kv = new_kv next_token = self._sample_next( logits, generated, temperature, top_p, top_k, repetition_penalty, do_sample, ) generated = torch.cat([generated, next_token], dim=1) cur_ids = next_token yield next_token if eos_token_id is not None and (next_token == eos_token_id).all(): break # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def _dtype_bytes(dt: torch.dtype) -> int: return {torch.float32: 4, torch.float16: 2, torch.bfloat16: 2}.get(dt, 4) # ───────────────────────────────────────────────────────────────────────────── # Self-test # ───────────────────────────────────────────────────────────────────────────── if __name__ == "__main__": import sys, time from transformers import AutoTokenizer path = sys.argv[1] if len(sys.argv) > 1 else "./model" chunk = int(sys.argv[2]) if len(sys.argv) > 2 else 75 print(f"\n=== ChunkedModel self-test model={path} chunk={chunk} MB ===\n") tok = AutoTokenizer.from_pretrained(path, local_files_only=True) m = ChunkedModel(path, chunk_mb=chunk) for prompt in [ "Say hello in one sentence.", "What is 2 + 2? Answer in one word.", ]: msgs = [{"role": "user", "content": prompt}] try: text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False) except TypeError: text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) ids = tok(text, return_tensors="pt")["input_ids"] t0 = time.time() out = m.generate(ids, max_new_tokens=40, temperature=0.7, eos_token_id=tok.eos_token_id) ans = tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True).strip() print(f"Q: {prompt}") print(f"A: {ans}") print(f" ({time.time()-t0:.1f}s)\n") print("=== done ===")