diffusers/main/ru/api/parallel.md

Параллелизм

Стратегии параллелизма помогают ускорить работу диффузионных трансформеров за счёт распределения вычислений по нескольким устройствам, что позволяет сократить время инференса/обучения. Подробнее см. в руководстве Distributed inference.

ParallelConfig[[diffusers.ParallelConfig]]

diffusers.ParallelConfig[[diffusers.ParallelConfig]]

diffusers.ParallelConfig(context_parallel_config: diffusers.models._modeling_parallel.ContextParallelConfig | None = None, tensor_parallel_config: diffusers.models._modeling_parallel.TensorParallelConfig | None = None, _rank: int = None, _world_size: int = None, _device: device = None, _mesh: DeviceMesh = None)

Источник

Параметры:

context_parallel_config (ContextParallelConfig, необязательный) : Конфигурация для параллелизма контекста.

tensor_parallel_config (TensorParallelConfig, необязательный) : Конфигурация для тензорного параллелизма.

Конфигурация для применения различных видов параллелизма.

ContextParallelConfig[[diffusers.ContextParallelConfig]]

diffusers.ContextParallelConfig[[diffusers.ContextParallelConfig]]

diffusers.ContextParallelConfig(ring_degree: int | None = None, ulysses_degree: int | None = None, convert_to_fp32: bool = True, rotate_method: typing.Literal['allgather', 'alltoall'] = 'allgather', mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, ulysses_anything: bool = False, ring_anything: bool = False, _rank: int = None, _world_size: int = None, _device: device = None, _mesh: DeviceMesh = None, _flattened_mesh: DeviceMesh = None, _ring_mesh: DeviceMesh = None, _ulysses_mesh: DeviceMesh = None, _ring_local_rank: int = None, _ulysses_local_rank: int = None)

Источник

Параметры:

ring_degree (int, optional, по умолчанию 1) : Количество устройств для использования в Ring Attention. Последовательность распределяется по устройствам. Каждое устройство вычисляет внимание между своим локальным Q и KV-фрагментами, которые последовательно передаются по кольцу. Это снижает потребление памяти (одновременно хранится только 1/N KV), позволяет перекрывать вычисления с коммуникацией, но требует N итераций для обработки всех токенов. Лучше всего подходит для длинных последовательностей при ограниченной памяти/пропускной способности. Количество устройств для ring attention внутри области контекстного параллелизма. Должно быть делителем общего числа устройств в сетке контекстного параллелизма.

ulysses_degree (int, optional, по умолчанию 1) : Количество устройств для использования в Ulysses Attention. Последовательность распределяется по устройствам. Каждое устройство вычисляет локальные QKV, затем с помощью all-gather собирает все фрагменты KV для вычисления полного внимания за один проход. Требует больше памяти (хранит все KV) и высокоскоростной all-to-all коммуникации, но обеспечивает низкую задержку. Лучше всего подходит для последовательностей средней длины при наличии высокопропускной межпроцессорной связи.

convert_to_fp32 (bool, необязательный, по умолчанию True) : Преобразовывать ли выходные данные и LSE в float32 для численной стабильности ring attention.

rotate_method (str, необязательный, по умолчанию "allgather") : Метод, используемый для ротации состояний ключей/значений между устройствами в ring attention. В настоящее время поддерживается только "allgather".

ulysses_anything (bool, необязательный, по умолчанию False) : Включает ли режим "Ulysses Anything", который поддерживает произвольные длины последовательностей и количества головок, не делящиеся нацело на ulysses_degree. При включении ulysses_degree должно быть больше 1, а ring_degree — равно 1.

ring_anything (bool, необязательный, по умолчанию False) : Включает ли режим "Ring Anything", который поддерживает произвольные длины последовательностей. При включении ring_degree должно быть больше 1, а ulysses_degree — равно 1.

mesh (torch.distributed.device_mesh.DeviceMesh, optional) : Пользовательская сетка устройств для использования в контекстном параллелизме. Если указана, эта сетка будет использоваться вместо создания новой. Это полезно при сочетании контекстного параллелизма с другими стратегиями параллелизма (например, FSDP, тензорным параллелизмом), которые используют одну и ту же сетку устройств. Сетка должна содержать измерения "ring" и "ulysses". Для неиспользуемых измерений используйте размер 1 (например, mesh_shape=(2, 1, 4) с mesh_dim_names=("ring", "ulysses", "fsdp") для ring attention только с FSDP).

Конфигурация для параллелизма контекста.

diffusers.hooks.apply_context_parallel[[diffusers.hooks.apply_context_parallel]]

diffusers.hooks.apply_context_parallel(module: Module, parallel_config: ContextParallelConfig, plan: dict[str, dict[str, dict[str | int, diffusers.models._modeling_parallel.ContextParallelInput | list[diffusers.models._modeling_parallel.ContextParallelInput] | tuple[diffusers.models._modeling_parallel.ContextParallelInput, ...]] | diffusers.models._modeling_parallel.ContextParallelOutput | list[diffusers.models._modeling_parallel.ContextParallelOutput] | tuple[diffusers.models._modeling_parallel.ContextParallelOutput, ...] | diffusers.models._modeling_parallel.Conditional]])

Источник

Примените контекстный параллелизм к модели.

TensorParallelConfig[[diffusers.TensorParallelConfig]]

diffusers.TensorParallelConfig[[diffusers.TensorParallelConfig]]

diffusers.TensorParallelConfig(tp_degree: int = 1, mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, _rank: int = None, _world_size: int = None, _device: device = None, _mesh: DeviceMesh = None, _tp_degree: int = None)

Источник

Параметры:

tp_degree (int, по умолчанию 1) : Количество устройств, на которые выполняется шардирование. Должно быть делителем количества головок внимания (и скрытых размерностей FFN) модели, для которой настраивается параллелизм.

mesh (torch.distributed.device_mesh.DeviceMesh, необязательный) : Пользовательская сетка устройств для использования. Если указано, tp_degree выводится из mesh.size(), а аргумент игнорируется. Полезно при сочетании TP с другими стратегиями параллелизма (например, CP), использующими одну и ту же сетку.

Конфигурация для параллелизма тензоров.

Тензорное параллелизм распределяет матрицы весов (по столбцам и по строкам) между устройствами. Каждое устройство вычисляет частичный результат; AllReduce/AllGather на границах слоёв восстанавливает полный выход. Использует torch.distributed.tensor.parallelize_module со стилями шардинга ColwiseParallel / RowwiseParallel. Поддерживаемые типы устройств: "cuda", "neuron" и "tpu".

diffusers.hooks.apply_tensor_parallel[[diffusers.hooks.apply_tensor_parallel]]

diffusers.hooks.apply_tensor_parallel(model: Module, config: TensorParallelConfig, tp_plan: dict, weights_already_sharded: bool = False)

Источник

Параметры:

model (torch.nn.Module) : модель, которую нужно разбить на части на месте.

config (TensorParallelConfig) : конфигурация тензорного параллелизма. Её сетка устройств должна уже быть настроена с помощью config.setup(...).

tp_plan (dict) : Плоское сопоставление глоб-имён модулей со стилем "colwise"/"rowwise" (или упакованным вариантом), как правило, это _tp_plan модели.

weights_already_sharded (bool, по умолчанию False) : Указывает, являются ли запланированные параметры уже фрагментами DTensor, как это бывает после потоковой загрузки с помощью from_pretrained. Если True, регистрируются только forward-хуки. Этот параметр передаётся явно, а не определяется автоматически, потому что запланированный параметр, отсутствующий в контрольной точке, всё равно останется мета-тензором, и автоматическое определение ошибочно сочтёт модель «не разбитой на фрагменты», хотя на самом деле она разбита наполовину.

Примените тензорное параллелизм к модели из её плоского _tp_plan.