# Paralelismo de tensor

[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) distribui as matrizes de pesos de um modelo entre os dispositivos. Cada dispositivo mantém uma fatia por coluna (`"colwise"`) ou por linha (`"rowwise"`) de cada camada e calcula um resultado parcial. Um `AllReduce` ou `AllGather` na fronteira da camada combina os resultados parciais para formar a saída completa. Diferentemente do [context parallelism](./context_parallelism), ele reduz a memória de *weight* por dispositivo, o que ajuda quando um modelo não cabe em um único dispositivo. Consulte [Distributed inference](../training/distributed_inference) para compará-lo com as outras estratégias.

## Carregar um modelo fragmentado

Passe um [TensorParallelConfig](/docs/diffusers/main/pt/api/parallel#diffusers.TensorParallelConfig) para o argumento `parallel_config` do [from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained) do modelo. `tp_degree` é o número de dispositivos para os quais o modelo será dividido e deve dividir o número de cabeças de atenção do modelo. A classe do modelo deve definir um `_tp_plan`, que mapeia nomes de módulos para um estilo de divisão. Para verificar se um modelo suporta paralelismo de tensor, procure por esse atributo em sua classe, como [Flux2Transformer2DModel](/docs/diffusers/main/pt/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

Para adicionar um plano a outro modelo, consulte [Adicionar paralelismo de tensor a um modelo](#add-tensor-parallelism-to-a-model).

Carregar dessa forma faz o checkpoint ser fragmentado conforme é lido. Cada rank lê apenas sua própria fatia de cada peso e a coloca diretamente em sua GPU, de modo que nenhum rank mantém um peso de tamanho completo. A memória por rank diminui à medida que `tp_degree` aumenta.

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

Salve o script como `tensor_parallel_flux2.py` e execute-o com `torchrun`.

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

O script define `tp_degree` como `world_size`, portanto `--nproc-per-node 4` distribui o transformer em 4 GPUs.

Apenas o transformer é fragmentado. Cada rank mantém uma cópia completa do codificador de texto, portanto, cada GPU precisa de memória suficiente para o codificador de texto mais sua parte do transformer.

## Desempenho de carregamento

Carregar com `parallel_config` é mais rápido e usa menos memória de CPU por rank do que carregar o modelo completo e, em seguida, chamar `enable_parallelism()`. A tabela abaixo mostra os resultados para [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (somente o transformer, 32B de parâmetros, bf16) em 4x A10G (23GB).

| tp_degree | método | tempo de carregamento | CPU de pico/por rank |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB |

## Limitações

O paralelismo de tensor não funciona com algumas opções e recursos de carregamento, e não é possível salvar um modelo fragmentado.

- Carregar com um `parallel_config` de paralelismo de tensor não funciona com `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True` ou pesos que não sejam safetensors. Cada um deles gera um erro em vez de carregar silenciosamente o checkpoint completo.
- A combinação de paralelismo de tensor com quantização, offloading ou adaptadores LoRA também gera um erro, independentemente de como o modelo é particionado.
- Para dividir um modelo que já está na memória, chame `enable_parallelism()` com a mesma configuração. Ele carrega o modelo completo primeiro e depois o divide, então cada rank precisa de memória para o checkpoint completo. Carregar com `parallel_config` evita esse custo.
- Salvar um modelo paralelo de tensor gera um erro em [save_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.save_pretrained). Salve o modelo antes de fragmentá-lo.

## Adicionar paralelismo de tensor a um modelo

Esta seção é destinada a contribuidores que desejam adicionar suporte a paralelismo de tensor a um modelo do Diffusers.

O paralelismo de tensores funciona apenas em modelos que definem um `_tp_plan`, um atributo de classe que mapeia nomes de módulos para um estilo de particionamento. Escrever um geralmente significa associar cada projeção que *expande* a dimensão oculta com a projeção que a *contrai* de volta.

Cada chave pode conter no máximo um `*`, e a parte antes do `*` deve resolver para um [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), de modo que uma entrada cubra todos os blocos. Uma chave sem `*` se aplica ao próprio modelo. Os caminhos são relativos ao modelo.

### Colwise e rowwise

| Estilo | Shards | Cada rank | Usar para |
|---|---|---|---|
| `"colwise"` | características de saída (dim 0 de `weight`) | calcula uma fatia da saída | `to_q`, `to_k`, `to_v`, in-projection do FFN |
| `"rowwise"` | recursos de entrada (dim 1 de `weight`) | calcula uma soma parcial | `to_out.0`, projeção de saída do FFN |
| `"rowwise_split_input"` | recursos de entrada (dim 1 de `weight`) | calcula uma soma parcial a partir de sua fatia da entrada replicada | uma projeção independente cuja entrada é replicada, como uma projeção de modulação AdaLN que lê o embedding do timestep |

Sempre combine `"colwise"` e `"rowwise"` nessa ordem. Uma projeção `"colwise"` deixa sua saída fragmentada, a projeção `"rowwise"` subsequente consome esse fragmento diretamente e um único `AllReduce` na fronteira do bloco reconstrói o resultado. Fragmentar o par de qualquer outra forma força um gather no meio e comunica muito mais.

Na atenção, as projeções de Q, K e V por coluna atribuem a cada rank um subconjunto de cabeças, portanto `tp_degree` deve dividir o número de cabeças. Os duplicados do stream do codificador (`add_q_proj`, `to_add_out`, `ff_context`) seguem o mesmo padrão de seus equivalentes no stream de imagem.

### Projeções fundidas

Quando um `Linear` empacota vários tensores lógicos ao longo da dimensão que está sendo fragmentada, fatias simples de `"colwise"`/`"rowwise"` atravessam a concatenação e desalinharam as peças. Use `PackedColwiseParallel`/`PackedRowwiseParallel` em vez disso, que fragmenta cada bloco empacotado de forma independente.

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

Pass `blocks`, uma lista de inteiros proporcionais cuja soma divide a dimensão empacotada. Por exemplo, use `[1, 1]` para uma projeção de gate e up do SwiGLU com metades iguais, ou `[1, 1, 1, 3, 3]` para uma projeção fundida de Q, K, V, gate e up com `mlp_ratio=3`.

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

Quando os tamanhos dos blocos são conhecidos apenas a partir da configuração, omita `blocks` e armazene os tamanhos absolutos no `Linear` durante o `__init__`, como `_tp_packed_col_blocks` ou `_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(),
```

O plano em [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) usa ambas as formas. O par SwiGLU de duplo fluxo passa `blocks` diretamente. A projeção QKV e MLP de fluxo único fundida armazena seus tamanhos no `Linear` em `Flux2ParallelSelfAttention.__init__`, porque eles dependem de `mlp_ratio` e `mlp_mult_factor`.

### O que deixar de fora

Módulos excluídos do plano são replicados em cada rank. Exclua as camadas de normalização, as incorporações de patch e de texto, e o `norm_out` e `proj_out` finais. Eles são pequenos, portanto, fragmentá-los poupa pouco memória e adiciona comunicação. A modulação AdaLN (`img_mod` e `txt_mod`) geralmente é pequena o suficiente para também ser excluída. Quando é grande, como o `adaln_proj.linear` do MiniMax-H3, fragmente-a com `"rowwise_split_input"`.

### Restrições e verificação

- `tp_degree` deve dividir `config.num_attention_heads`.
- Cada bloco empacotado deve ser *individualmente* divisível por `tp_degree`, e não apenas sua soma.

[from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained) e `enable_parallelism()` verificam ambos antes de carregar ou dividir qualquer peso.

Valide um novo plano numericamente, em vez de a olho nu: gere com uma semente fixa em um único dispositivo e, em seguida, novamente sob paralelismo de tensor, e compare as saídas. Um `"colwise"`/`"rowwise"` mal posicionado geralmente ainda executa e produz uma imagem plausível, porém incorreta.

> [!TIP]
> Comece a partir de um plano existente para uma arquitetura semelhante.
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) não possui projeções fundidas, portanto cada entrada é apenas `"colwise"` ou `"rowwise"`. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) adiciona uma projeção row-wise empacotada. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) usa ambos os estilos empacotados.