Paralelismo de tensores para inferencia

Paralelismo de tensores divide una capa del modelo en fragmentos para que varios aceleradores de hardware trabajen en ella simultáneamente. Esto permite ejecutar modelos que superan la capacidad de memoria de una sola GPU y lograr un mayor rendimiento. Se requiere una comunicación rápida dentro del nodo, ya que las GPUs intercambian resultados parciales en cada capa.

Un modelo admite la paralelización de tensores si su configuración define base_model_tp_plan. Compruebe un modelo cargado con la propiedad supports_tp_plan.

from transformers import AutoModelForCausalLM

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

Esta guía explica cómo habilitar el paralelismo de tensores en Transformers y las estrategias de partición disponibles.

Particionamiento de un modelo

Configura el número de dispositivos de paralelismo de tensores con tp_size en DistributedConfig.

Establezca DistributedConfig(tp_size=N) para utilizar el plan predefinido del modelo.
Defina un tp_plan manual y páselo a DistributedConfig junto con tp_size.

También puede establecer tp_plan="auto" para solicitar explícitamente el plan predefinido. Cuando se omite tp_size y se establece un tp_plan, tp_size se deriva de WORLD_SIZE dividido por los demás tamaños de paralelismo. Pasar tp_plan directamente a from_pretrained() está deprecado y se eliminará en v5.18.

auto plan

plan manual

Estrategias de particionamiento

ALL_PARALLEL_STYLES, una ParallelInterface definida en distributed/parallel_utils.py, asigna cada nombre de estrategia que puedes usar en un plan a una instancia de estrategia configurada. No interactúas con ella directamente para particionar un modelo, pero es la lista autoritativa de nombres disponibles. distributed/tensor_parallel.py registra las estrategias de paralelismo de tensor en ella:

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())

Cada estrategia de paralelismo de tensores es una subclase de TensorParallelLayer en ese archivo. Las estrategias de paralelismo de expertos, como ep_dispatch_experts, se encuentran en distributed/expert_parallel.py, heredan de ExpertParallelLayer, que tiene los mismos hooks, y se registran en el mismo ALL_PARALLEL_STYLES. Las dos que más utilizarás, ColwiseParallel y RowwiseParallel, aceptan colocaciones de input_layouts y output_layouts, lo que permite que una sola clase cubra varios nombres de plan. colwise deja su salida fragmentada en la última dimensión para que una capa rowwise posterior la consuma, mientras que colwise_gather_output la reúne mediante all-gather en un tensor completo.

La tabla siguiente describe cada estrategia.

Estrategia	Descripción
`ColwiseParallel`	Fragmenta un peso lineal en su dimensión de características de salida (`Shard(0)` para un peso 2D de `nn.Linear`) y fragmenta un sesgo 1D. Redistribuye la entrada a `Replicate()` y luego coloca la salida en `output_layouts`.
`RowwiseParallel`	Fragmenta un peso lineal en su dimensión de características de entrada (`Shard(-1)`) y replica el sesgo, que se suma después de la reducción. El forward de cada rank produce una salida `Partial()` que contiene su parte de la suma, y `transform_output_post_forward` redistribuye esa parte hacia `output_layouts`. Reducir a `Replicate()` emite un all-reduce, y reducir a `Shard(1)` emite un reduce-scatter.
`EmbeddingParallel`	Fragmenta un `nn.Embedding` con los ids de tokens replicados. `shard_dim=0` (`embedding_rowwise`, la opción habitual) fragmenta el vocabulario y realiza un all-reduce de la búsqueda parcial de cada rank para obtener la salida completa. `shard_dim=1` (`embedding_colwise`) fragmenta la dimensión oculta y realiza un all-gather de la porción de la salida de cada rank.
`PackedColwiseParallel`	Una variante de `ColwiseParallel` para pesos fusionados, por ejemplo, `up_proj` y `gate_proj` empaquetados en `gate_up_proj`. Utilice `split_factor` cuando un peso empaquete un número de capas distinto de dos.
`PackedRowwiseParallel`	La contraparte por filas, para pesos empaquetados a lo largo de la dimensión final. Replica parámetros 1D.
`SequenceParallel`	Replica los parámetros del módulo y divide su entrada en `sequence_dim` (por defecto, `1`). Se utiliza para normalizaciones que operan por token, como `LayerNorm` y `RMSNorm`.
`ReplicatedWithGradAllReduce`	Replica un parámetro, pero realiza un all-reduce de su gradiente. Es necesario para las normas que se encuentran entre una capa column-wise y una capa row-wise y normalizan a lo largo de un eje fragmentado, donde cada rank solo ve sus propias heads.
`AllReduceParallel`	Realiza una all-reduce de la salida de forward de `Partial()` de un módulo hacia `Replicate()`. Úsela como punto de sincronización para un módulo cuyo cómputo termina en una suma parcial.
`MlaKvAProjParallel`	Divide la salida de `kv_a_proj_with_mqa` de la atención MLA de estilo DeepSeek-V2 y realiza un all-reduce del gradiente de la mitad de RoPE, lo que evita `kv_b_proj` y de lo contrario mantendría un gradiente parcial. Requiere `qk_rope_head_dim` en la configuración del modelo.
`MoeExpertsParallel`	Expertos MoE con paralelismo de tensores. Realiza un all-reduce de la salida de los expertos en la propagación hacia adelante y añade los all-reduces de la propagación hacia atrás para los estados ocultos y los pesos de enrutamiento.
`MoeIdentityParallel`	Divide previamente la entrada de un experto cero o identidad por el tamaño de la malla, cancelando el all-reduce que `moe_tp_experts` aplica aguas abajo.
`LocalParamsParallel`	Ejecuta el método forward de un módulo con sus propios parámetros fragmentados como tensores locales. Úsalo en el propietario de un `nn.Parameter` simple, como los attention sinks, que su propia entrada de plan fragmenta.
`MoEParamShard`	Fragmenta los pesos de los expertos MoE en una colocación dada. Sustenta el nombre `grouped_gemm`, donde `shards_expert_dim=True` también reescribe `module.num_experts` al número de expertos por rank.
`EpRouterParallel`	Enmascara las puntuaciones del router para los expertos no locales y remapea los IDs globales de expertos a locales, de modo que cada rank solo ejecute los expertos que posee. Requiere que `num_experts` sea divisible por el tamaño de la malla.
`RouterParallelMegaMoe`	Variante de router para DeepGEMM Mega MoE, que despacha los expertos dentro del kernel, el cual requiere la salida del router sin modificar.
`MoeTensorParalellMegaMoeExperts`	Capa de expertos solo para inferencia para DeepGEMM Mega MoE. Omite las sincronizaciones de gradientes y pasa el grupo de procesos al módulo para que el kernel pueda configurar sus búferes compartidos en la primera pasada hacia adelante.

Estrategias empaquetadas

El empaquetado de pesos combina varias capas lineales en una única capa más grande. Las estrategias PackedColwiseParallel y PackedRowwiseParallel dividen correctamente los pesos empaquetados. Las estrategias básicas ColwiseParallel o RowwiseParallel dividen los pesos empaquetados de forma incorrecta.

El siguiente ejemplo agrupa up_proj y gate_proj en un único módulo gate_up_proj y requiere la estrategia packed_rowwise para particionar 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))

Use la multiplicación de matrices por lotes en el paso forward para calcular la salida del módulo 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

Un Shard simple divide esa dimensión en bloques contiguos, por lo que un rank recibiría el final de gate y el inicio de up en lugar de una porción de cada uno. Las estrategias empaquetadas utilizan _StridedShard con split_factor (2 por defecto) para alternar la división, de modo que cada rank reciba una porción correspondiente de ambas mitades, de forma que chunk siga alineado tras el sharding.

Estrategias de particionamiento personalizadas

Hereda de TensorParallelLayer en distributed/tensor_parallel.py para crear una estrategia de partición personalizada. Sobrescribe solo los hooks que tu estrategia necesite, ya que todos tienen un valor predeterminado sin efecto.

Hook	Propósito
`validate_param`	Rechazar un parámetro que esta estrategia no puede fragmentar, antes de cargar cualquier peso.
`shard_param`	Sustituya un parámetro por un marcador `DTensor` para que el cargador sepa qué fragmento corresponde a este rango.
`transform_inputs_pre_forward`	Redistribuya las entradas del módulo al formato que su método forward espera.
`context_around_forward`	Envuelva el método forward en un gestor de contexto, por ejemplo, para exponer tensores locales a un kernel.
`transform_output_post_forward`	Redistribuir o reducir la salida del módulo.
`should_use_local_tensors`	Informe de que el método forward de este módulo necesita tensores simples en lugar de `DTensor`s.
`install_forward`	Sustituya `module.forward` por completo. Solo sobrescriba este método cuando los hooks anteriores no sean suficientes, como hace `ReplicatedWithGradAllReduce` para registrar un hook de retropropagación.

El siguiente ejemplo recorre una versión reducida de ColwiseParallel.

Hereda de TensorParallelLayer y almacena las ubicaciones con las que trabaja la estrategia. La clase base no define __init__, por lo que no hay nada sobre lo que llamar a 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

Implementa shard_param para envolver un parámetro como un marcador de posición DTensor. Se ejecuta en tensores meta, por lo que distribute_tensor solo construye metadatos y no mueve datos. Pasa src_data_rank=None porque aún no hay un tensor completo del que dispersar.

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,
    )

Implementa las transformaciones de entrada y salida. install_forward las llama alrededor del forward original del módulo, por lo que solo deben mover los tensores entre disposiciones. La partición por columnas espera una entrada replicada y produce una salida fragmentada en la última dimensión.

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

El ColwiseParallel incluido añade rutas rápidas sobre esta base que omiten el viaje de ida y vuelta de DTensor para la inferencia de nn.Linear planos y para los módulos cuantizados que requieren tensores locales. Lea el código fuente antes de copiarlo si su estrategia necesita el mismo tratamiento.

Registra la estrategia para que tp_plan pueda nombrarla. El registro toma una instancia de la estrategia, por lo que las ubicaciones que pasas a __init__ se almacenan junto con el nombre. Observa el () a continuación.

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,
)

Benchmarks

La paralelización de tensores acelera significativamente la inferencia, especialmente para tamaños de lote grandes o secuencias largas.

Este gráfico muestra la aceleración esperada para una única pasada hacia adelante en Llama con una longitud de secuencia de 512.

Implementación del diseño

Transformers implementa el paralelismo de tensores de manera agnóstica al framework. Se basa en DeviceMesh y DTensor de torch.distributed para proporcionar una interfaz simple y extensible.

DeviceMesh

DeviceMesh crea una cuadrícula multidimensional de dispositivos que se comunican entre sí. Las distintas estrategias de paralelización requieren distintos patrones de comunicación. Crea un DeviceMesh con varias subcuadrículas para gestionar estos patrones.

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"])

La mayoría de las estrategias de paralelización de torch.distributed se aplican a la malla en sí o a sus submallas. La malla gestiona automáticamente los patrones de comunicación.

DTensor

DTensor (Distributed Tensor) gestiona la lógica distribuida sobre las operaciones habituales de tensores. Los pesos del modelo bajo paralelismo de tensores se almacenan como DTensors, lo que permite que una estrategia describa la comunicación como un cambio de distribución en lugar de una operación colectiva explícita.

El atributo placements indica a PyTorch cómo se distribuye un tensor entre los dispositivos de una DeviceMesh. Acepta los siguientes valores:

Shard(dimension) divide un DTensor a lo largo de una dimensión dada sobre el DeviceMesh con el que se construyó. La partición por columnas divide la dimensión de características de salida del peso y la única dimensión del sesgo.
```
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
```
La partición por filas fragmenta la dimensión de las características de entrada y replica el sesgo, ya que el sesgo se suma una sola vez después de la reducción, en lugar de en cada rank.
```
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() replica un DTensor en todo el DeviceMesh, creando una copia completa del tensor en cada dispositivo.
Partial() marca un tensor como pendiente de una reducción. La salida de la función forward de una capa por filas es Partial(), y redistribuirla a Replicate() es lo que emite el all-reduce.

Recursos

La sección Ultra-Scale Playbook sobre paralelismo de tensores ofrece más detalles.
Consulta la guía de expert parallelism si estás utilizando un modelo de mixture-of-experts (MoE). Estos modelos admiten tensor parallelism y expert parallelism.
Lee el artículo de blog Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand para obtener un resumen rápido del paralelismo de tensores y aprender cómo difieren las configuraciones de paralelismo por columnas y por filas.
Consulte la guía de entrenamiento de Tensor parallelism para aprender a utilizarlo en un entorno de entrenamiento.

Update on GitHub