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

Tensor parallelism распределяет матрицы весов модели по устройствам. Каждое устройство хранит столбцовый ("colwise") или строчный ("rowwise") срез каждого слоя и вычисляет частичный результат. Операция AllReduce или AllGather на границе слоя объединяет частичные результаты в полный вывод. В отличие от context parallelism, он снижает потребление памяти под weight на каждое устройство, что полезно, когда модель не помещается на одном устройстве. Для сравнения с другими стратегиями см. раздел Distributed inference.

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

Pass a TensorParallelConfig to the parallel_config argument of the model’s 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.

Чтобы добавить план к другой модели, см. Добавление тензорного параллелизма в модель.

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

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.

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 (только трансформер, 32 млрд параметров, bf16) на 4x A10G (23 ГБ).

Ограничения

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(). Сохраните модель до её шардирования.

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

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

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

Каждый ключ может содержать не более одного *, а часть до * должна указывать на nn.ModuleList, чтобы одна запись покрывала все блоки. Ключ без * применяется к самой модели. Пути указаны относительно модели.

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

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

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

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

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

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

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

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

План в 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() и enable_parallelism() проверяют оба условия до загрузки или разделения весов.

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

Начните с существующего плана для похожей архитектуры.

• В файле transformer_qwenimage.py нет слиянных проекций, поэтому каждая запись — это просто "colwise" или "rowwise".
• Добавляет одно упакованное построчное проекционное преобразование в transformer_flux.py.
• transformer_flux2.py использует оба стиля упаковки.