fix(offload): aparcar experto inactivo en disco (mmap) en vez de copiar 14.6GB a RAM (v0.3.14)

This commit is contained in:
neuregex
2026-06-05 17:55:01 -06:00
parent 171453bc42
commit a56bc058dc
3 changed files with 20 additions and 8 deletions
+4 -1
View File
@@ -211,7 +211,10 @@ def _load_prefp8(model_dir, subfolder):
for shard in sorted(glob.glob(os.path.join(model_dir, subfolder, "*.safetensors"))):
sd.update(_safe_load(shard, keep))
missing, unexpected = m.load_state_dict(sd, strict=False, assign=True)
m._bernini_keep = keep # mantiene vivos mmap+file mientras viva el modelo
m._bernini_keep = keep # mantiene vivos mmap+file mientras viva el modelo
m._bernini_cpu_sd = sd # pesos read-only mmap (CPU, perezosos): para volver a
# "aparcar" el experto en disco al hacer offload (to_idle)
# sin copiar 14.6GB a RAM.
if missing:
print(f"[BerniniR] fp8 load {subfolder}: {len(missing)} missing (ej: {missing[:2]})")
if unexpected:
+15 -6
View File
@@ -119,12 +119,21 @@ class BerniniExpert(torch.nn.Module):
self.t.scale_shift_table.data = self.t.scale_shift_table.data.to(device)
def to_idle(self):
"""Manda TODO el experto (residentes + swap + no-bloque) a CPU (inactivo). Si YA
está en CPU no hace nada: llamar .to('cpu') sobre los tensores fp8 de SOLO-LECTURA
(respaldados por mmap) dispara un 'access violation' -> torch recrea el Parameter
sobre memoria no-escribible. Solo movemos de verdad cuando viene de GPU (ahí
.to('cpu') crea copias escribibles nuevas, sin tocar el mmap original)."""
if self.device.type != "cpu":
"""Aparca el experto FUERA de la GPU sin meterlo en RAM. Si ya está en CPU, nada.
Si está en GPU, RE-ASIGNA los pesos read-only mmap originales (CPU, perezosos,
~0 RAM) y libera la VRAM -> el experto inactivo vuelve a 'vivir en disco'.
Antes hacía `self.t.to('cpu')`, que copia ~14.6GB a RAM escribible y, con la RAM
llena, revienta (access violation / WinError 1450/1455). Re-asignar el state_dict
mmap (assign=True) NO copia: solo reapunta los params al mmap y suelta los tensores
GPU (se liberan con empty_cache en el sampler). Al reactivarse, to_active vuelve a
streamear los pesos a GPU desde el mmap (file cache, no commit de RAM)."""
if self.device.type == "cpu":
return
cpu_sd = getattr(self.t, "_bernini_cpu_sd", None)
if cpu_sd is not None:
self.t.load_state_dict(cpu_sd, strict=False, assign=True) # re-apunta a mmap; libera GPU
else:
self.t.to("cpu")
@property
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui-berninir"
description = "ByteDance Bernini-R (Wan2.2-T2V-A14B + source-id RoPE + APG multi-condition guidance) for ComfyUI: t2v/t2i, image & video editing, reference-to-video."
version = "0.3.13"
version = "0.3.14"
license = { file = "LICENSE" }
requires-python = ">=3.10"
dependencies = [