Paralelismo de tensor

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, 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 para compará-lo com as outras estratégias.

Carregar um modelo fragmentado

Passe um TensorParallelConfig para o argumento parallel_config do 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.

Para adicionar um plano a outro modelo, consulte Adicionar paralelismo de tensor a um modelo.

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.

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.

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 (somente o transformer, 32B de parâmetros, bf16) em 4x A10G (23GB).

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

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.

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

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

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

O plano em 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() 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.

Comece a partir de um plano existente para uma arquitetura semelhante.

• transformer_qwenimage.py não possui projeções fundidas, portanto, cada entrada é apenas "colwise" ou "rowwise".
• transformer_flux.py adiciona uma projeção empacotada linha a linha.
• transformer_flux2.py usa ambos os estilos empacotados.