# Paralelismo de tensor para inferência

[Tensor parallelism](./perf_train_gpu_many#tensor-parallelism) divide uma camada do modelo em fragmentos para que vários aceleradores de hardware trabalhem nela simultaneamente. Isso permite executar modelos que excedem a capacidade de memória de uma única GPU e alcançar maior throughput. Você precisará de comunicação intra-nó rápida, pois as GPUs trocam resultados parciais em cada camada.

Um modelo suporta paralelismo de tensor se o seu config define `base_model_tp_plan`. Verifique um modelo carregado com a propriedade `supports_tp_plan`.

```py
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)
```

Este guia aborda a ativação do paralelismo de tensor no Transformers e as estratégias de particionamento disponíveis.

## Particionamento de um modelo

Configure o número de dispositivos de paralelismo de tensor com `tp_size` em [DistributedConfig](/docs/transformers/main/pt/expert_parallelism#transformers.DistributedConfig).

- Defina `DistributedConfig(tp_size=N)` para usar o plano predefinido do modelo.
- Defina um `tp_plan` manual e passe-o para [DistributedConfig](/docs/transformers/main/pt/expert_parallelism#transformers.DistributedConfig) junto com `tp_size`.

Você também pode definir `tp_plan="auto"` para solicitar explicitamente o plano predefinido. Quando `tp_size` é omitido e um `tp_plan` está definido, `tp_size` é derivado de `WORLD_SIZE` dividido pelos outros tamanhos de paralelismo. Passar `tp_plan` diretamente para [from_pretrained()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained) está obsoleto e será removido na v5.18.

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)
print(model.tp_plan)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

# distributed run
outputs = model(inputs)
```

Inicie o script de inferência com [torchrun](https://pytorch.org/docs/stable/elastic/run.html). Use um processo por GPU.

```bash
torchrun --nproc-per-node 4 demo.py
```

Defina um plano de paralelismo de tensor para cada camada em `tp_plan` e passe-o por meio de [DistributedConfig](/docs/transformers/main/pt/expert_parallelism#transformers.DistributedConfig). O exemplo abaixo usa particionamento por coluna e por linha. Consulte a seção [Estratégias de particionamento](#partitioning-strategies) para outras estratégias suportadas.

A particionamento manual exige um profundo entendimento da arquitetura do modelo e das interações entre estratégias. Escolhas inadequadas de particionamento resultam em modelos lentos que falham ou produzem resultados incorretos. O [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) explica as estratégias de particionamento em detalhes.

As chaves são nomes de módulos ou parâmetros, em que `*` representa os índices das camadas. Um nome de estratégia não reconhecido gera um `ValueError` listando os nomes suportados, e o Transformers registra um aviso para regras de plano que não corresponderam a nada e para parâmetros que nenhuma regra cobriu.

```py
from transformers import AutoModelForCausalLM, DistributedConfig

tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
    ...
}

distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype="auto",
    distributed_config=distributed_config,
)
print(model.tp_plan)
```

## Estratégias de particionamento

`ALL_PARALLEL_STYLES`, uma `ParallelInterface` definida em [distributed/parallel_utils.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/parallel_utils.py), mapeia cada nome de estratégia que você pode usar em um plano para uma instância de estratégia configurada. Você não interage diretamente com ela para fazer o sharding de um modelo, mas é a lista autoritativa de nomes disponíveis. [distributed/tensor_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py) registra as estratégias de paralelismo de tensor nela:

```py
ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())
```

Toda estratégia de paralelismo de tensor é uma subclasse de `TensorParallelLayer` naquele arquivo. As estratégias de paralelismo de especialistas, como `ep_dispatch_experts`, estão em [distributed/expert_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/expert_parallel.py), são subclasses de `ExpertParallelLayer`, que possui os mesmos hooks, e se registram no mesmo `ALL_PARALLEL_STYLES`. As duas que você mais usará, `ColwiseParallel` e `RowwiseParallel`, aceitam posicionamentos de `input_layouts` e `output_layouts`, o que permite que uma única classe cubra vários nomes de plano. `colwise` deixa sua saída fragmentada na última dimensão para que uma camada `rowwise` subsequente a consuma, enquanto `colwise_gather_output` faz um all-gather para reconstituí-la em um tensor completo.

A tabela abaixo descreve cada estratégia.

| Estratégia | Descrição |
|---|---|
| `ColwiseParallel` | Faz o particionamento de um peso linear na dimensão de recursos de saída (`Shard(0)` para um peso 2D de `nn.Linear`) e particiona um viés 1D. Redistribui a entrada para `Replicate()` e, em seguida, posiciona a saída em `output_layouts`. |
| `RowwiseParallel` | Faz o sharding de um peso linear na dimensão de recursos de entrada (`Shard(-1)`) e replica o bias, que é adicionado após a redução. O forward de cada rank produz uma saída `Partial()` contendo sua parte da soma, e `transform_output_post_forward` redistribui isso para `output_layouts`. Reduzir para `Replicate()` emite um all-reduce, e reduzir para `Shard(1)` emite um reduce-scatter. |
| `EmbeddingParallel` | Faz o sharding de um `nn.Embedding` com os ids de tokens replicados. `shard_dim=0` (`embedding_rowwise`, a escolha usual) faz o sharding do vocabulário e faz o all-reduce da busca parcial de cada rank para a saída completa. `shard_dim=1` (`embedding_colwise`) faz o sharding da dimensão oculta e faz o all-gather do slice de saída de cada rank. |
| `PackedColwiseParallel` | Uma variante de `ColwiseParallel` para pesos fundidos, por exemplo `up_proj` e `gate_proj` empacotados em `gate_up_proj`. Use `split_factor` quando um peso empacota um número de camadas diferente de dois. |
| `PackedRowwiseParallel` | O equivalente por linhas, para pesos empacotados ao longo da dimensão final. Replica parâmetros 1D. |
| `SequenceParallel` | Replica os parâmetros do módulo e faz o sharding da sua entrada em `sequence_dim` (padrão é `1`). Usado para normalizações que operam por token, como `LayerNorm` e `RMSNorm`. |
| `ReplicatedWithGradAllReduce` | Replica um parâmetro, mas faz o all-reduce do seu gradiente. Necessário para normas que ficam entre uma camada column-wise e uma camada row-wise e normalizam ao longo de um eixo sharded, onde cada rank só vê seus próprios heads. |
| `AllReduceParallel` | Realiza um all-reduce da saída do forward do `Partial()` de um módulo para `Replicate()`. Use-o como um ponto de sincronização para um módulo cujo cálculo termina em uma soma parcial. |
| `MlaKvAProjParallel` | Divide a saída de `kv_a_proj_with_mqa` da atenção MLA no estilo DeepSeek-V2 e faz o all-reduce do gradiente da metade do RoPE, o que contorna `kv_b_proj` e, caso contrário, manteria um gradiente parcial. Exige `qk_rope_head_dim` na configuração do modelo. |
| `MoeExpertsParallel` | Especialistas MoE com paralelismo de tensor. Executa o all-reduce da saída dos especialistas na propagação direta e adiciona os all-reduces na propagação reversa para os estados ocultos e os pesos de roteamento. |
| `MoeIdentityParallel` | Divide previamente a entrada de um especialista zero ou identidade pelo tamanho da malha, cancelando o all-reduce que `moe_tp_experts` aplica a jusante. |
| `LocalParamsParallel` | Executa o forward de um módulo com seus próprios parâmetros fragmentados como tensores locais. Use-o no proprietário de um `nn.Parameter` puro, como attention sinks, que a própria entrada do plano fragmenta. |
| `MoEParamShard` | Fragmenta os pesos dos especialistas de MoE em um placement dado. Suporta o nome `grouped_gemm`, onde `shards_expert_dim=True` também reescreve `module.num_experts` para a contagem de especialistas por rank. |
| `EpRouterParallel` | Mascara as pontuações do roteador para especialistas não locais e remapeia os IDs globais de especialistas para os locais, de modo que cada rank execute apenas os especialistas que possui. Exige que `num_experts` seja divisível pelo tamanho da malha. |
| `RouterParallelMegaMoe` | Variante de roteador para DeepGEMM Mega MoE, que despacha os especialistas dentro do kernel, que deseja a saída do roteador inalterada. |
| `MoeTensorParalellMegaMoeExperts` | Camada de especialistas somente para inferência do DeepGEMM Mega MoE. Pula as sincronizações de gradiente e passa o grupo de processos para o módulo, permitindo que o kernel configure seus buffers compartilhados no primeiro forward. |

### Estratégias empacotadas

O empacotamento de pesos combina várias camadas lineares em uma única camada maior. As estratégias `PackedColwiseParallel` e `PackedRowwiseParallel` dividem os pesos empacotados corretamente. As estratégias básicas `ColwiseParallel` ou `RowwiseParallel` dividem os pesos empacotados incorretamente.

O exemplo abaixo agrupa `up_proj` e `gate_proj` em um único módulo `gate_up_proj` e exige a estratégia `packed_rowwise` para dividir `gate_up_proj`.

```python
class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))
```

Use a multiplicação de matrizes em lote na passagem `forward` para calcular a saída do módulo `gate_up_proj`.

```python
def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up
```

Um `Shard` simples divide essa dimensão em blocos contíguos, de modo que um rank receberia o final de `gate` e o início de `up`, em vez de um recorte de cada. As estratégias empacotadas usam `_StridedShard` com `split_factor` (`2` por padrão) para intercalar a divisão, garantindo que cada rank receba um recorte correspondente de ambas as metades, para que o `chunk` continue alinhado após o sharding.

## Estratégias de particionamento personalizadas

Herdar de `TensorParallelLayer` em [distributed/tensor_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py) para criar uma estratégia de particionamento personalizada. Sobrescreva apenas os hooks que sua estratégia precisa, pois todos possuem um padrão sem operação (no-op).

| Hook | Propósito |
|---|---|
| `validate_param` | Rejeite um parâmetro que esta estratégia não consegue dividir em fragmentos, antes de qualquer carregamento de pesos. |
| `shard_param` | Substitua um parâmetro por um placeholder `DTensor` para que o loader saiba qual shard pertence a este rank. |
| `transform_inputs_pre_forward` | Redistribua as entradas do módulo para o layout que seu forward espera. |
| `context_around_forward` | Envolva o forward em um gerenciador de contexto, por exemplo, para expor tensores locais a um kernel. |
| `transform_output_post_forward` | Redistribua ou reduza a saída do módulo. |
| `should_use_local_tensors` | Relate que o forward deste módulo precisa de tensores simples em vez de `DTensor`s. |
| `install_forward` | Substitua `module.forward` por completo. Sobrescreva este método apenas quando os hooks acima não forem suficientes, como `ReplicatedWithGradAllReduce` faz para registrar um hook de backward. |

O exemplo abaixo percorre uma versão reduzida de `ColwiseParallel`.

1. Herde de `TensorParallelLayer` e armazene os placements com os quais a estratégia trabalha. A classe base não define `__init__`, portanto, não há nada para chamar `super()`.

    ```python
    class ColwiseParallel(TensorParallelLayer):
        def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
            self.input_layouts = input_layouts or Replicate()
            self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
            self.use_local_output = use_local_output
    ```

2. Implemente `shard_param` para encapsular um parâmetro como um placeholder de `DTensor`. Ele é executado em meta tensors, portanto, `distribute_tensor` apenas constrói metadados e não move dados. Passe `src_data_rank=None` porque ainda não há um tensor completo para dispersar.

    ```python
    def shard_param(self, module, param, mesh):
        meta = module._parameters.get(param)
        if meta is None:
            return
        # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
        placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
        module._parameters[param] = torch.nn.Parameter(
            distribute_tensor(meta, mesh, [placement], src_data_rank=None),
            requires_grad=meta.requires_grad,
        )
    ```

3. Implemente as transformações de entrada e saída. `install_forward` as chama em torno do forward original do módulo, portanto, elas apenas precisam mover os tensores entre layouts. A particionamento em colunas espera uma entrada replicada e produz uma saída fragmentada na última dimensão.

    ```python
    def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
        x = args[0]
        if not isinstance(x, DTensor):
            x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
        if x.placements != (Replicate(),):
            x = x.redistribute(placements=[Replicate()])
        return (x,) + args[1:], kwargs

    def transform_output_post_forward(self, module, output, mesh):
        if not isinstance(output, DTensor):
            output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
        if output.placements != (self.output_layouts,):
            output = output.redistribute(placements=[self.output_layouts])
        return output.to_local() if self.use_local_output else output
    ```

    O `ColwiseParallel` incluído adiciona caminhos rápidos sobre isso que pulam o ciclo de ida e volta do `DTensor` para inferência de `nn.Linear` simples e para módulos quantizados que precisam de tensores locais. Leia o código-fonte antes de copiá-lo se a sua estratégia precisar do mesmo tratamento.

4. Registre a estratégia para que um `tp_plan` possa nomeá-la. O registro recebe uma instância da estratégia, portanto, as colocações que você passa para `__init__` são armazenadas junto com o nome. Observe o `()` abaixo.

    ```python
    import torch

    from transformers import AutoModelForCausalLM, DistributedConfig
    from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES

    ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
    tp_plan = {
        "model.layers.*.self_attn.q_proj": "colwise_custom",
        ...
    }
    distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Meta-Llama-3-8B-Instruct",
        dtype=torch.bfloat16,
        distributed_config=distributed_config,
    )
    ```

## Benchmarks

O paralelismo de tensor acelera significativamente a inferência, especialmente para tamanhos de lote grandes ou sequências longas.

Este gráfico mostra a aceleração esperada para uma única passagem direta no [Llama](./model_doc/llama) com um comprimento de sequência de 512.

    

## Implementação do design

O Transformers implementa a paralelização de tensores de uma forma agnóstica ao framework. Ele depende de [DeviceMesh](https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html) e [DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html) do [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) para fornecer uma interface simples e extensível.

### DeviceMesh

`DeviceMesh` cria uma grade multidimensional de dispositivos que se comunicam entre si. Diferentes estratégias de paralelização exigem diferentes padrões de comunicação. Crie um `DeviceMesh` com vários sub-meshes para lidar com esses padrões.

```python
import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])
```

A maioria das estratégias de paralelização do `torch.distributed` se aplica à própria malha ou a suas submalhas. A malha gerencia automaticamente os padrões de comunicação.

### DTensor

`DTensor` (Distributed Tensor) gerencia a lógica de distribuição sobre as operações habituais de tensor. Os pesos do modelo sob paralelismo de tensor são armazenados como `DTensor`s, o que permite que uma estratégia descreva a comunicação como uma mudança de layout, em vez de uma operação coletiva explícita.

O atributo `placements` informa ao PyTorch como um tensor é distribuído entre os dispositivos em um `DeviceMesh`. Ele aceita os seguintes valores:

- `Shard(dimension)` divide um `DTensor` em uma dimensão especificada ao longo do `DeviceMesh` em que foi construído. A particionamento por colunas divide a dimensão de recursos de saída do peso e a única dimensão do viés.

    ```python
    weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
    bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension
    ```

    A particionamento por linhas divide a dimensão dos recursos de entrada e replica o viés, pois o viés é adicionado uma única vez após a redução, em vez de em cada rank.

    ```python
    weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
    bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs
    ```

- `Replicate()` replica um `DTensor` em todo o `DeviceMesh`, criando uma cópia completa do tensor em cada dispositivo.

- `Partial()` marca um tensor como pendente de uma redução. A saída do forward de uma camada row-wise é `Partial()`, e redistribuí-la para `Replicate()` é o que emite o all-reduce.

## Recursos

- A seção [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) sobre paralelismo de tensor fornece mais detalhes.

- Consulte o guia de [parallelismo de especialistas](./expert_parallelism) se você estiver usando um modelo de mixture-of-experts (MoE). Esses modelos suportam parallelismo de tensor e parallelismo de especialistas.

- Leia o post no blog [Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand](https://huggingface.co/blog/qgallouedec/tp) para uma visão geral rápida do paralelismo de tensor e saiba como as configurações de paralelismo por coluna e por linha diferem.

- Consulte o guia de treinamento de [Tensor parallelism](./tensor_parallelism) para aprender a usá-lo em um ambiente de treinamento.