FLUX.2-dev GGUF Q4_K_M: 29 GB de VRAM en Apple Silicon

¿Por qué FLUX.2-dev GGUF Q4_K_M usa 29GB de VRAM en Apple Silicon?

El modelo FLUX.2-dev en formato cuantizado GGUF Q4_K_M debería usar teóricamente alrededor de 19GB de memoria, pero en Apple Silicon con chip M3 Pro alcanza un pico de 28.92GB. El análisis revela que los gigabytes adicionales provienen del sobrecosto de MPS, la descuantización de pesos y las transferencias de datos entre CPU y GPU en memoria unificada.

Arquitectura de carga y descarga secuencial

FLUX.2-dev utiliza el codificador de texto Mistral-Small-3.2-24B, que requiere 45GB de VRAM en precisión completa. Para Apple Silicon, se emplea una versión de 4 bits mediante MLX, consumiendo 16GB. El proceso sigue un patrón de descarga secuencial:

• Cargar el codificador MLX (~13GB).
• Codificar el prompt en embeddings en la CPU.
• Descargar el codificador y limpiar la caché.
• Cargar el modelo DiT (GGUF Q4_K_M) desde disco (~19GB).
• Inicializar el VAE y el programador en la CPU.
• Transferir todo el pipeline a MPS.

El pico teórico es el máximo de los tamaños individuales de componentes, pero la realidad difiere debido al comportamiento de la memoria unificada.

Entorno de pruebas

• macOS: 26.3.1 (Tahoe)
• Chip: Apple M3 Pro, 36GB de memoria unificada
• Python: 3.14.3
• PyTorch: 2.11.0
• diffusers: 0.38.0.dev
• mlx-lm: 0.31.1

La medición se realizó usando phys_footprint (similar al Monitor de Actividad) y métricas de MPS de PyTorch. psutil.rss no considera asignaciones de MPS.

Desglose detallado del uso de memoria

Medidas paso a paso

| Etapa | phys_footprint | Driver MPS | MPS actual |

|-------|----------------|------------|-------------|

| Base | 0.27 GB | 0 | 0 |

| Codificador MLX cargado | 13.22 GB | 0 | 0 |

| GGUF en CPU | 19.23 GB | 0 | 0 |

| pipe.to(mps) | 23.79 GB | 20.06 GB | 18.75 GB |

| Mitad de inferencia | 26.50 GB | 25.62 GB | 18.78 GB |

El pico de 28.92GB ocurre durante pipe.to("mps") debido a la presencia simultánea de copias en CPU y MPS de los datos.

Anatomía del pico de 28GB

• DiT Q4_K_M (uint8): 18.59 GB
• VAE (bf16): 0.16 GB
• Sobrecosto de pool de MPS: +1.31 GB
• IOKit/Metal graphCache: ~3.73 GB (shaders, tablas de páginas, buffers en estado de limbo)

Los pesos GGUF permanecen en uint8 en MPS; la descuantización a bf16 ocurre capa por capa durante el paso forward. Se acumulan tensors temporales bf16 (hasta 6.92GB) en caché.

Estrategias de optimización para reducir el pico de memoria

• Transferencia incremental a MPS: Mover parámetros uno a uno con gc.collect() cada 30 parámetros.

def _move_to_device_incremental(pipe, device):
    n = 0
    for attr in list(vars(pipe)):
        comp = getattr(pipe, attr, None)
        if not isinstance(comp, torch.nn.Module):
            continue
        for param in comp.parameters():
            param.data = param.data.to(device)
            n += 1
            if n % 30 == 0:
                gc.collect()
        for buf in comp.buffers():
            buf.data = buf.data.to(device)
        gc.collect()
    torch.mps.synchronize()
    torch.mps.empty_cache()

• Vaciar caché entre pasos: Usar un callback durante los pasos de inferencia con torch.mps.empty_cache() para liberar la caché de descuantización.

• Limpieza tras generación: Tras la inferencia, ejecutar gc.collect() seguido de torch.mps.empty_cache().

Resultados de las optimizaciones

| Métrica | Antes | Después | Diferencia |

|--------|--------|-------|-------|

| Pico de phys_footprint | 28.92 GB | 27.36 GB | -1.56 GB |

| Estado estable tras inferencia | 26.25 GB | 21.76 GB | -4.49 GB |

| Tiempo de generación (1 paso) | 2:04 | 2:11 | +7s |

La reducción del pico se debe a la recolección de basura incremental durante la carga; la mejora en estado estable proviene de la limpieza de caché. El sobrecosto no optimizable de MPS/IOKit permanece en ~3.5GB.

Conclusiones clave:

• Estado estable mínimo: ~24GB, pico de inferencia ~27GB.
• Costo principal: 18.59GB para pesos DiT uint8 + 3.5GB de sobrecosto MPS.
• La descuantización a bf16 añade 3–5GB de buffers temporales.
• La transferencia CPU→MPS duplica el pico por memoria unificada.
• Las optimizaciones logran -1.5GB en pico y -4.5GB en estado estable sin pérdida de calidad.

El mínimo real para FLUX.2-dev GGUF en M3 Pro es de 24+ GB de memoria unificada. Para configuraciones más pequeñas, se requiere una cuantización más agresiva o descarga a CPU.

— Editorial Team