# 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 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](./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](../training/distributed_inference) para compararlo con las demás estrategias.

## Cargar un modelo fragmentado

Pase un [TensorParallelConfig](/docs/diffusers/v0.43.0/es/api/parallel#diffusers.TensorParallelConfig) al argumento `parallel_config` de [from_pretrained()](/docs/diffusers/v0.43.0/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. 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](/docs/diffusers/v0.43.0/es/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

Para añadir un plan a otro modelo, consulte [Añadir paralelismo de tensores a un modelo](#add-tensor-parallelism-to-a-model).

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`.

```py
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`.

```shell
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](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 |

## 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()](/docs/diffusers/v0.43.0/es/api/models/overview#diffusers.ModelMixin.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`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), 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

| 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 una suma parcial a partir de su porción de la entrada replicada | una proyección independiente cuya entrada está replicada, como una proyección de modulación AdaLN que lee el embedding del timestep |

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.

```py
# 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`.

```py
"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`.

```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(),
```

El plan en [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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()](/docs/diffusers/v0.43.0/es/api/models/overview#diffusers.ModelMixin.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.

> [!TIP]
> Comience con un plan existente para una arquitectura similar.
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) no tiene proyecciones fusionadas, por lo que cada entrada es simplemente `"colwise"` o `"rowwise"`. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) añade una proyección row-wise empaquetada. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) utiliza ambos estilos empaquetados.