Распределённый вывод

Распределённый вывод распределяет нагрузку между несколькими GPU. Это полезная техника для размещения более крупных моделей в памяти, позволяющая обрабатывать несколько запросов для повышения пропускной способности.

В этом руководстве показано, как использовать Accelerate и PyTorch Distributed для распределённого вывода.

Accelerate

Accelerate — это библиотека, предназначенная для упрощения инференса и обучения на нескольких ускорителях за счет автоматизации настройки, что позволяет пользователям сосредоточиться на своем коде PyTorch.

Установите Accelerate с помощью следующей команды.

uv pip install accelerate

Инициализируйте класс accelerate.PartialState в файле Python, чтобы создать распределённую среду. Класс accelerate.PartialState управляет процессами, контролем и распределением устройств, а также координацией процессов.

Переместите DiffusionPipeline на accelerate.PartialState.device, чтобы назначить GPU каждому процессу.

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

Используйте утилиту split_between_processes в качестве контекстного менеджера для автоматического распределения промптов между процессами.

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

Вызовите accelerate launch, чтобы запустить скрипт, и используйте аргумент --num_processes, чтобы задать количество используемых GPU.

accelerate launch run_distributed.py --num_processes=2

См. этот минимальный скрипт для запуска инференса на нескольких GPU. Чтобы узнать больше, ознакомьтесь с руководством Distributed Inference with 🤗 Accelerate.

PyTorch Distributed

PyTorch DistributedDataParallel обеспечивает data parallelism, при котором одна и та же модель дублируется на каждом устройстве для параллельной обработки разных батчей данных.

Импортируйте torch.distributed и torch.multiprocessing в файл Python, чтобы настроить группу распределённых процессов и запустить процессы для вывода на каждом GPU.

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16,
)

Создайте функцию для инференса с init_process_group. Этот метод создаёт распределённую среду с типом бэкенда, rank текущего процесса и world_size или количеством участвующих процессов (например, 2 GPU будут world_size=2).

Переместите конвейер на rank и используйте get_rank, чтобы назначить GPU каждому процессу. Каждый процесс обрабатывает отдельный промпт.

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    pipeline.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")

Используйте mp.spawn для создания количества процессов, заданного в world_size.

def main():
    world_size = 2
    mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()

Вызовите torchrun, чтобы запустить скрипт инференса, и используйте аргумент --nproc_per_node, чтобы задать количество используемых GPU.

torchrun --nproc_per_node=2 run_distributed.py

device_map

Аргумент device_map обеспечивает распределённый вывод, автоматически размещая компоненты модели на разных GPU. Это особенно полезно, когда модель не помещается на одном GPU. С помощью device_map можно выборочно загружать и выгружать необходимые компоненты модели на каждом этапе, как показано в примере ниже (предполагается наличие двух GPU).

Установите device_map="balanced", чтобы равномерно распределить текстовые энкодеры по всем доступным GPU. Аргумент max_memory позволяет задать максимальный объём памяти для каждого текстового энкодера. Не загружайте другие компоненты конвейера, чтобы избежать лишнего потребления памяти.

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

После вычисления текстовых эмбеддингов удалите их с GPU, чтобы освободить место для диффузионного трансформера.

import gc 

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

Установите device_map="auto", чтобы автоматически распределить модель по двум GPU. Эта стратегия размещает модель сначала на самом быстром устройстве, а при необходимости — на более медленных, таких как CPU или жесткий диск. Недостатком хранения параметров модели на медленных устройствах является увеличение задержки при выводе.

from diffusers import AutoModel
import torch 

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)

Выполните pipeline.hf_device_map, чтобы увидеть, как различные модели распределены по устройствам. Это полезно для отслеживания размещения моделей на устройствах. Вы также можете вызвать hf_device_map для трансформерной модели, чтобы увидеть, как она распределена.

Добавьте модель трансформера в конвейер и установите output_type="latent", чтобы генерировать латенты.

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

Удалите pipeline и transformer из памяти и загрузите VAE для декодирования латентов. VAE обычно достаточно мал, чтобы быть загруженным на одно устройство.

import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", dtype=torch.bfloat16).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")

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

Параллелизм контекста

Контекстная параллельность распределяет входные последовательности между несколькими GPU для снижения потребления памяти. Каждый GPU обрабатывает свой срез последовательности.

Используйте set_attention_backend(), чтобы переключиться на более оптимизированный бэкенд внимания. Полный список доступных бэкендов приведен в этой таблице.

Большинство бэкендов внимания совместимы с контекстным параллелизмом. Если бэкенд несовместим, создайте issue.

Ring Attention

Представления ключей (K) и значений (V) обмениваются между устройствами с помощью Ring Attention. Это гарантирует, что каждый фрагмент видит K/V всех остальных токенов. Каждый GPU вычисляет внимание для своих локальных K/V и передает результат следующему GPU в кольце. Ни один GPU не хранит полную последовательность, что снижает задержку коммуникации.

Передайте ContextParallelConfig в аргумент parallel_config трансформерной модели. Конфигурация поддерживает аргумент ring_degree, который определяет, сколько устройств использовать для Ring Attention.

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save(f"output.png")

    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

Скрипт выше необходимо запускать с помощью распределённого лаунчера, например torchrun, совместимого с PyTorch. Параметр --nproc-per-node устанавливается в значение количества доступных GPU.

torchrun --nproc-per-node 2 above_script.py

Ulysses Attention

Ulysses Attention splits a sequence across GPUs and performs an all-to-all communication (every device sends/receives data to every other device). Each GPU ends up with all tokens for only a subset of attention heads. Each GPU computes attention locally on all tokens for its head, then performs another all-to-all to regroup results by tokens for the next layer.

ContextParallelConfig поддерживает Ulysses Attention через аргумент ulysses_degree. Он определяет, сколько устройств использовать для Ulysses Attention.

Передайте ContextParallelConfig в enable_parallelism().

# Depending on the number of GPUs available.
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

Единое внимание

Unified Sequence Parallelism объединяет Ring Attention и Ulysses Attention в единый подход для эффективной обработки длинных последовательностей. Сначала применяется коммуникация all-to-all из Ulysses для перераспределения голов и токенов последовательности, затем Ring Attention используется для обработки перераспределённых данных, а в конце all-to-all выполняется в обратном порядке для восстановления исходной раскладки.

Этот гибридный подход использует преимущества обоих методов:

• Ulysses Attention эффективно параллелизует вычисления по головам внимания
• Ring Attention обрабатывает очень длинные последовательности с минимальными накладными расходами по памяти
• Вместе они обеспечивают 2D-параллелизацию как по размерности голов, так и по размерности последовательности

ContextParallelConfig поддерживает Unified Attention за счёт указания обоих параметров ulysses_degree и ring_degree. Общее количество используемых устройств равно ulysses_degree * ring_degree; они образуют двумерную сетку, в которой группы Ulysses и Ring ортогональны (не пересекаются). Передайте ContextParallelConfig с обоими параметрами ulysses_degree и ring_degree, установленными в значения больше 1, в enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

Unified Attention следует использовать, когда имеется достаточно устройств для их размещения в 2D-сетке (не менее 4 устройств).

Мы провели бенчмарк с Ulysess, Ring и Unified Attention с помощью этого скрипта на узле с 4 GPU H100. Результаты представлены ниже:

Из приведённой выше таблицы видно, что Ulysses обеспечивает более высокую пропускную способность, однако количество устройств, которые он может использовать, по-прежнему ограничено числом attention heads. Это ограничение устраняется с помощью unified attention.

Ulysses Anything Attention

Механизм Ulysses Attention по умолчанию требует, чтобы длина последовательности скрытых состояний делилась на количество устройств. Это налагает существенные ограничения на практическое применение Ulysses Attention. Ulysses Anything Attention — это вариант Ulysses Attention, поддерживающий произвольные длины последовательностей и произвольное количество головок внимания, что повышает универсальность Ulysses Attention в реальных сценариях использования.

ContextParallelConfig поддерживает Ulysses Anything Attention за счёт указания обоих параметров ulysses_degree и ulysses_anything. Обратите внимание, что Ulysses Anything Attention в настоящее время не поддерживается в Unified Attention. Передайте ContextParallelConfig с ulysses_degree, установленным в значение больше 1, и ulysses_anything=True в enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

Чтобы избежать нескольких принудительных синхронизаций CUDA, вызванных передачами H2D и D2H, добавьте бэкенд gloo в init_process_group. Это значительно снизит задержки коммуникации.

Мы провели бенчмарк FLUX.1-dev с Ulysses, Ring, Unified Attention и Ulysses Anything Attention с помощью этого скрипта на узле с 4 GPU L20. Результаты представлены ниже:

Из приведённой выше таблицы видно, что Ulysses Anything Attention обеспечивает лучшую совместимость с произвольными длинами последовательностей, сохраняя при этом ту же производительность, что и стандартный Ulysses Attention.

Ring Anything Attention

По умолчанию Ring Attention требует, чтобы длина последовательности скрытых состояний делилась нацело на степень кольца. Ring Anything Attention — это вариант Ring Attention, поддерживающий произвольные (не делящиеся нацело) длины последовательностей. Он дополняет локальные KV каждого ранга до глобальной максимальной длины последовательности, выполняет all-gather для дополненного буфера KV и затем обрезает его до истинной длины каждого ранга перед выполнением внимания.

Ring Anything Attention не поддерживается в Unified Attention. Чтобы включить Ring Anything Attention, установите ring_degree > 1 и ring_anything=True.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Добавьте бэкенд gloo в init_process_group, чтобы избежать нескольких принудительных синхронизаций CUDA при передачах H2D и D2H.

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

Ring Anything Attention в настоящее время поддерживает только инференс, а маски внимания, отличные от None, не поддерживаются. attn_mask должен быть None.

См. ниже бенчмарки FLUX.1-dev на узле с 4 GPU RTX 4090 (48 ГБ).

Из приведённой выше таблицы следует, что Ring Anything Attention обеспечивает совместимость с последовательностями произвольной длины, сохраняя при этом производительность, сопоставимую со стандартным Ring Attention. Подробнее о мотивации и компромиссах Ring Anything Attention см. в этом комментарии.

parallel_config

Передайте parallel_config при инициализации модели, чтобы включить контекстную параллельность.

CKPT_ID = "black-forest-labs/FLUX.1-dev"

cp_config = ContextParallelConfig(ring_degree=2)
transformer = AutoModel.from_pretrained(
    CKPT_ID, 
    subfolder="transformer", 
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)

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

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

Передайте TensorParallelConfig в аргумент parallel_config метода from_pretrained() модели. tp_degree — это количество устройств, на которые выполняется разбиение, и оно должно делить количество attention-голов модели. Модель должна определять _tp_plan (плоское отображение глобальных имён модулей на стиль "colwise"/"rowwise").

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

По сравнению с загрузкой полной модели и последующим вызовом enable_parallelism(), этот способ загружается быстрее и потребляет меньше CPU-памяти на каждый ранк, при этом разрыв увеличивается с ростом tp_degree. Приведённые ниже цифры относятся к black-forest-labs/FLUX.2-dev (только трансформер, 32 млрд параметров, bf16) на 4x A10G (23 ГБ).

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()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto `cuda:rank`.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, torch_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()

torchrun --nproc-per-node 4 tensor_parallel_flux.py

tp_degree берётся из world_size выше, поэтому --nproc-per-node 4 распределяет трансформер по 4 устройствам.

Загрузка с тензорно-параллельным parallel_config пока не поддерживается вместе с device_map, quantization_config, low_cpu_mem_usage=False, use_flashpack=True или весами, не хранящимися в формате safetensors; в каждом из этих случаев возникает ошибка, а не происходит тихого перехода к загрузке полной контрольной точки.

Комбинирование тензорного параллелизма с квантованием, offloading или адаптерами LoRA пока не поддерживается, поэтому эти методы вызывают ValueError независимо от способа шардирования модели.

Мы рекомендуем загружать тензорно-параллельные модели именно таким образом. Вы также можете разбить модель, уже находящуюся в памяти, вызвав enable_parallelism() с той же конфигурацией, но при этом сначала загружается весь чекпоинт, а затем он переформатируется, что приводит к полному потреблению памяти чекпоинта на каждом ранке.

Сохранение тензорно-параллельной модели пока не поддерживается, и save_pretrained() вызывает ошибку при попытке сохранить такую модель. Сохраните модель до её разделения на шарды.

Создание плана параллелизма тензоров

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

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

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

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

Для механизма внимания это означает, что каждый ранг владеет подмножеством головок, поэтому 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] для проекции gate+up в SwiGLU с равными половинами или [1, 1, 1, 3, 3] для объединённой проекции Q+K+V+gate+up с mlp_ratio=3.

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

Если размеры блоков известны только из конфигурации, опустите аргумент и вместо этого сохраните абсолютные размеры на 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.

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

Всё, что отсутствует в плане, остаётся реплицированным на каждом ранге, что является правильным выбором для слоёв нормализации, модуляции AdaLN (img_mod/txt_mod), эмбеддингов патчей и текста, а также финальных norm_out/proj_out. Эти компоненты малы, поэтому их шардирование экономит мало памяти, но увеличивает объём коммуникации.

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

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

Оба проверяются с помощью from_pretrained() и enable_parallelism() до загрузки или разделения весов.

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

Начните с существующего плана для похожей архитектуры. QwenImageTransformer2DModel полностью неслитая, и каждая запись — это обычный "colwise"/"rowwise", FluxTransformer2DModel добавляет одну упакованную row-wise проекцию, а Flux2Transformer2DModel покрывает оба упакованных стиля — см. transformer_qwenimage.py, transformer_flux.py и transformer_flux2.py соответственно.

Выбор стратегии

Стратегии, описанные выше, решают разные задачи, и полезный вопрос заключается не в том, какая из них быстрее в абстрактном смысле, а в том, какого ресурса вам не хватает.

Некоторые практические рекомендации:

• Пропускная способность при работе с множеством запросов, модель уже помещается в память. Используйте параллелизм данных. Это единственная стратегия, которая линейно масштабирует пропускную способность без изменения модели и не влияет на задержку для одиночного запроса.
• Веса одного компонента не помещаются в память. В первую очередь стоит обратиться к тензорному параллелизму, так как он снижает как потребление памяти, так и задержку. Он требует обмена данными на каждой границе блока, поэтому для его эффективной работы нужен быстрый интерфейс, например NVLink; при использовании PCIe трафик на каждом слое может оказаться больше, чем вычислительная нагрузка, которую он экономит, и в этом случае device_map становится лучшим выбором. device_map также решает проблему, когда компоненты по отдельности помещаются в память, но их совокупный размер слишком велик.
• Проблема не в весах, а в активациях. Это режим длинных последовательностей — большие изображения, множество кадров — и контекстная параллельность является прямым решением. Выберите бэкенд из бенчмарков Ulysses/Ring выше.
• И веса, и последовательность слишком велики. Сочетайте тензорный и контекстный параллелизм. TensorParallelConfig принимает аргумент mesh, что позволяет использовать одну сетку устройств для обоих. Эта комбинация пока не проверена, поэтому рассматривайте её как экспериментальную.