Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /ru /optimization /tensor_parallelism.md

mishig

about 2 hours ago

preview code

download

raw

15.6 kB

# Параллелизм тензоров

[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) распределяет матрицы весов модели по устройствам. Каждое устройство хранит столбцовый (`"colwise"`) или строчный (`"rowwise"`) срез каждого слоя и вычисляет частичный результат. Операция `AllReduce` или `AllGather` на границе слоя объединяет частичные результаты в полный вывод. В отличие от [context parallelism](./context_parallelism), он снижает потребление памяти под *weight* на каждое устройство, что полезно, когда модель не помещается на одном устройстве. Для сравнения с другими стратегиями см. раздел [Distributed inference](../training/distributed_inference).

## Загрузка модели с разбивкой на части

Pass a [TensorParallelConfig](/docs/diffusers/v0.43.0/ru/api/parallel#diffusers.TensorParallelConfig) to the `parallel_config` argument of the model's [from_pretrained()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.from_pretrained). `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model's class must define a `_tp_plan`, which maps module names to a sharding style. To check whether a model supports tensor parallelism, look for this attribute on its class, such as [Flux2Transformer2DModel](/docs/diffusers/v0.43.0/ru/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

Чтобы добавить план к другой модели, см. [Добавление тензорного параллелизма в модель](#add-tensor-parallelism-to-a-model).

Загрузка таким образом разбивает контрольную точку на фрагменты по мере её чтения. Каждый ранг читает только свой срез каждого веса и размещает его непосредственно на своём GPU, поэтому ни один ранг никогда не хранит вес полного размера. Память на ранг уменьшается по мере увеличения `tp_degree`.

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

Сохраните скрипт как `tensor_parallel_flux2.py` и запустите его с помощью `torchrun`.

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

Скрипт устанавливает `tp_degree` в `world_size`, поэтому `--nproc-per-node 4` распределяет трансформер по 4 GPU.

Шардируется только трансформер. Каждый ранг хранит полную копию текстового энкодера, поэтому на каждой GPU должно быть достаточно памяти для текстового энкодера и его доли трансформера.

## Производительность загрузки

Загрузка с `parallel_config` быстрее и потребляет меньше CPU-памяти на ранг, чем загрузка полной модели с последующим вызовом `enable_parallelism()`. В таблице ниже приведены результаты для [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (только трансформер, 32 млрд параметров, bf16) на 4x A10G (23 ГБ).

| tp_degree | метод | время загрузки | пиковое значение CPU/ранк |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6,8 ГБ |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64,1 ГБ |
## Ограничения
Tensor parallelism не работает с некоторыми параметрами загрузки и функциями, а сохранить модель с разбивкой на шунки невозможно.
- Загрузка с тензорно-параллельным `parallel_config` не работает с `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True` или весами, не хранящимися в формате safetensors. Каждый из этих параметров приводит к ошибке, а не к тихой загрузке полной контрольной точки.
- Сочетание тензорного параллелизма с квантованием, offloading или адаптерами LoRA также приводит к ошибке, независимо от способа шардирования модели.
- Чтобы разбить на шарды модель, уже загруженную в память, вызовите `enable_parallelism()` с той же конфигурацией. Сначала загружается полная модель, а затем она разбивается на шарды, поэтому каждому ранку требуется память для полного чекпоинта. Загрузка с `parallel_config` позволяет избежать этих затрат.
- Сохранение тензорно-параллельной модели вызывает ошибку в [save_pretrained()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.save_pretrained). Сохраните модель до её шардирования.

## Добавление тензорного параллелизма в модель

Этот раздел предназначен для участников, которые хотят добавить поддержку тензорного параллелизма в модель Diffusers.

Tensor parallelism работает только с моделями, которые определяют `_tp_plan` — атрибут класса, сопоставляющий имена модулей со стилем шардирования. В основном его создание сводится к тому, чтобы связать каждую проекцию, которая *расширяет* скрытое измерение, с проекцией, которая *сжимает* его обратно.

Каждый ключ может содержать не более одного `*`, а часть до `*` должна указывать на [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), чтобы одна запись покрывала все блоки. Ключ без `*` применяется к самой модели. Пути указаны относительно модели.

### По столбцам и по строкам

| Стиль | Шарды | Каждый ранг | Использование |
|---|---|---|---|
| `"colwise"` | выходные признаки (`weight` размер 0) | вычисляет срез выходных данных | `to_q`, `to_k`, `to_v`, входная проекция FFN |
| `"rowwise"` | входные признаки (размерность `weight` 1) | вычисляет частичную сумму | `to_out.0`, выходная проекция FFN |
| `"rowwise_split_input"` | входные признаки (размерность `weight` 1) | вычисляет частичную сумму из своего среза реплицированного входа | независимая проекция, вход которой реплицирован, например, проекция модуляции AdaLN, читающая эмбеддинг временного шага |

Всегда используйте `"colwise"` и `"rowwise"` именно в таком порядке. Проекция `"colwise"` оставляет её вывод в разбитом виде, следующая проекция `"rowwise"` напрямую потребляет этот фрагмент, а один вызов `AllReduce` на границе блока восстанавливает результат. Любое другое разбиение пары приводит к необходимости собрать данные посередине и значительно увеличивает объём коммуникации.

В механизме внимания столбчатые проекции Q, K и V распределяют подмножество голов по каждому рангу, поэтому `tp_degree` должно делить количество голов. Дубликаты энкодерного потока (`add_q_proj`, `to_add_out`, `ff_context`) следуют тому же паттерну, что и их аналоги в потоке изображений.

### Сплоченные проекции

Когда один `Linear` упаковывает несколько логических тензоров вдоль измерения, по которому выполняется шардирование, обычные срезы `"colwise"`/`"rowwise"` проходят прямо через конкатенацию и смещают фрагменты. Вместо этого используйте `PackedColwiseParallel`/`PackedRowwiseParallel`, которые шардируют каждый упакованный блок независимо.

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

Передайте `blocks` — список пропорциональных целых чисел, сумма которых делит упакованное измерение. Например, используйте `[1, 1]` для SwiGLU-гейта и проекции up с равными половинами, или `[1, 1, 1, 3, 3]` для объединённых проекций Q, K, V, gate и up при `mlp_ratio=3`.

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

Если размеры блоков известны только из конфигурации, опустите `blocks` и вместо этого храните абсолютные размеры на `Linear` во время `__init__` как `_tp_packed_col_blocks` или `_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(),
```

План в [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) использует обе формы. Пара double-stream SwiGLU передаёт `blocks` напрямую. Сплочённые single-stream проекции QKV и MLP хранят свои размеры в `Linear` внутри `Flux2ParallelSelfAttention.__init__`, поскольку они зависят от `mlp_ratio` и `mlp_mult_factor`.

### Что исключить

Модули, не включённые в план, дублируются на каждом ранге. Не включайте слои нормализации, эмбеддинги патчей и текста, а также финальные `norm_out` и `proj_out`. Они небольшие, поэтому их шардирование экономит мало памяти и увеличивает объём коммуникации. Модуляция AdaLN (`img_mod` и `txt_mod`) обычно достаточно мала, чтобы тоже не включать её в план. Если она большая, как `adaln_proj.linear` в MiniMax-H3, шардируйте её с `"rowwise_split_input"`.

### Ограничения и проверка

- `tp_degree` должен делить `config.num_attention_heads`.
- Каждый упакованный блок должен *отдельно* делиться на `tp_degree`, а не только их сумма.

[from_pretrained()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.from_pretrained) и `enable_parallelism()` проверяют оба условия до загрузки или разделения весов.

Численно проверяйте новый план, а не на глаз: сгенерируйте с фиксированным seed на одном устройстве, затем снова с tensor parallelism и сравните результаты. Неправильно размещённые `"colwise"`/`"rowwise"` обычно всё равно запускаются и дают правдоподобное, но неверное изображение.

> [!TIP]
> Начните с существующего плана для похожей архитектуры.
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) не содержит объединённых проекций, поэтому каждый элемент — это просто `"colwise"` или `"rowwise"`. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) добавляет одну упакованную row-wise проекцию. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) использует оба упакованных стиля.

Xet Storage Details

Size:: 15.6 kB
Xet hash:: 9487744b1247cfe86a25a9d62efad6b1337f2ae43259257da0355c7a389802af

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

	# Параллелизм тензоров

	[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) распределяет матрицы весов модели по устройствам. Каждое устройство хранит столбцовый (`"colwise"`) или строчный (`"rowwise"`) срез каждого слоя и вычисляет частичный результат. Операция `AllReduce` или `AllGather` на границе слоя объединяет частичные результаты в полный вывод. В отличие от [context parallelism](./context_parallelism), он снижает потребление памяти под weight на каждое устройство, что полезно, когда модель не помещается на одном устройстве. Для сравнения с другими стратегиями см. раздел [Distributed inference](../training/distributed_inference).

	## Загрузка модели с разбивкой на части

	Pass a [TensorParallelConfig](/docs/diffusers/v0.43.0/ru/api/parallel#diffusers.TensorParallelConfig) to the `parallel_config` argument of the model's [from_pretrained()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.from_pretrained). `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model's class must define a `_tp_plan`, which maps module names to a sharding style. To check whether a model supports tensor parallelism, look for this attribute on its class, such as [Flux2Transformer2DModel](/docs/diffusers/v0.43.0/ru/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

	Чтобы добавить план к другой модели, см. [Добавление тензорного параллелизма в модель](#add-tensor-parallelism-to-a-model).

	Загрузка таким образом разбивает контрольную точку на фрагменты по мере её чтения. Каждый ранг читает только свой срез каждого веса и размещает его непосредственно на своём GPU, поэтому ни один ранг никогда не хранит вес полного размера. Память на ранг уменьшается по мере увеличения `tp_degree`.

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

	Сохраните скрипт как `tensor_parallel_flux2.py` и запустите его с помощью `torchrun`.

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

	Скрипт устанавливает `tp_degree` в `world_size`, поэтому `--nproc-per-node 4` распределяет трансформер по 4 GPU.

	Шардируется только трансформер. Каждый ранг хранит полную копию текстового энкодера, поэтому на каждой GPU должно быть достаточно памяти для текстового энкодера и его доли трансформера.

	## Производительность загрузки

	Загрузка с `parallel_config` быстрее и потребляет меньше CPU-памяти на ранг, чем загрузка полной модели с последующим вызовом `enable_parallelism()`. В таблице ниже приведены результаты для [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (только трансформер, 32 млрд параметров, bf16) на 4x A10G (23 ГБ).

	\| tp_degree \| метод \| время загрузки \| пиковое значение CPU/ранк \|
	\|---\|---\|---\|---\|
	\| 4 \| `from_pretrained(parallel_config=...)` \| 12.5s \| 6,8 ГБ \|
	\| 4 \| `from_pretrained` + `enable_parallelism` \| 30.4s \| 64,1 ГБ \|

	## Ограничения

	Tensor parallelism не работает с некоторыми параметрами загрузки и функциями, а сохранить модель с разбивкой на шунки невозможно.

	- Загрузка с тензорно-параллельным `parallel_config` не работает с `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True` или весами, не хранящимися в формате safetensors. Каждый из этих параметров приводит к ошибке, а не к тихой загрузке полной контрольной точки.
	- Сочетание тензорного параллелизма с квантованием, offloading или адаптерами LoRA также приводит к ошибке, независимо от способа шардирования модели.
	- Чтобы разбить на шарды модель, уже загруженную в память, вызовите `enable_parallelism()` с той же конфигурацией. Сначала загружается полная модель, а затем она разбивается на шарды, поэтому каждому ранку требуется память для полного чекпоинта. Загрузка с `parallel_config` позволяет избежать этих затрат.
	- Сохранение тензорно-параллельной модели вызывает ошибку в [save_pretrained()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.save_pretrained). Сохраните модель до её шардирования.

	## Добавление тензорного параллелизма в модель

	Этот раздел предназначен для участников, которые хотят добавить поддержку тензорного параллелизма в модель Diffusers.

	Tensor parallelism работает только с моделями, которые определяют `_tp_plan` — атрибут класса, сопоставляющий имена модулей со стилем шардирования. В основном его создание сводится к тому, чтобы связать каждую проекцию, которая расширяет скрытое измерение, с проекцией, которая сжимает его обратно.

	Каждый ключ может содержать не более одного ``, а часть до `` должна указывать на [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), чтобы одна запись покрывала все блоки. Ключ без `*` применяется к самой модели. Пути указаны относительно модели.

	### По столбцам и по строкам

	\| Стиль \| Шарды \| Каждый ранг \| Использование \|
	\|---\|---\|---\|---\|
	\| `"colwise"` \| выходные признаки (`weight` размер 0) \| вычисляет срез выходных данных \| `to_q`, `to_k`, `to_v`, входная проекция FFN \|
	\| `"rowwise"` \| входные признаки (размерность `weight` 1) \| вычисляет частичную сумму \| `to_out.0`, выходная проекция FFN \|
	\| `"rowwise_split_input"` \| входные признаки (размерность `weight` 1) \| вычисляет частичную сумму из своего среза реплицированного входа \| независимая проекция, вход которой реплицирован, например, проекция модуляции AdaLN, читающая эмбеддинг временного шага \|

	Всегда используйте `"colwise"` и `"rowwise"` именно в таком порядке. Проекция `"colwise"` оставляет её вывод в разбитом виде, следующая проекция `"rowwise"` напрямую потребляет этот фрагмент, а один вызов `AllReduce` на границе блока восстанавливает результат. Любое другое разбиение пары приводит к необходимости собрать данные посередине и значительно увеличивает объём коммуникации.

	В механизме внимания столбчатые проекции Q, K и V распределяют подмножество голов по каждому рангу, поэтому `tp_degree` должно делить количество голов. Дубликаты энкодерного потока (`add_q_proj`, `to_add_out`, `ff_context`) следуют тому же паттерну, что и их аналоги в потоке изображений.

	### Сплоченные проекции

	Когда один `Linear` упаковывает несколько логических тензоров вдоль измерения, по которому выполняется шардирование, обычные срезы `"colwise"`/`"rowwise"` проходят прямо через конкатенацию и смещают фрагменты. Вместо этого используйте `PackedColwiseParallel`/`PackedRowwiseParallel`, которые шардируют каждый упакованный блок независимо.

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

	Передайте `blocks` — список пропорциональных целых чисел, сумма которых делит упакованное измерение. Например, используйте `[1, 1]` для SwiGLU-гейта и проекции up с равными половинами, или `[1, 1, 1, 3, 3]` для объединённых проекций Q, K, V, gate и up при `mlp_ratio=3`.

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

	Если размеры блоков известны только из конфигурации, опустите `blocks` и вместо этого храните абсолютные размеры на `Linear` во время `__init__` как `_tp_packed_col_blocks` или `_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(),
	```

	План в [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) использует обе формы. Пара double-stream SwiGLU передаёт `blocks` напрямую. Сплочённые single-stream проекции QKV и MLP хранят свои размеры в `Linear` внутри `Flux2ParallelSelfAttention.__init__`, поскольку они зависят от `mlp_ratio` и `mlp_mult_factor`.

	### Что исключить

	Модули, не включённые в план, дублируются на каждом ранге. Не включайте слои нормализации, эмбеддинги патчей и текста, а также финальные `norm_out` и `proj_out`. Они небольшие, поэтому их шардирование экономит мало памяти и увеличивает объём коммуникации. Модуляция AdaLN (`img_mod` и `txt_mod`) обычно достаточно мала, чтобы тоже не включать её в план. Если она большая, как `adaln_proj.linear` в MiniMax-H3, шардируйте её с `"rowwise_split_input"`.

	### Ограничения и проверка

	- `tp_degree` должен делить `config.num_attention_heads`.
	- Каждый упакованный блок должен отдельно делиться на `tp_degree`, а не только их сумма.

	[from_pretrained()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.from_pretrained) и `enable_parallelism()` проверяют оба условия до загрузки или разделения весов.

	Численно проверяйте новый план, а не на глаз: сгенерируйте с фиксированным seed на одном устройстве, затем снова с tensor parallelism и сравните результаты. Неправильно размещённые `"colwise"`/`"rowwise"` обычно всё равно запускаются и дают правдоподобное, но неверное изображение.

	> [!TIP]
	> Начните с существующего плана для похожей архитектуры.
	>
	> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) не содержит объединённых проекций, поэтому каждый элемент — это просто `"colwise"` или `"rowwise"`. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) добавляет одну упакованную row-wise проекцию. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) использует оба упакованных стиля.