Upload 6 files

Dockerfile

@@ -1,9 +1,14 @@
 FROM python:3.11-slim
 
 ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
     PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
     HF_HOME=/home/user/.cache/huggingface \
-    OMP_NUM_THREADS=2
+    TOKENIZERS_PARALLELISM=false \
+    MALLOC_ARENA_MAX=2
+
+# Los hilos de PyTorch los fija app.py según las vCPU reales del contenedor (variable TORCH_THREADS para forzarlos).
 
 RUN useradd -m -u 1000 user
 WORKDIR /home/user/app

README.md

@@ -96,7 +96,7 @@ Los errores usan el formato de OpenAI: `{"error": {"message", "type", "param", "
 | `MAX_TOKENS_CAP` | `1024` | Tope de `max_tokens` |
 | `DEFAULT_MAX_TOKENS` | `512` | Valor si no envías `max_tokens` |
 | `ALLOWED_ORIGINS` | `*` | Orígenes CORS (coma). Pon el dominio de tu app en producción |
-| `TORCH_THREADS` | núm. de CPUs | Hilos de PyTorch |
+| `TORCH_THREADS` | vCPU reales del contenedor | Hilos de PyTorch (se detectan por cgroup, no por los núcleos del servidor) |
 
 ## Cosas que debes saber
 

app.py

@@ -10,6 +10,7 @@ Endpoints:
 Docs interactivas: /docs
 """
 import asyncio
+import gc
 import json
 import os
 import sys
@@ -57,15 +58,46 @@ class State:
 STATE = State()
 
 
+def effective_cpus():
+    """vCPUs realmente disponibles para este contenedor. os.cpu_count() devuelve los nucleos del servidor
+    anfitrion (p. ej. 8-48) aunque el Space solo tenga 2 vCPU; usar ese numero de hilos satura la CPU
+    y hace la generacion MUCHO mas lenta."""
+    n = os.cpu_count() or 2
+    try:
+        n = min(n, len(os.sched_getaffinity(0)))
+    except Exception:  # noqa: BLE001
+        pass
+    try:  # cgroup v2
+        quota, period = open("/sys/fs/cgroup/cpu.max").read().split()[:2]
+        if quota != "max":
+            n = min(n, max(1, int(float(quota) / float(period) + 0.5)))
+    except Exception:  # noqa: BLE001
+        try:  # cgroup v1
+            quota = int(open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read())
+            period = int(open("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read())
+            if quota > 0:
+                n = min(n, max(1, int(quota / period + 0.5)))
+        except Exception:  # noqa: BLE001
+            pass
+    return max(1, n)
+
+
 def load_model():
     for fn in ("model.safetensors", "milox_config.json", "tokenizer.json"):
         path = hf_hub_download(REPO_ID, fn, token=HF_TOKEN)
     model_dir = os.path.dirname(path)
-    torch.set_num_threads(int(os.getenv("TORCH_THREADS", os.cpu_count() or 2)))
+    threads = int(os.getenv("TORCH_THREADS", effective_cpus()))
+    torch.set_num_threads(threads)
+    try:
+        torch.set_num_interop_threads(1)               # no hace falta paralelismo entre operaciones
+    except Exception:  # noqa: BLE001
+        pass
+    print(f"Hilos de PyTorch: {threads}", flush=True)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if device == "cuda" else torch.float32
     model = mm.load_milox(model_dir, device=device, dtype=dtype, int8=(USE_INT8 and device == "cpu"))
     tok = mm.ChatTokenizer(os.path.join(model_dir, "tokenizer.json"))
+    gc.collect()                                       # libera temporales de la carga
     return model, tok, device
 
 
@@ -271,7 +303,7 @@ class StopHit(Exception):
     pass
 
 
-def run_generation(ids, p, stops, stop_ids, on_text, cancel):
+def run_generation(ids, p, stops, stop_ids, on_text, cancel, live=True):
     """Se ejecuta en un hilo. Llama on_text(delta) con texto nuevo. Devuelve dict(text, completion_tokens, finish)."""
     tok, model = STATE.tok, STATE.model
     if p["seed"] is not None:
@@ -287,6 +319,8 @@ def run_generation(ids, p, stops, stop_ids, on_text, cancel):
         if cancel.is_set():
             raise Abort()
         out.append(t)
+        if not live and not stops:                    # respuesta no-stream sin stop: se decodifica una sola vez al final
+            return
         text = tok.decode(out)
         if text.endswith("\ufffd"):                   # carácter UTF-8 incompleto: espera al siguiente token
             return
@@ -336,7 +370,7 @@ async def respond(kind, ids, p, stops, stream, include_usage):
         parts = []
         try:
             async with STATE.lock:
-                res = await asyncio.to_thread(run_generation, ids, p, stops, stop_ids, parts.append, cancel)
+                res = await asyncio.to_thread(run_generation, ids, p, stops, stop_ids, parts.append, cancel, False)
         except asyncio.CancelledError:
             cancel.set()
             raise

milox_model.py

@@ -133,6 +133,22 @@ def apply_rope(q, k, cos, sin):
     return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
 
 
+class LayerCache:
+    """KV-cache preasignado por capa: escribe en un buffer fijo en vez de hacer torch.cat en cada token
+    (menos copias de memoria, menos fragmentacion de RAM). Solo se usa al generar (batch = 1)."""
+    __slots__ = ("k", "v", "n")
+
+    def __init__(self, n_kv, cap, hd, dtype, device):
+        self.k = torch.empty(1, n_kv, cap, hd, dtype=dtype, device=device)
+        self.v = torch.empty(1, n_kv, cap, hd, dtype=dtype, device=device)
+        self.n = 0
+
+
+def _past_len(past):
+    p0 = past[0]
+    return p0.n if isinstance(p0, LayerCache) else p0[0].shape[2]
+
+
 class Attention(nn.Module):
     def __init__(self, c: MiloXConfig):
         super().__init__()
@@ -152,6 +168,31 @@ class Attention(nn.Module):
         k = self.k_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
         v = self.v_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
         q, k = apply_rope(q, k, cos, sin)
+        if isinstance(past, LayerCache):
+            n0 = past.n
+            past.k[:, :, n0:n0 + T] = k
+            past.v[:, :, n0:n0 + T] = v
+            past.n = n0 + T
+            k = past.k[:, :, :past.n]
+            v = past.v[:, :, :past.n]
+            if T == 1 and n0 > 0:
+                # Decodificacion token a token: GQA sin repetir K/V. Las cabezas de q de un mismo grupo
+                # (h = j*rep + r) se apilan en la dimension de consulta; es el mismo calculo, sin copiar
+                # rep veces toda la cache en cada token y en cada capa.
+                qg = q.reshape(B, self.nkv, self.rep, self.hd)
+                y = F.scaled_dot_product_attention(qg, k, v)
+                return self.o_proj(y.reshape(B, 1, self.nh * self.hd)), past
+            if self.rep > 1:
+                k = k.repeat_interleave(self.rep, dim=1)
+                v = v.repeat_interleave(self.rep, dim=1)
+            if n0 == 0:
+                y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+            else:
+                S = k.shape[2]
+                mask = torch.ones(T, S, dtype=torch.bool, device=x.device).tril(diagonal=S - T)
+                y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+            y = y.transpose(1, 2).contiguous().view(B, T, self.nh * self.hd)
+            return self.o_proj(y), past
         if past is not None:
             k = torch.cat([past[0], k], dim=2)
             v = torch.cat([past[1], v], dim=2)
@@ -213,7 +254,7 @@ def _blk_fwd(blk, x, cos, sin):
 
 
 class MiloX(nn.Module):
-    def __init__(self, cfg: MiloXConfig):
+    def __init__(self, cfg: MiloXConfig, init_weights=True):
         super().__init__()
         self.cfg = cfg
         self.model = Backbone(cfg)
@@ -224,11 +265,12 @@ class MiloX(nn.Module):
         self.register_buffer("rope_cos", cos, persistent=False)
         self.register_buffer("rope_sin", sin, persistent=False)
         self.grad_ckpt = False
-        self.apply(self._init_weights)
-        std = 0.02 / math.sqrt(2 * cfg.n_layers)
-        for n, p in self.named_parameters():
-            if n.endswith("o_proj.weight") or n.endswith("down_proj.weight"):
-                nn.init.normal_(p, mean=0.0, std=std)
+        if init_weights:  # al cargar un modelo exportado no hace falta (los pesos se sobrescriben)
+            self.apply(self._init_weights)
+            std = 0.02 / math.sqrt(2 * cfg.n_layers)
+            for n, p in self.named_parameters():
+                if n.endswith("o_proj.weight") or n.endswith("down_proj.weight"):
+                    nn.init.normal_(p, mean=0.0, std=std)
 
     @staticmethod
     def _init_weights(m):
@@ -245,7 +287,7 @@ class MiloX(nn.Module):
 
     def forward(self, idx, targets=None, past=None, use_cache=False, full_logits=False):
         B, T = idx.shape
-        off = 0 if past is None else past[0][0].shape[2]
+        off = 0 if past is None else _past_len(past)
         assert off + T <= self.cfg.max_seq_len, f"Secuencia demasiado larga ({off + T} > {self.cfg.max_seq_len})"
         cos = self.rope_cos[off:off + T]
         sin = self.rope_sin[off:off + T]
@@ -268,30 +310,37 @@ class MiloX(nn.Module):
             loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100)
         return logits, loss, presents
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def generate(self, ids, max_new_tokens=256, temperature=0.7, top_k=40, top_p=0.9,
                  repetition_penalty=1.1, stop_ids=(), on_token=None):
         """Genera tokens con KV-cache. `ids` es una lista de ints. Devuelve la lista de tokens nuevos."""
         self.eval()
-        device = self.model.embed_tokens.weight.device
+        emb = self.model.embed_tokens.weight
+        device, dtype = emb.device, emb.dtype
         max_len = self.cfg.max_seq_len
         if len(ids) >= max_len:
             ids = ids[-(max_len - 1):]
+        cap = min(max_len, len(ids) + max_new_tokens)          # cache del tamano justo, no del maximo
+        hd = self.cfg.dim // self.cfg.n_heads
+        past = [LayerCache(self.cfg.n_kv_heads, cap, hd, dtype, device) for _ in range(self.cfg.n_layers)]
         cur = torch.tensor([ids], dtype=torch.long, device=device)
-        past, out, seen = None, [], list(ids[-256:])
+        out = []
         stop_ids = set(stop_ids)
-        block = [i for i in range(N_SPECIAL) if i not in stop_ids]
+        block = torch.tensor([i for i in range(N_SPECIAL) if i not in stop_ids], dtype=torch.long, device=device)
+        # tokens "vistos" para la penalizacion (ultimos 256 del prompt + todo lo generado), como mascara fija
+        seen = torch.zeros(self.cfg.vocab_size, dtype=torch.bool, device=device)
+        if ids:
+            seen[torch.tensor(ids[-256:], dtype=torch.long, device=device)] = True
         for _ in range(max_new_tokens):
-            total = cur.shape[1] + (0 if past is None else past[0][0].shape[2])
+            total = cur.shape[1] + _past_len(past)
             if total > max_len:
                 break
             logits, _, past = self(cur, past=past, use_cache=True)
             logits = logits[0, -1].float()
             logits[block] = float("-inf")
-            if repetition_penalty != 1.0 and seen:
-                sidx = torch.tensor(sorted(set(seen)), device=device)
-                sc = logits[sidx]
-                logits[sidx] = torch.where(sc > 0, sc / repetition_penalty, sc * repetition_penalty)
+            if repetition_penalty != 1.0:
+                logits = torch.where(seen, torch.where(logits > 0, logits / repetition_penalty,
+                                                       logits * repetition_penalty), logits)
             if temperature <= 0:
                 nxt = int(torch.argmax(logits))
             else:
@@ -310,7 +359,7 @@ class MiloX(nn.Module):
             if nxt in stop_ids:
                 break
             out.append(nxt)
-            seen.append(nxt)
+            seen[nxt] = True
             if on_token is not None:
                 on_token(nxt)
             cur = torch.tensor([[nxt]], dtype=torch.long, device=device)
@@ -410,11 +459,24 @@ class ChatSession:
 def load_milox(path, device="cpu", dtype=torch.float32, int8=False):
     from safetensors.torch import load_file
     cfg = MiloXConfig.from_json(os.path.join(path, "milox_config.json"))
-    model = MiloX(cfg)
     sd = load_file(os.path.join(path, "model.safetensors"))
     if cfg.tie_embeddings and "lm_head.weight" not in sd:
         sd["lm_head.weight"] = sd["model.embed_tokens.weight"]
-    model.load_state_dict(sd, strict=True)
+    try:
+        # Carga eficiente en RAM: el modelo se crea "vacio" (meta) y adopta directamente los tensores del
+        # archivo; evita inicializar pesos al azar y tener 2-3 copias del modelo a la vez durante el arranque.
+        with torch.device("meta"):
+            model = MiloX(cfg, init_weights=False)
+        model.load_state_dict(sd, strict=True, assign=True)
+        cos, sin = build_rope(cfg.dim // cfg.n_heads, cfg.max_seq_len, cfg.rope_theta)
+        model.rope_cos, model.rope_sin = cos, sin
+        if cfg.tie_embeddings:
+            model.lm_head.weight = model.model.embed_tokens.weight
+    except Exception as e:  # torch antiguo u otro problema: camino clasico
+        print("Carga eficiente no disponible, se usa la carga clasica:", e)
+        model = MiloX(cfg, init_weights=False)
+        model.load_state_dict(sd, strict=True)
+    del sd
     model.eval()
     if int8:
         try: