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 y calcula un resultado parcial. Un AllReduce o AllGather en el límite de la capa combina los resultados parciales para obtener la salida completa. A diferencia de context parallelism, reduce la memoria de weight por dispositivo, lo cual es útil cuando un modelo no cabe en un solo dispositivo. Consulte Distributed inference para compararlo con las demás estrategias.

Cargar un modelo fragmentado

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. La clase del modelo debe definir un _tp_plan, que mapea los nombres de los módulos a un estilo de partición. Para comprobar si un modelo admite la paralelización de tensores, busque este atributo en su clase, como Flux2Transformer2DModel.

Para añadir un plan a otro modelo, consulte Añadir paralelismo de tensores a un modelo.

Al cargar de esta manera, el checkpoint se fragmenta a medida que se lee. Cada rank lee solo su propia porción de cada peso y la coloca directamente en su GPU, por lo que ningún rank mantiene un peso de tamaño completo. La memoria por rank disminuye a medida que aumenta tp_degree.

import os
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()
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

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

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, 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()

Guarde el script como tensor_parallel_flux2.py y ejecútelo con torchrun.

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

El script establece tp_degree en world_size, por lo que --nproc-per-node 4 divide el transformer en 4 GPUs.

Solo el transformer se fragmenta. Cada rank mantiene una copia completa del codificador de texto, por lo que cada GPU necesita suficiente memoria para el codificador de texto más su parte del transformer.

Rendimiento de carga

Cargar con parallel_config es más rápido y utiliza menos memoria de CPU por rank que cargar el modelo completo y luego llamar a enable_parallelism(). La tabla siguiente muestra los resultados para black-forest-labs/FLUX.2-dev (solo el transformer, 32B parámetros, bf16) en 4x A10G (23GB).

Limitaciones

La paralelización de tensores no funciona con algunas opciones y funciones de carga, y no se puede guardar un modelo fragmentado.

• Cargar con una parallel_config de paralelismo de tensores no funciona 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 cargar silenciosamente el punto de control completo.
• Combinar el paralelismo de tensores con cuantización, offloading o adaptadores LoRA también genera un error, sin importar cómo se fragmente el modelo.
• Para fragmentar un modelo que ya está en memoria, llama a enable_parallelism() con la misma configuración. Primero carga el modelo completo y luego lo fragmenta, por lo que cada rank necesita memoria para el checkpoint completo. Cargar con parallel_config evita ese coste.
• Guardar un modelo de paralelismo de tensores genera un error en save_pretrained(). Guarda el modelo antes de dividirlo.

Añadir paralelismo de tensores a un modelo

Esta sección está dirigida a los colaboradores que deseen agregar soporte para paralelismo de tensores a un modelo de Diffusers.

La paralelización de tensores solo funciona en modelos que definen un _tp_plan, un atributo de clase que asigna 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 la parte anterior al * debe resolverse a un nn.ModuleList, de modo que una entrada cubra todos los bloques. Una clave sin * se aplica al propio modelo. Las rutas son relativas al modelo.

Colwise y rowwise

Acompáñe siempre "colwise" y "rowwise" 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 forma obliga a hacer un gather en el medio y comunica mucho más.

En la atención, las proyecciones columnales de Q, K y V asignan a cada rango un subconjunto de cabezas, por lo que tp_degree debe dividir el número de cabezas. Los duplicados del flujo del codificador (add_q_proj, to_add_out, ff_context) siguen el mismo patrón que sus equivalentes en el flujo 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

Pase blocks, una lista de enteros proporcionales cuya suma divide la dimensión empaquetada. Por ejemplo, use [1, 1] para una proyección de puerta y una de subida de SwiGLU de mitades iguales, o [1, 1, 1, 3, 3] para una proyección fusionada de Q, K, V, puerta y subida con mlp_ratio=3.

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

Cuando los tamaños de los bloques solo se conocen desde la configuración, omita blocks 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(),

El plan en transformer_flux2.py utiliza ambas formas. El par SwiGLU de doble flujo pasa blocks directamente. La proyección QKV y MLP de flujo único fusionada almacena sus tamaños en el Linear dentro de Flux2ParallelSelfAttention.__init__, ya que dependen de mlp_ratio y mlp_mult_factor.

Qué dejar fuera

Los módulos que se dejan fuera del plan se replican en cada rank. Deje fuera las capas de normalización, las incrustaciones de parche y de texto, y las capas finales norm_out y proj_out. Son pequeñas, por lo que fragmentarlas ahorra poca memoria y añade comunicación. La modulación AdaLN (img_mod y txt_mod) suele ser lo suficientemente pequeña para dejarla fuera también. Cuando es grande, como adaln_proj.linear de MiniMax-H3, frágméntela con "rowwise_split_input".

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.

from_pretrained() y enable_parallelism() verifican ambas cosas antes de cargar o dividir 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.

• transformer_qwenimage.py no tiene proyecciones fusionadas, por lo que cada entrada es simplemente "colwise" o "rowwise".
• transformer_flux.py añade una proyección empaquetada por filas.
• transformer_flux2.py utiliza ambos estilos empaquetados.