Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /main /pt /optimization /tensor_parallelism.md

mishig

3 days ago

preview code

download

raw

11.2 kB

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

Xet Storage Details

Size:: 11.2 kB
Xet hash:: e8f36e375b0e02689e16e594dc87b5c4b1f6a0c04534a0fb865c72b2ddc6150f

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

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