# Paralelismo de tensores para inferencia

[Paralelismo de tensores](./perf_train_gpu_many#tensor-parallelism) 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`.

```py
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](/docs/transformers/main/es/expert_parallelism#transformers.DistributedConfig).

- Establezca `DistributedConfig(tp_size=N)` para utilizar el plan predefinido del modelo.
- Defina un `tp_plan` manual y páselo a [DistributedConfig](/docs/transformers/main/es/expert_parallelism#transformers.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()](/docs/transformers/main/es/main_classes/model#transformers.PreTrainedModel.from_pretrained) está deprecado y se eliminará en v5.18.

```py
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)
```

Lance el script de inferencia con [torchrun](https://pytorch.org/docs/stable/elastic/run.html). Use un proceso por GPU.

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

Defina un plan de paralelismo de tensores para cada capa en `tp_plan` y páselo a través de [DistributedConfig](/docs/transformers/main/es/expert_parallelism#transformers.DistributedConfig). El siguiente ejemplo utiliza particionamiento por columnas y por filas. Consulte la sección [Estrategias de particionamiento](#partitioning-strategies) para conocer otras estrategias admitidas.

La partición manual requiere un profundo conocimiento de la arquitectura del modelo y de las interacciones entre estrategias. Una elección inadecuada de partición genera modelos lentos que fallan o producen resultados incorrectos. El [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) explica las estrategias de partición en detalle.

Las claves son nombres de módulos o parámetros, donde `*` representa los índices de las capas. Un nombre de estrategia no reconocido genera un `ValueError` que enumera los nombres admitidos, y Transformers registra una advertencia para las reglas del plan que no coincidieron con nada y para los parámetros que ninguna regla cubrió.

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

## Estrategias de particionamiento

`ALL_PARALLEL_STYLES`, una `ParallelInterface` definida en [distributed/parallel_utils.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/parallel_utils.py), asocia cada nombre de estrategia que puedes usar en un plan con 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](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py) registra las estrategias de paralelismo de tensor en ella:

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

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](https://github.com/huggingface/transformers/blob/main/src/transformers/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 para convertirla 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 fragmenta 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 reducción total de la salida de la función forward del `Partial()` de un módulo hacia `Replicate()`. Úsela como punto de sincronización para un módulo cuyo cálculo 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`.

```python
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`.

```python
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](https://github.com/huggingface/transformers/blob/main/src/transformers/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`.

1. Hereda de `TensorParallelLayer` y guarda 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()`.

    ```python
    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. Implemente `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. Pase `src_data_rank=None` porque aún no hay un tensor completo desde el cual dispersar.

    ```python
    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. Implementa las transformaciones de entrada y salida. `install_forward` las invoca 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.

    ```python
    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.

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

    ```python
    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](./model_doc/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](https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html) y [DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html) de [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) 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.

```python
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 `DTensor`s, 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.

    ```python
    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.

    ```python
    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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) sobre paralelismo de tensores ofrece más detalles.

- Consulta la guía de [expert parallelism](./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](https://huggingface.co/blog/qgallouedec/tp) 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](./tensor_parallelism) para aprender a utilizarlo en un entorno de entrenamiento.