Тензорное параллелизм для инференса

Tensor parallelism разбивает слой модели на части, чтобы несколько аппаратных ускорителей могли обрабатывать его одновременно. Это позволяет запускать модели, размер которых превышает объём памяти одного GPU, и достигать более высокой пропускной способности. Для этого требуется быстрая коммуникация внутри узла, поскольку GPU обмениваются промежуточными результатами на каждом слое.

Модель поддерживает тензорный параллелизм, если её конфигурация определяет base_model_tp_plan. Проверить загруженную модель можно с помощью свойства supports_tp_plan.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

В этом руководстве описано, как включить параллелизм тензоров в Transformers, а также какие стратегии разделения доступны.

Разбиение модели

Настройте количество устройств для тензорного параллелизма с помощью tp_size в DistributedConfig.

• Установите DistributedConfig(tp_size=N), чтобы использовать предопределённый план модели.
• Определите ручной tp_plan и передайте его в DistributedConfig вместе с tp_size.

Вы также можете установить tp_plan="auto", чтобы явно запросить предопределённый план. Если tp_size не указан, а tp_plan задан, то tp_size вычисляется как WORLD_SIZE, делённое на остальные размеры параллелизма. Передача tp_plan непосредственно в from_pretrained() устарела и будет удалена в v5.18.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)
print(model.tp_plan)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

# distributed run
outputs = model(inputs)

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

from transformers import AutoModelForCausalLM, DistributedConfig

tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
    ...
}

distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype="auto",
    distributed_config=distributed_config,
)
print(model.tp_plan)

Стратегии разделения

ALL_PARALLEL_STYLES — это ParallelInterface, определённый в distributed/parallel_utils.py, который сопоставляет каждое имя стратегии, доступное в плане, с настроенным экземпляром стратегии. Вы не взаимодействуете с ним напрямую для шардирования модели, но это авторитетный список доступных имён. distributed/tensor_parallel.py регистрирует в нём стратегии тензорного параллелизма:

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

Каждая стратегия тензорного параллелизма является подклассом TensorParallelLayer в этом файле. Стратегии экспертного параллелизма, такие как ep_dispatch_experts, находятся в distributed/expert_parallel.py, наследуются от ExpertParallelLayer, который имеет те же хуки, и регистрируются в том же ALL_PARALLEL_STYLES. Две стратегии, которые вы будете использовать чаще всего, ColwiseParallel и RowwiseParallel, принимают размещения input_layouts и output_layouts, что позволяет одному классу покрывать несколько имён планов. colwise оставляет выходной тензор разбитым по последнему измерению для последующего слоя rowwise, тогда как colwise_gather_output собирает его обратно в полный тензор с помощью all-gather.

В таблице ниже описана каждая стратегия.

Стратегии упаковки

Упаковка весов объединяет несколько линейных слоёв в один, более крупный слой. Стратегии PackedColwiseParallel и PackedRowwiseParallel корректно распределяют упакованные веса. Базовые стратегии ColwiseParallel или RowwiseParallel распределяют упакованные веса некорректно.

Приведённый ниже пример объединяет up_proj и gate_proj в единый модуль gate_up_proj и требует использование стратегии packed_rowwise для шардирования gate_up_proj.

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

Используйте пакетное матричное умножение в проходе forward для вычисления выходных данных модуля gate_up_proj.

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

Обычный Shard разбивает это измерение на непрерывные блоки, поэтому ранг получает хвост gate и начало up, а не срез каждой из частей. Пакетные стратегии используют _StridedShard с split_factor (по умолчанию 2), чтобы чередовать разбиение, обеспечивая каждому рангу соответствующий срез обеих половин, благодаря чему chunk остаётся согласованным после шардирования.

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

Наследуйте от TensorParallelLayer в distributed/tensor_parallel.py, чтобы создать собственную стратегию разделения. Переопределяйте только те хуки, которые нужны вашей стратегии, так как у каждого из них есть реализация по умолчанию, не выполняющая никаких действий.

Приведённый ниже пример разбирает упрощённую версию ColwiseParallel.

1. Наследуйте от TensorParallelLayer и храните размещения, с которыми работает стратегия. Базовый класс не определяет __init__, поэтому вызывать super() не нужно.

   Скопировано

   class ColwiseParallel(TensorParallelLayer):
       def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
           self.input_layouts = input_layouts or Replicate()
           self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
           self.use_local_output = use_local_output

2. Реализуйте shard_param, чтобы обернуть один параметр в плейсхолдер DTensor. Он работает с мета-тензорами, поэтому distribute_tensor лишь формирует метаданные и не перемещает данные. Передайте src_data_rank=None, так как полного тензора для рассредоточения пока нет.

   Скопировано

   def shard_param(self, module, param, mesh):
       meta = module._parameters.get(param)
       if meta is None:
           return
       # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
       placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
       module._parameters[param] = torch.nn.Parameter(
           distribute_tensor(meta, mesh, [placement], src_data_rank=None),
           requires_grad=meta.requires_grad,
       )

3. Реализуйте преобразования входных и выходных данных. install_forward вызывает их вокруг исходного forward модуля, поэтому им достаточно лишь перемещать тензоры между раскладками. Разбиение по столбцам предполагает реплицированный вход и формирует выход, разбитый по последнему измерению.

   Скопировано

   def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
       x = args[0]
       if not isinstance(x, DTensor):
           x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
       if x.placements != (Replicate(),):
           x = x.redistribute(placements=[Replicate()])
       return (x,) + args[1:], kwargs
   
   def transform_output_post_forward(self, module, output, mesh):
       if not isinstance(output, DTensor):
           output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
       if output.placements != (self.output_layouts,):
           output = output.redistribute(placements=[self.output_layouts])
       return output.to_local() if self.use_local_output else output

   Встроенный ColwiseParallel добавляет быстрые пути поверх этого, которые пропускают цикл DTensor для обычного вывода nn.Linear и для квантованных модулей, которым нужны локальные тензоры. Если вашей стратегии требуется такая же обработка, изучите исходный код перед тем, как его копировать.

4. Зарегистрируйте стратегию, чтобы tp_plan мог ссылаться на неё по имени. При регистрации передаётся экземпляр стратегии, поэтому размещения, переданные в __init__, сохраняются вместе с именем. Обратите внимание на () ниже.

   Скопировано

   import torch
   
   from transformers import AutoModelForCausalLM, DistributedConfig
   from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES
   
   ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
   tp_plan = {
       "model.layers.*.self_attn.q_proj": "colwise_custom",
       ...
   }
   distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
   model = AutoModelForCausalLM.from_pretrained(
       "meta-llama/Meta-Llama-3-8B-Instruct",
       dtype=torch.bfloat16,
       distributed_config=distributed_config,
   )

Бенчмарки

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

На этом графике показано ожидаемое ускорение одного прямого прохода для Llama при длине последовательности 512.

Реализация дизайна

Transformers реализует тензорное параллелизм независимо от фреймворка. Он опирается на DeviceMesh и DTensor из torch.distributed, чтобы обеспечить простой и расширяемый интерфейс.

DeviceMesh

DeviceMesh создаёт многомерную сетку устройств, которые взаимодействуют между собой. Разные стратегии параллелизации требуют разных паттернов коммуникации. Создайте DeviceMesh с несколькими подсетками, чтобы обрабатывать эти паттерны.

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

Большинство стратегий параллелизации torch.distributed применяются к самой сетке или её подсетке. Сетка автоматически обрабатывает паттерны коммуникации.

DTensor

DTensor (Distributed Tensor) обрабатывает распределённую логику поверх обычных тензорных операций. Веса модели при тензорном параллелизме хранятся в виде DTensor, что позволяет стратегии описывать коммуникацию как смену раскладки, а не как явный коллективный вызов.

Атрибут placements указывает PyTorch, как тензор распределён по устройствам в DeviceMesh. Он принимает следующие значения:

• Shard(dimension) распределяет DTensor по указанному измерению в рамках DeviceMesh, в котором он был создан. Разбиение по столбцам разделяет измерение выходных признаков веса и единственное измерение смещения.

  Скопировано

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension

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

  Скопировано

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs

• Replicate() дублирует DTensor по DeviceMesh, создавая полную копию тензора на каждом устройстве.

• Partial() помечает тензор как ожидающий редукцию. Выход forward слоя, работающего построчно, имеет тип Partial(), а его перераспределение в Replicate() инициирует all-reduce.

Ресурсы

• Раздел Ultra-Scale Playbook о тензорном параллелизме содержит более подробную информацию.

• Если вы используете модель типа mixture-of-experts (MoE), обратитесь к руководству по expert parallelism. Такие модели поддерживают tensor parallelism и expert parallelism.

• Ознакомьтесь с блогом Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand, чтобы быстро понять, что такое тензорная параллельность, и узнать, чем отличаются конфигурации column и row parallel.

• См. руководство по обучению Tensor parallelism, чтобы узнать, как использовать его в режиме обучения.