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

[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) использует оба упакованных стиля.