# 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](https://huggingface.co/docs/accelerate/index) y [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) 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. ```bash uv pip install accelerate ``` Inicie una clase [accelerate.PartialState](https://huggingface.co/docs/accelerate/main/en/package_reference/state#accelerate.PartialState) en un archivo de Python para crear un entorno distribuido. La clase [accelerate.PartialState](https://huggingface.co/docs/accelerate/main/en/package_reference/state#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](/docs/diffusers/main/es/api/pipelines/overview#diffusers.DiffusionPipeline) a `accelerate.PartialState.device` para asignar una GPU a cada proceso. ```py 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](https://huggingface.co/docs/accelerate/main/en/package_reference/state#accelerate.PartialState.split_between_processes) como un gestor de contexto para distribuir automáticamente los prompts entre el número de procesos. ```py 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. ```bash accelerate launch run_distributed.py --num_processes=2 ``` > [!TIP] > Consulte este [script](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) minimalista para ejecutar inferencia en varias GPU. Para obtener más información, consulte la guía [Distributed Inference with 🤗 Accelerate](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate). ## PyTorch Distributed PyTorch [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) habilita el [data parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=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. ```py 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](https://pytorch.org/docs/stable/distributed.html?highlight=init_process_group#torch.distributed.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. ```py 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](https://pytorch.org/docs/stable/multiprocessing.html#torch.multiprocessing.spawn) para crear el número de procesos definido en `world_size`. ```py 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. ```bash 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 uniformemente 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. ```py 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. ```py 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. ```py from diffusers import AutoModel import torch transformer = AutoModel.from_pretrained( "black-forest-labs/FLUX.1-dev", subfolder="transformer", device_map="auto", dtype=torch.bfloat16 ) ``` > [!TIP] > 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. ```py 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. ```py 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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism) 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()](/docs/diffusers/main/es/api/models/overview#diffusers.ModelMixin.set_attention_backend) para cambiar a un backend de atención más optimizado. Consulte esta [tabla](../optimization/attention_backends#available-backends) 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](https://github.com/huggingface/diffusers/issues/new) si un backend no es compatible. ### Ring Attention Key Las representaciones de clave (K) y valor (V) se comunican entre dispositivos mediante [Ring Attention](https://huggingface.co/papers/2310.01889). 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](/docs/diffusers/main/es/api/parallel#diffusers.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. ```py 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](https://docs.pytorch.org/docs/stable/elastic/run.html), compatible con PyTorch. `--nproc-per-node` se establece en el número de GPUs disponibles. ```shell torchrun --nproc-per-node 2 above_script.py ``` ### Ulysses Attention [Ulysses Attention](https://huggingface.co/papers/2309.14509) 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](/docs/diffusers/main/es/api/parallel#diffusers.ContextParallelConfig) admite Ulysses Attention a través del argumento `ulysses_degree`. Este determina cuántos dispositivos se usarán para Ulysses Attention. Pase [ContextParallelConfig](/docs/diffusers/main/es/api/parallel#diffusers.ContextParallelConfig) a `enable_parallelism()`. ```py # Depending on the number of GPUs available. pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2)) ``` ### Atención unificada [Unified Sequence Parallelism](https://huggingface.co/papers/2405.07719) 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](/docs/diffusers/main/es/api/parallel#diffusers.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). Pasa el [ContextParallelConfig](/docs/diffusers/main/es/api/parallel#diffusers.ContextParallelConfig) con `ulysses_degree` y `ring_degree` establecidos en valores mayores que 1 a `enable_parallelism()`. ```py pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2)) ``` > [!TIP] > 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](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532) en un nodo de 4 GPUs H100. Los resultados se resumen a continuación: | Backend de CP | Tiempo / Iteración (ms) | Pasos / Seg | Memoria máxima (GB) | |--------------------|------------------|-------------|------------------| | ulysses | 6670.789 | 7.50 | 33.85 | | ring | 13076.492 | 3.82 | 56.02 | | unified_balanced | 11068.705 | 4.52 | 33.85 | 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](https://github.com/huggingface/diffusers/pull/12996) 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](/docs/diffusers/main/es/api/parallel#diffusers.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](/docs/diffusers/main/es/api/parallel#diffusers.ContextParallelConfig) con `ulysses_degree` establecido en un valor mayor que 1 y `ulysses_anything=True` a `enable_parallelism()`. ```py pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True)) ``` > [!TIP] 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](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999) en un nodo de 4 GPUs L20. Los resultados se resumen a continuación: | Backend de CP | Tiempo / Iteración (ms) | Pasos / Seg | Memoria máxima (GB) | Forma (HxW)| |--------------------|------------------|-------------|------------------|------------| | ulysses | 281.07 | 3.56 | 37.11 | 1024x1024 | | ring | 351.34 | 2.85 | 37.01 | 1024x1024 | | unified_balanced | 324.37 | 3.08 | 37.16 | 1024x1024 | | ulysses_anything | 280.94 | 3.56 | 37.11 | 1024x1024 | | ulysses | fallido | fallido | fallido | 1008x1008 | | ring | fallido | fallido | fallido | 1008x1008 | | unified_balanced | fallido | fallido | fallido | 1008x1008 | | ulysses_anything | 278.40 | 3.59 | 36.99 | 1008x1008 | 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](https://huggingface.co/papers/2310.01889) por defecto requiere que la longitud de secuencia de los estados ocultos sea divisible uniformemente entre el grado del anillo. [Ring Anything Attention](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4302195582) 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. ```py pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True)) ``` > [!TIP] > Agregue el backend `gloo` a [init_process_group](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) para evitar múltiples sincronizaciones forzadas de CUDA derivadas de las transferencias H2D y D2H. ```py import torch.distributed as dist dist.init_process_group(backend="cpu:gloo,cuda:nccl") ``` > [!NOTE] > 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). | Backend de CP | Tiempo / Iteración (ms) | Pasos / Seg | Memoria máxima (GB) | Forma (HxW)| |--------------------|------------------|-------------|------------------|------------| | ulysses | 259.07 | 3.86 | 33.83 | 1024x1024 | | ring | 338.98 | 2.95 | 33.83 | 1024x1024 | | unified_balanced | 321.54 | 3.11 | 33.83 | 1024x1024 | | ulysses_anything | 259.07 | 3.86 | 33.83 | 1024x1024 | | ring_anything | 340.14 | 2.94 | 33.83 | 1024x1024 | | ulysses | fallido | fallido | fallido | 1008x1008 | | ring | fallido | fallido | fallido | 1008x1008 | | unified_balanced | fallido | fallido | fallido | 1008x1008 | | ulysses_anything | 253.16 | 3.95 | 33.75 | 1008x1008 | | ring_anything | 335.57 | 2.98 | 33.75 | 1008x1008 | A partir de la tabla anterior, Ring Anything Attention ofrece compatibilidad con longitudes de secuencia arbitrarias manteniendo 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](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4304104462). ### parallel_config Pasa `parallel_config` durante la inicialización del modelo para habilitar el paralelismo de contexto. ```py 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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=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](/docs/diffusers/main/es/api/parallel#diffusers.TensorParallelConfig) al argumento `parallel_config` de [from_pretrained()](/docs/diffusers/main/es/api/models/overview#diffusers.ModelMixin.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](https://huggingface.co/black-forest-labs/FLUX.2-dev) (solo el transformer, 32B parámetros, bf16) en 4x A10G (23GB). | tp_degree | método | tiempo de carga | CPU pico/rango | |---|---|---|---| | 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB | | 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB | ```py 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() ``` ```shell 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. > [!CAUTION] > 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 provoca una excepción 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 haya fragmentado 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()](/docs/diffusers/main/es/api/models/overview#diffusers.ModelMixin.save_pretrained) lanza una excepción en uno de ellos. Guarda el modelo antes de dividirlo. ### 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`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) 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 | Estilo | Fragmentos | Cada rango | Uso | |---|---|---|---| | `"colwise"` | características de salida (dimensión 0 de `weight`) | calcula una porción de la salida | `to_q`, `to_k`, `to_v`, proyección de entrada del FFN | | `"rowwise"` | características de entrada (dimensión 1 de `weight`) | calcula una suma parcial | `to_out.0`, proyección de salida del FFN | | `"rowwise_split_input"` | características de entrada (dimensión 1 de `weight`) | calcula un resultado parcial usando su porción de la entrada y los pesos; `AllReduce` suma estos resultados para que cada rank obtenga todas las características de salida | una proyección independiente que lee una activación replicada, por ejemplo, una proyección de modulación AdaLN a partir del embedding del timestep | 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. ```py # 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`. ```py "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`. ```py # in the attention module's __init__ self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim] ``` ```py # in the model's _tp_plan "single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(), ``` Ambas formas aparecen en un solo plan en [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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()](/docs/diffusers/main/es/api/models/overview#diffusers.ModelMixin.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. > [!TIP] > Comience con un plan existente para una arquitectura similar. [QwenImageTransformer2DModel](/docs/diffusers/main/es/api/models/qwenimage_transformer2d#diffusers.QwenImageTransformer2DModel) está completamente no fusionado y cada entrada es un simple `"colwise"`/`"rowwise"`, [FluxTransformer2DModel](/docs/diffusers/main/es/api/models/flux_transformer#diffusers.FluxTransformer2DModel) añade una única proyección row-wise empaquetada, y [Flux2Transformer2DModel](/docs/diffusers/main/es/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel) cubre ambos estilos empaquetados; consulte [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py), [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) y [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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. | Estrategia | Divisiones | Reduce | Latencia para un solo prompt | Mejor cuando | |---|---|---|---|---| | [Accelerate](#accelerate) / [DDP](#pytorch-distributed) | prompts entre réplicas | nada: cada dispositivo mantiene una copia completa | sin cambios | el modelo ya cabe y tienes muchos prompts | | [`device_map`](#devicemap) | componentes entre dispositivos | memoria de pesos | ligeramente peor | el modelo no cabe y el interconector es lento | | [Paralelismo de contexto](#context-parallelism) | la secuencia de entrada | memoria de activación | lower | las secuencias son largas — alta resolución o video | | [Paralelismo de tensores](#tensor-parallelism) | matrices de pesos | memoria de pesos | lower | los pesos de un componente no caben y el interconector es rápido | 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](#ulysses-attention) anteriores. - **Tanto los pesos como la secuencia son demasiado grandes.** Combina la paralelización de tensores y de contexto. [TensorParallelConfig](/docs/diffusers/main/es/api/parallel#diffusers.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.