Paralelismo de tensor para inferência

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.

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.

Defina DistributedConfig(tp_size=N) para usar o plano predefinido do modelo.
Defina um tp_plan manual e passe-o para 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() está obsoleto e será removido na v5.18.

plano automático

plano manual

Estratégias de particionamento

ALL_PARALLEL_STYLES, uma ParallelInterface definida em 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 registra as estratégias de paralelismo de tensor nela:

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, 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 reconstruí-la como 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 particionamento de um peso linear na dimensão de recursos de entrada (`Shard(-1)`) e replica o viés, 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 token ids 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 cuja computação 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 particionados como tensores locais. Use-o no proprietário de um `nn.Parameter` simples, como attention sinks, que a própria entrada do plano particiona.
`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.

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.

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 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`	Rejeitar um parâmetro que essa estratégia não consegue dividir em shards, 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 o 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.

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

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

Implemente shard_param para envolver um parâmetro como um placeholder 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.

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

Implemente as transformações de entrada e saída. install_forward as chama em torno do forward original do módulo, de modo que 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.

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.

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.

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 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 e DTensor do torch.distributed 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.

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 DTensors, 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.

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.

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 a redistribuição para Replicate() é o que emite o all-reduce.

Recursos

A seção Ultra-Scale Playbook sobre paralelismo de tensor fornece mais detalhes.
Consulte o guia de expert parallelism se você estiver usando um modelo de mixture-of-experts (MoE). Esses modelos suportam tensor parallelism e expert parallelism.
Leia o post no blog Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand 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 para aprender a usá-lo em um ambiente de treinamento.

Update on GitHub