Inferencia distribuida

La inferencia distribuida divide la carga de trabajo entre varias GPUs. Es una técnica útil para ajustar modelos más grandes en la memoria y puede procesar múltiples prompts para lograr un mayor rendimiento.

Esta guía te mostrará cómo usar Accelerate y PyTorch Distributed para la inferencia distribuida.

Accelerate

Accelerate es una biblioteca diseñada para simplificar la inferencia y el entrenamiento en múltiples aceleradores al gestionar la configuración, lo que permite a los usuarios centrarse en su código de PyTorch.

Instale Accelerate con el siguiente comando.

uv pip install accelerate

Inicie una clase accelerate.PartialState en un archivo de Python para crear un entorno distribuido. La clase accelerate.PartialState gestiona la administración de procesos, el control y la distribución de dispositivos, y la coordinación de procesos.

Mueva el DiffusionPipeline a accelerate.PartialState.device para asignar una GPU a cada proceso.

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

Use la utilidad split_between_processes como un gestor de contexto para distribuir automáticamente los prompts entre el número de procesos.

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

Ejecuta accelerate launch para ejecutar el script y usa el argumento --num_processes para establecer el número de GPUs a utilizar.

accelerate launch run_distributed.py --num_processes=2

Consulte este script minimalista para ejecutar inferencia en varias GPU. Para obtener más información, consulte la guía Distributed Inference with 🤗 Accelerate.

PyTorch Distributed

PyTorch DistributedDataParallel habilita el data parallelism, que replica el mismo modelo en cada dispositivo para procesar distintos lotes de datos en paralelo.

Importe torch.distributed y torch.multiprocessing en un archivo de Python para configurar el grupo de procesos distribuidos y generar los procesos para la inferencia en cada GPU.

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16,
)

Cree una función para la inferencia con init_process_group. Este método crea un entorno distribuido con el tipo de backend, el rank del proceso actual y el world_size o número de procesos participantes (por ejemplo, 2 GPUs serían world_size=2).

Mueva el pipeline a rank y use get_rank para asignar una GPU a cada proceso. Cada proceso maneja un prompt diferente.

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    pipeline.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")

Use mp.spawn para crear el número de procesos definido en world_size.

def main():
    world_size = 2
    mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()

Llame a torchrun para ejecutar el script de inferencia y use el argumento --nproc_per_node para establecer el número de GPUs a utilizar.

torchrun --nproc_per_node=2 run_distributed.py

device_map

El argumento device_map habilita la inferencia distribuida colocando automáticamente los componentes del modelo en GPUs separadas. Esto es especialmente útil cuando un modelo no cabe en una sola GPU. Puedes usar device_map para cargar y descargar selectivamente los componentes del modelo necesarios en una etapa dada, como se muestra en el ejemplo siguiente (asume que hay dos GPUs disponibles).

Establezca device_map="balanced" para distribuir de manera uniforme los codificadores de texto en todas las GPU disponibles. Puede usar el argumento max_memory para asignar una cantidad máxima de memoria para cada codificador de texto. No cargue ningún otro componente de la pipeline para evitar el uso de memoria.

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

Una vez que se calculan las incrustaciones de texto, elimínalas de la GPU para hacer espacio para el transformador de difusión.

import gc 

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

Establezca device_map="auto" para distribuir automáticamente el modelo en las dos GPU. Esta estrategia coloca el modelo primero en el dispositivo más rápido y, si es necesario, lo coloca en un dispositivo más lento, como una CPU o un disco duro. El inconveniente de almacenar los parámetros del modelo en dispositivos más lentos es una latencia de inferencia más alta.

from diffusers import AutoModel
import torch 

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)

Ejecuta pipeline.hf_device_map para ver cómo se distribuyen los distintos modelos entre los dispositivos. Esto es útil para hacer un seguimiento de la ubicación de los modelos en los dispositivos. También puedes llamar a hf_device_map en el modelo transformer para ver cómo se distribuye.

Agregue el modelo transformer a la pipeline y establezca output_type="latent" para generar los latentes.

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

Elimine el pipeline y el transformer de la memoria y cargue un VAE para decodificar los latentes. El VAE suele ser lo suficientemente pequeño como para cargarse en un solo dispositivo.

import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", dtype=torch.bfloat16).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")

Al cargar y descargar selectivamente los modelos necesarios en cada etapa y distribuir los modelos más grandes entre varias GPU, es posible ejecutar inferencia con modelos grandes en GPU de consumo.

Paralelismo de contexto

Paralelismo de contexto divide las secuencias de entrada entre varias GPU para reducir el uso de memoria. Cada GPU procesa su propia porción de la secuencia.

Use set_attention_backend() para cambiar a un backend de atención más optimizado. Consulte esta tabla para obtener una lista completa de los backends disponibles.

La mayoría de los backends de atención son compatibles con la paralelización de contexto. Abra un issue si un backend no es compatible.

Ring Attention

Las representaciones de clave (K) y valor (V) se comunican entre dispositivos mediante Ring Attention. Esto garantiza que cada división vea las K/V de todos los demás tokens. Cada GPU calcula la atención para su K/V local y la pasa a la siguiente GPU en el anillo. Ninguna GPU individual almacena la secuencia completa, lo que reduce la latencia de comunicación.

Pase un ContextParallelConfig al argumento parallel_config del modelo transformer. La configuración admite el argumento ring_degree, que determina cuántos dispositivos se usarán para Ring Attention.

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save(f"output.png")

    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

El script anterior debe ejecutarse con un lanzador distribuido, como torchrun, compatible con PyTorch. --nproc-per-node se establece en el número de GPUs disponibles.

torchrun --nproc-per-node 2 above_script.py

Ulysses Attention

Ulysses Attention divide una secuencia entre varias GPU y realiza una comunicación all-to-all (cada dispositivo envía/recibe datos a cada otro dispositivo). Cada GPU termina con todos los tokens solo para un subconjunto de cabezas de atención. Cada GPU calcula la atención localmente sobre todos los tokens de su cabeza y luego realiza otro all-to-all para reagrupar los resultados por tokens para la siguiente capa.

ContextParallelConfig admite Ulysses Attention a través del argumento ulysses_degree. Este determina cuántos dispositivos se usarán para Ulysses Attention.

Pase ContextParallelConfig a enable_parallelism().

# Depending on the number of GPUs available.
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

Atención unificada

Unified Sequence Parallelism combina Ring Attention y Ulysses Attention en un único enfoque para el procesamiento eficiente de secuencias largas. Primero aplica la comunicación all-to-all de Ulysses para redistribuir las cabezas y los tokens de la secuencia, luego utiliza Ring Attention para procesar los datos redistribuidos y, finalmente, invierte el all-to-all para restaurar el diseño original.

Este enfoque híbrido aprovecha las fortalezas de ambos métodos:

• Ulysses Attention paraleliza de manera eficiente entre las cabezas de atención
• Ring Attention maneja secuencias muy largas con un sobrecoste de memoria mínimo
• Juntas, permiten la paralelización 2D a lo largo de las dimensiones de cabezas y de secuencia

ContextParallelConfig admite Unified Attention especificando tanto ulysses_degree como ring_degree. El número total de dispositivos utilizados es ulysses_degree * ring_degree, dispuestos en una cuadrícula 2D donde los grupos Ulysses y Ring son ortogonales (sin solapamiento). Pase el ContextParallelConfig con tanto ulysses_degree como ring_degree establecidos en un valor mayor que 1 a enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

Unified Attention debe utilizarse cuando hay suficientes dispositivos para organizarlos en una cuadrícula 2D (al menos 4 dispositivos).

Ejecutamos un benchmark con Ulysess, Ring y Unified Attention mediante este script en un nodo de 4 GPUs H100. Los resultados se resumen a continuación:

De la tabla anterior, es claro que Ulysses ofrece un mayor rendimiento, pero el número de dispositivos que puede utilizar sigue limitado al número de cabezas de atención, una limitación que se resuelve con la atención unificada.

Ulysses Anything Attention

El mecanismo de atención Ulysses por defecto requiere que la longitud de secuencia de los estados ocultos sea divisible por el número de dispositivos. Esto impone limitaciones significativas a la aplicación práctica de la atención Ulysses. Ulysses Anything Attention es una variante de la atención Ulysses que admite longitudes de secuencia arbitrarias y un número arbitrario de cabezas de atención, lo que mejora la versatilidad de la atención Ulysses en su uso práctico.

ContextParallelConfig admite Ulysses Anything Attention especificando tanto ulysses_degree como ulysses_anything. Tenga en cuenta que Ulysses Anything Attention no es compatible actualmente con Unified Attention. Pase el ContextParallelConfig con ulysses_degree establecido en un valor mayor que 1 y ulysses_anything=True a enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

Para evitar múltiples sincronizaciones forzadas de CUDA causadas por las transferencias H2D y D2H, agregue el backend gloo en init_process_group. Esto reducirá significativamente la latencia de comunicación.

Ejecutamos un benchmark de FLUX.1-dev con Ulysses, Ring, Unified Attention y Ulysses Anything Attention con este script en un nodo de 4 GPUs L20. Los resultados se resumen a continuación:

De la tabla anterior, es claro que Ulysses Anything Attention ofrece una mejor compatibilidad con longitudes de secuencia arbitrarias, manteniendo el mismo rendimiento que la Ulysses Attention estándar.

Ring Anything Attention

El Ring Attention por defecto requiere que la longitud de secuencia de los estados ocultos sea divisible uniformemente entre el grado del anillo. Ring Anything Attention es una variante de Ring Attention que admite longitudes de secuencia arbitrarias (no uniformemente divisibles). Rellena el KV local de cada rank hasta la longitud de secuencia máxima global, realiza un all-gather del búfer KV relleno y lo recorta a la longitud real de cada rank antes de ejecutar la atención.

Ring Anything Attention no es compatible con Unified Attention. Establezca ring_degree > 1 y ring_anything=True para habilitar Ring Anything Attention.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Agregue el backend gloo a init_process_group para evitar múltiples sincronizaciones forzadas de CUDA derivadas de las transferencias H2D y D2H.

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

Ring Anything Attention solo admite inferencia por el momento y no soporta máscaras de atención distintas de None. attn_mask debe ser None.

Vea los benchmarks de FLUX.1-dev a continuación en un nodo de 4 GPUs RTX 4090 (48GB).

A partir de la tabla anterior, Ring Anything Attention ofrece compatibilidad con longitudes de secuencia arbitrarias mientras mantiene un rendimiento comparable al de Ring Attention estándar. Para más detalles sobre la motivación y las compensaciones de Ring Anything Attention, consulte este comentario.

parallel_config

Pasa parallel_config durante la inicialización del modelo para habilitar el paralelismo de contexto.

CKPT_ID = "black-forest-labs/FLUX.1-dev"

cp_config = ContextParallelConfig(ring_degree=2)
transformer = AutoModel.from_pretrained(
    CKPT_ID, 
    subfolder="transformer", 
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)

Paralelismo de tensores

Tensor parallelism distribuye las matrices de pesos de un modelo entre varios dispositivos. Cada dispositivo mantiene una porción por columnas ("colwise") o por filas ("rowwise") de cada capa, calcula un resultado parcial y un AllReduce/AllGather en el límite de la capa reconstruye la salida completa. A diferencia del paralelismo de contexto, reduce la memoria de pesos por dispositivo, lo cual es útil para modelos que no caben en un solo dispositivo.

Pase un TensorParallelConfig al argumento parallel_config de from_pretrained() del modelo. tp_degree es el número de dispositivos en los que se dividirá el modelo y debe dividir el número de cabezas de atención del modelo. El modelo debe definir un _tp_plan (un mapeo plano de glob de nombres de módulos a un estilo "colwise"/"rowwise").

Al cargar de esta manera, el checkpoint se fragmenta mientras se lee: cada rank lee solo su propia porción de cada peso fragmentado y lo coloca directamente en su propio dispositivo. Nunca se materializa nada a tamaño completo, por lo que la memoria por rank disminuye a medida que aumenta tp_degree.

En comparación con cargar el modelo completo y luego llamar a enable_parallelism(), esta opción carga más rápido y usa menos memoria de CPU por rank, y la diferencia aumenta a medida que tp_degree sube. Los números a continuación corresponden a black-forest-labs/FLUX.2-dev (solo el transformer, 32B parámetros, bf16) en 4x A10G (23GB).

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto `cuda:rank`.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, torch_dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

torchrun --nproc-per-node 4 tensor_parallel_flux.py

tp_degree se toma de world_size anterior, por lo que --nproc-per-node 4 divide el transformer en 4 dispositivos.

La carga con una parallel_config de paralelismo de tensores aún no es compatible con device_map, quantization_config, low_cpu_mem_usage=False, use_flashpack=True o pesos que no sean de safetensors; cada uno de estos parámetros genera un error en lugar de volver silenciosamente a la carga del punto de control completo.

La combinación de paralelismo de tensores con cuantización, offloading o adaptadores LoRA tampoco está soportada todavía, por lo que esos lanzan ValueError sin importar cómo se fragmente el modelo.

Recomendamos cargar los modelos de paralelismo de tensores de esta manera. También puedes fragmentar un modelo que ya está en memoria llamando a enable_parallelism() con la misma configuración, pero eso carga todo primero y lo vuelve a fragmentar, por lo que consume la memoria completa del checkpoint en cada rank.

No se admite aún el guardado de un modelo paralelo de tensores, y save_pretrained() lanza una excepción en uno. Guarda el modelo antes de fragmentarlo.

Escribir un plan de paralelismo de tensores

La paralelización de tensores solo funciona en modelos que definen un _tp_plan, un atributo de clase plano que mapea patrones de nombres de módulos a un estilo de particionamiento. Escribir uno consiste principalmente en emparejar cada proyección que expande la dimensión oculta con la proyección que la contrae de nuevo.

Cada clave puede contener como máximo un *, y el prefijo anterior debe resolverse en un nn.ModuleList para que una única entrada cubra todos los bloques. Una clave sin * se aplica al propio modelo. Las rutas son relativas al modelo.

Colwise y rowwise

Acompánelos siempre en ese orden. Una proyección "colwise" deja su salida fragmentada, la proyección "rowwise" siguiente consume ese fragmento directamente y un único AllReduce en el límite del bloque reconstruye el resultado. Fragmentar el par de cualquier otra manera obliga a realizar una agregación intermedia y comunica mucho más.

En el caso de la atención, esto significa que cada rank posee un subconjunto de cabezas, por lo que tp_degree debe dividir el número de cabezas. Las duplicaciones del stream del codificador (add_q_proj, to_add_out, ff_context) siguen el mismo patrón que sus equivalentes en el stream de imagen.

Proyecciones fusionadas

Cuando un Linear empaqueta varios tensores lógicos a lo largo de la dimensión que se fragmenta, las rebanadas simples "colwise"/"rowwise" cortan directamente a través de la concatenación y desalinean las piezas. En su lugar, usa PackedColwiseParallel/PackedRowwiseParallel, que fragmenta cada bloque empaquetado de forma independiente.

# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel

blocks es una lista de enteros proporcionales cuya suma divide la dimensión empaquetada: [1, 1] para una proyección gate+up de SwiGLU con mitades iguales, o [1, 1, 1, 3, 3] para una proyección fusionada Q+K+V+gate+up con mlp_ratio=3.

"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),

Cuando los tamaños de bloque solo se conocen desde la configuración, omita el argumento y almacene los tamaños absolutos en Linear durante __init__ en su lugar, como _tp_packed_col_blocks o _tp_packed_row_blocks.

# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]

# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),

Ambas formas aparecen en un solo plan en transformer_flux2.py: el par SwiGLU de doble flujo pasa blocks literalmente, mientras que la proyección QKV+MLP de flujo único fusionada almacena sus tamaños en el Linear de Flux2ParallelSelfAttention.__init__ porque dependen de mlp_ratio y mlp_mult_factor.

Qué dejar fuera

Todo lo que no esté presente en el plan se mantiene replicado en cada rank, que es la opción correcta para las capas de normalización, la modulación AdaLN (img_mod/txt_mod), los embeddings de parches y de texto, y el norm_out/proj_out final. Estos son pequeños, por lo que fragmentarlos ahorra poca memoria a cambio de añadir comunicación.

Restricciones y verificación

• tp_degree debe dividir a config.num_attention_heads.
• Cada bloque empaquetado debe ser divisible individualmente por tp_degree, no solo su suma.

Ambos se validan mediante from_pretrained() y enable_parallelism() antes de que se cargue o se distribuya cualquier peso.

Valide un nuevo plan numéricamente en lugar de a simple vista: genere con una semilla fija en un solo dispositivo y luego bajo paralelismo de tensores, y compare las salidas. Un "colwise"/"rowwise" mal colocado suele ejecutarse y producir una imagen plausible pero incorrecta.

Comience con un plan existente para una arquitectura similar. QwenImageTransformer2DModel está completamente no fusionado y cada entrada es un simple "colwise"/"rowwise", FluxTransformer2DModel añade una única proyección row-wise empaquetada, y Flux2Transformer2DModel cubre ambos estilos empaquetados; consulte transformer_qwenimage.py, transformer_flux.py y transformer_flux2.py respectivamente.

Elegir una estrategia

Las estrategias anteriores resuelven problemas distintos, y la pregunta útil no es cuál es la más rápida en abstracto, sino de qué recurso estás quedando corto.

Algunas orientaciones prácticas:

• Alto rendimiento con muchos prompts, el modelo ya cabe en memoria. Utilice paralelismo de datos. Es la única estrategia que escala el rendimiento de forma lineal sin modificar el modelo y que no afecta a la latencia de un solo prompt.
• Los pesos de un solo componente no caben. Recurre primero al paralelismo de tensores, ya que reduce tanto la memoria como la latencia. Comunica en cada límite de bloque, por lo que requiere un interconexión rápida como NVLink; sobre PCIe, ese tráfico por capa puede superar el cómputo que ahorra, y device_map se convierte en la mejor opción. device_map también maneja el caso en que los componentes son individualmente adecuados pero colectivamente demasiado grandes.
• El problema son las activaciones, no los pesos. Este es el régimen de secuencias largas: imágenes grandes, muchos fotogramas, y el paralelismo de contexto es la respuesta directa. Elija un backend de los benchmarks de Ulysses/Ring anteriores.
• Tanto los pesos como la secuencia son demasiado grandes. Combina la paralelización de tensores y de contexto. TensorParallelConfig acepta un argumento mesh para que ambos puedan compartir una sola malla de dispositivos. Esta combinación aún no está validada, por lo que debe considerarse experimental.