fix(offload): aparcar experto inactivo en disco (mmap) en vez de copiar 14.6GB a RAM (v0.3.14)
This commit is contained in:
+4
-1
@@ -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
@@ -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
@@ -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 = [
|
||||
|
||||
Reference in New Issue
Block a user