Al trabajar con modelos de generación de imágenes como Stable Diffusion XL (SDXL), es común encontrarse con el error CUDA out of memory al intentar producir visuales en resoluciones elevadas (como 1024x1024 o superiores). Una GPU con 8GB o 12GB de VRAM suele ser insuficiente para estas cargas de trabajo, lo que obliga a los usuarios a invertir en hardware de alta gama. Sin embargo, existe una alternativa más eficiente: la coordinación de múltiples GPUs de consumo. Este artículo documenta la implementación de dos estrategias fundamentales para sortear las limitaciones de memoria: 1. Balanceo de carga: Distribuir solicitudes de inferencia concurrentes a través de varios dispositivos para maximizar el rendimiento. 2. Fragmentación de modelo (Model Sharding): Dividir la estructura del modelo para que varias GPUs compartan la carga de memoria de una única inferencia compleja. Combinando ambas técnicas, es posible ejecutar inferencias que requerirían más de 16GB de VRAM en una sola tarjeta, distribuyendo el trabajo entre, por ejemplo, dos GPUs de 8GB. ### Arquitectura: Concurrencia vs. Paralelismo de Modelo

Balanceo de carga: Ejecución paralela de tareas

El balanceo de carga opera de manera similar a un sistema de múltiples cajas en un supermercado. En lugar de enviar todas las solicitudes de generación a una única GPU (donde se formarían cuellos de botella y se agotaría la memoria), un despachador asigna cada nueva solicitud a la GPU que tenga menor ocupación en ese momento. Esto incrementa drásticamente la capacidad de procesamiento por lotes (batch processing). #### Fragmentación de modelo: Distribución de pesos

A diferencia del balanceo de carga, donde cada GPU procesa una imagen distinta, la fragmentación divide los pesos de la red neuronal. Si SDXL es demasiado pesado para una tarjeta, se pueden cargar las primeras capas del modelo en la GPU 0 y las capas subsiguientes en la GPU 1. Durante la inferencia, los tensores de activación se transfieren de la GPU 0 a la GPU 1 para completar el cálculo. Esto permite ejecutar modelos enormes en hardware modesto, a costa de cierta latencia por la comunicación entre dispositivos. ### Implementación Práctica

La modificación del sistema se centra en los controladores de inferencia y carga de modelos. #### Preparación del Entorno

Se requiere un entorno con múltiples GPUs NVIDIA y las siguientes dependencias: ```

Instalación de PyTorch con soporte CUDA 11.8

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

Frameworks de inferencia y distribución

pip install diffusers transformers accelerate streamlit Pillow

Verificación del hardware disponible: ```
import torch

assert torch.cuda.is_available(), "CUDA no está disponible"
print(f"Dispositivos detectados: {torch.cuda.device_count()}")
for dev_idx in range(torch.cuda.device_count()):
    mem_total_gb = torch.cuda.get_device_properties(dev_idx).total_memory / 1e9
    print(f"  GPU {dev_idx}: {torch.cuda.get_device_name(dev_idx)} | VRAM: {mem_total_gb:.2f} GB")

Estrategia 1: Despachador de Carga Multidispositivo

Implementarmeos un enrutador que instanciará un pipeline de Diffusers en cada GPU y asignará las tareas mediante una política Round-Robin. ```

gpu_router.py

import torch from diffusers import StableDiffusionXLPipeline import threading

class DeviceLoadBalancer: """Enrutador de inferencia para múltiples GPUs usando Round-Robin."""

def __init__(self, ckpt_dir, target_gpus=None):
    self.ckpt_dir = ckpt_dir
    self.target_gpus = target_gpus or list(range(torch.cuda.device_count()))
    
    if not self.target_gpus:
        raise ValueError("No se especificaron GPUs válidas.")
        
    self.inference_units = []
    self.unit_locks = []
    self.robin_idx = 0
    self.idx_mutex = threading.Lock()
    
    print(f"Inicializando pipelines en GPUs {self.target_gpus}...")
    for dev_id in self.target_gpus:
        dev_str = f"cuda:{dev_id}"
        pipe = StableDiffusionXLPipeline.from_pretrained(
            self.ckpt_dir,
            torch_dtype=torch.float16,
            variant="fp16"
        ).to(dev_str)
        
        # Optimización de atención (si xformers está instalado)
        try:
            pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            pass
            
        self.inference_units.append({"pipe": pipe, "dev_id": dev_id})
        self.unit_locks.append(threading.Lock())
        print(f"Pipeline desplegado en {dev_str} ({torch.cuda.get_device_name(dev_id)})")
        
def _acquire_device(self):
    """Obtiene la siguiente unidad de inferencia disponible de forma segura."""
    with self.idx_mutex:
        unit = self.inference_units[self.robin_idx]
        self.robin_idx = (self.robin_idx + 1) % len(self.inference_units)
    return unit
    
def render_visual(self, txt_prompt, avoid_prompt=None, **gen_kwargs):
    """Ejecuta la inferencia en el dispositivo seleccionado con exclusión mutua."""
    unit = self._acquire_device()
    pipe_instance = unit["pipe"]
    target_dev = unit["dev_id"]
    
    lock_index = self.target_gpus.index(target_dev)
    with self.unit_locks[lock_index]:
        print(f"[Router] Solicitud asignada a GPU {target_dev}")
        config = {"num_inference_steps": 30, "guidance_scale": 7.5}
        config.update(gen_kwargs)
        
        result = pipe_instance(
            prompt=txt_prompt,
            negative_prompt=avoid_prompt,
            **config
        ).images[0]
        
    return result

Integración con la interfaz Streamlit: ```
# ui_app.py
import streamlit as st
import torch
from gpu_router import DeviceLoadBalancer

st.set_page_config(page_title="Estudio Visual Multi-GPU", layout="wide")

# Monitorio de hardware en la barra lateral
with st.sidebar:
    st.subheader("Hardware Disponible")
    total_gpus = torch.cuda.device_count()
    if total_gpus > 1:
        st.success(f"Modo multi-GPU activo ({total_gpus} dispositivos)")
    else:
        st.warning("Operando en modo mono-GPU")

@st.cache_resource
def init_load_balancer():
    model_path = "ruta/al/modelo/sdxl-base-1.0"
    return DeviceLoadBalancer(model_path)

engine = init_load_balancer()

def generate_artwork(concept, restrictions, iters, scale):
    if engine:
        return engine.render_visual(
            txt_prompt=concept,
            avoid_prompt=restrictions,
            num_inference_steps=iters,
            guidance_scale=scale
        )
    st.error("Motor de inferencia no inicializado")
    return None

Estrategia 2: Despliegue Fragmentado con Accelerate

Cuando el objetivo es generar una única imagen de resoluciones extremas (ej. 2048x2048), el balanceo de carga no sirve. Debemos particionar el modelo utilizando la librería accelerate de Hugging Face. ```

sharded_inference.py

from accelerate import init_empty_weights, load_checkpoint_and_dispatch from diffusers import StableDiffusionXLPipeline import torch

def deploy_sharded_sdxl(ckpt_dir, mapping_strategy="auto"): """Carga un pipeline SDXL distribuyendo sus capas entre múltiples dispositivos.""" print(f"Iniciando mapeo de modelo con estrategia '{mapping_strategy}'...")

# Definir arquitectura sin asignar memoria
with init_empty_weights():
    generator = StableDiffusionXLPipeline.from_pretrained(
        ckpt_dir,
        torch_dtype=torch.float16,
        variant="fp16"
    )

# Despachar pesos a las GPUs según el mapa definido
generator = load_checkpoint_and_dispatch(
    generator,
    ckpt_dir,
    device_map=mapping_strategy,
    # Limitar la asignación por GPU para evitar OOM durante la carga
    max_memory={idx: "10GiB" for idx in range(torch.cuda.device_count())},
    offload_folder="cpu_spill", # Respaldo en disco si la VRAM y RAM se saturan
    dtype=torch.float16
)

print("Despliegue fragmentado completado.")
return generator

Script de prueba

if name == "main": sdxl_path = "ruta/al/modelo/sdxl-base-1.0" sharded_pipe = deploy_sharded_sdxl(sdxl_path)

output = sharded_pipe("Futuristic metropolis floating in a nebula", num_inference_steps=30).images[0]
output.save("sharded_output.png")

Para alternar entre modos en la aplicación de usuario, se puede añadir un selector en la interfaz: ```
# Extensión de ui_app.py
with st.sidebar:
    execution_policy = st.radio(
        "Política de Inferencia:",
        ["Concurrencia (Balanceo de Carga)", "Alta Resolución (Fragmentación)"]
    )

@st.cache_resource
def get_engine(policy, path):
    if policy == "Concurrencia (Balanceo de Carga)":
        from gpu_router import DeviceLoadBalancer
        return DeviceLoadBalancer(path)
    else:
        from sharded_inference import deploy_sharded_sdxl
        return deploy_sharded_sdxl(path)

current_engine = get_engine(execution_policy, "ruta/al/modelo")

Perfilamiento y Análisis de Rendimiento

Las pruebas se ejecutaron en un equipo con dos NVIDIA RTX 3060 (12GB VRAM cada una), utilizando SDXL 1.0 Base a 1024x1024 y 30 pasos de muestreo. | Métrica / Escenario | GPU Única | Balanceo de Carga (2x GPU) | Fragmentación (2x GPU) | |---|---|---|---| | Tiempo por imagen | ~22 seg | ~22 seg (Ejecución aislada) | ~28 seg (Latencia inter-GPU) | | Lote de 5 imágenes | ~110 seg | ~65 seg | N/A (Diseñado para una sola imagen grande) | | Pico de VRAM por GPU | 9.5 GB | 9.5 GB (Por dispositivo) | 6.0 GB (Por dispositivo) | | Resolución máxima viable | 1024x1024 | 1024x1024 (Por solicitud) | 1536x1536 |

El balanceo de carga mejora el rendimiento en escenarios de procesamiento por lotes casi de forma lineal (1.7x de aceleración en este caso), sufriendo pequeñas pérdidas por la sobrecarga de la interfaz de usuario y la sincronización. Por su parte, la fragmentación reduce con éxito el consumo por tarjeta, permitiendo resoluciones inalcanzables para una sola GPU de 12GB, a cambio de un incremento del 27% en el tiempo de inferencia causado por el cuello de botella en la transferencia de activaciones a través del PCIe.