Paralelismo de secuencias en 🤗 accelerate

Esta guía cubrirá los fundamentos del uso de la paralelización de secuencias en 🤗accelerate.

Véase también el muy relacionado Context Parallelism.

¿Por qué la paralelización de secuencias?

Con el auge de los modelos de lenguaje de gran escala y, más recientemente, de los modelos de razonamiento, la longitud de la secuencia ha crecido rápidamente. Esto, combinado con la complejidad de memoria cuadrática de la atención, ha generado la necesidad de formas más eficientes para entrenar modelos con secuencias largas. Con una longitud de secuencia de 128k, el requisito de memoria de la matriz de atención es 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads para precisión bf16, dado una implementación estándar de atención. Concedido que, con el uso de flash attention o SDPA, que no materializan estos pesos de atención, esto disminuye drásticamente, pero el crecimiento en los requisitos de memoria sigue siendo considerable.

El paralelismo de secuencia Ulysses nos permite dividir las entradas del cálculo de atención a lo largo de la dimensión de secuencia y calcular la atención de forma normal, pero utilizando solo un subconjunto de cabezas de atención en cada GPU. Con esto, podemos entrenar modelos con secuencias largas, con algunas herramientas adicionales, escalando a una longitud de secuencia de más de 15M. Para ver cómo aumentar Ulysses SP con TiledMLP, Liger-Kernel, el descargo de puntos de control de activación a la CPU y algunos otros trucos, consulte el artículo: Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences.

En qué se diferencia Ulysses SP de FSDP CP

En el documento Context Parallelism puedes aprender a desplegar otra tecnología llamada Context Parallelism, que también realiza particiones en la dimensión de la secuencia, pero utiliza Ring Attention en lugar de particionar en la dimensión de las cabeceras.

Los siguientes artículos ofrecen una explicación muy detallada de las diferencias entre las dos tecnologías:

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

Un resumen rápido adaptado de uno de los artículos:

• Ulysses SP tiene un sobrecosto de comunicación relativamente bajo, pero está limitado por el número de Attention Heads y, por lo tanto, tiene ciertos requisitos para la topología de red (el número de attention heads debe ser divisible por el número de gpus participantes para una sola réplica). La comunicación all-to-all es sensible a la latencia y requiere Deepspeed.
• La comunicación P2P en anillo de FSDP CP Ring-Attention no tiene los requisitos de divisibilidad mencionados anteriormente, pero tiene un mayor volumen de comunicación.

Finalmente, debería ser posible combinar SP + CP, tal como se explica en el artículo USP: A Unified Sequence Parallelism Approach for Long Context Generative AI, para soportar una longitud de secuencia aún mayor, aunque esto aún no está integrado en 🤗accelerate.

Backends de paralelismo de secuencias admitidos

Actualmente, el único backend de paralelismo de secuencias es deepspeed, que proviene de la versión modernizada de Ulysses SP, parte de la tecnología Arctic Long Sequence Training. También existe un tutorial por si deseas integrarlo directamente en tu propio código.

Cómo usar el paralelismo de secuencias?

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

Como con cualquier otra función de 🤗accelerate, también puedes habilitar el paralelismo de secuencias pasando las banderas correspondientes a accelerate launch. En este caso, no hay diferencia:

accelerate launch --parallelism-config-sp-size 8  ...

También puedes establecer sp_size y otras configuraciones en el comando accelerate config, que las guardará en tu archivo de configuración de accelerate, de modo que no tengas que pasarlas cada vez que ejecutes tu script.

El paralelismo de secuencias se combina con el paralelismo de datos. No requiere GPUs adicionales. Por lo tanto, si tienes 8 GPUs, puedes hacer: --parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8. O bien, puedes usar la clase ParallelismConfig para configurarlos de forma programática.

Importante: Debe asegurarse de que dp_replicate_size × dp_shard_size × sp_size = num_processes. Por ejemplo, con 8 GPUs y sp_size=8, necesita dp_shard_size=1 (ya que 1 × 1 × 8 = 8). Con 4 GPUs y sp_size=2, podría usar dp_shard_size=2 (ya que 1 × 2 × 2 = 4) para un paralelismo 2D.

Configuración del backend SP de ALST/Ulysses

ALST/UlyssesSP implementa el paralelismo de secuencia mediante el paralelismo de cabezas de atención, tal como se explica en este artículo. Por simplicidad, reutilizamos el concepto y la configuración del paralelismo de secuencia, que, desde la perspectiva del usuario, es el mismo: se utilizan varias GPU para procesar un solo lote.

Para dar una idea de lo que ALST hizo posible: nos permitió entrenar en bf16 con 500K tokens en una sola GPU H100, 3.7M en un solo nodo y 15M en Llama-8B usando solo cuatro nodos. Esta función de HF Accelerate habilita solo 1 de los 3 componentes de ALST, por lo que la longitud de secuencia alcanzable será menor. Para obtener todo el poder de ALST, se necesitaría habilitar TiledMLP, Activation checkpoint offload to CPU y algunas otras cosas. Para más detalles, consulte este tutorial.

Para configurar el backend deepspeed:

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend: establécelo en deepspeed aquí
• sp_size es el grado de paralelismo de secuencia: en el ejemplo anterior es 4, por lo que se usarán 4 GPUs para procesar un solo lote (mientras se realiza DP=4 sobre las mismas GPUs).
• sp_seq_length y sp_seq_length_is_variable se utilizan para gestionar las longitudes de secuencia. Si sp_seq_length_is_variable=True, el backend trabajará con una longitud de secuencia que puede cambiar entre lotes, en cuyo caso el valor de sp_seq_length puede establecerse en cualquier valor divisible por el grado de paralelismo de secuencia o no establecerse en absoluto. En este caso, en cada forward, las variables de secuencia se derivarán de la entrada. Si es False, entonces seq_length debe coincidir con la dimensión de longitud de secuencia del lote, que luego deberá rellenarse para que sea siempre la misma. El valor predeterminado es True.
• sp_attn_implementation es uno de sdpa, flash_attention_2 o flash_attention_3. Esta implementación de paralelismo de secuencias utiliza position_ids en lugar de attention_mask; por lo tanto, eager no puede funcionar aquí hasta que admita el trabajo con position_ids. Además, tenga en cuenta que sdpa no maneja correctamente múltiples muestras combinadas en una; atenderá a la muestra completa como una sola. Si las muestras no están combinadas, sdpa funcionará correctamente. Por lo tanto, Flash Attention debería ser la opción ideal, ya que siempre funciona.

En lugar de establecer estos valores en el objeto DeepSpeedSequenceParallelConfig, también puede usar variables de entorno para lograr lo mismo; aquí están, en correspondencia con el final de la lista anterior.

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

Si no se pasa en el código, sp_size se puede establecer mediante el argumento de línea de comandos --parallelism_config_sp_size. Lo mismo aplica para otros argumentos. También se puede hacer una configuración de estilo archivo de configuración de accelerate, por ejemplo, para 2 GPUs:

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

Como se mencionó anteriormente, la paralelización de secuencias de Ulysses normalmente se superpone con la paralelización de datos: se utilizan los mismos rangos para alimentar flujos de datos únicos y también para realizar la paralelización de secuencias de Ulysses. Pero también se pueden crear réplicas de la siguiente manera:

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

Aquí utilizamos 4 gpus, con 2 réplicas de paralelismo de secuencia. Deepspeed-ZeRO es lo que impulsa el paralelismo de datos aquí.

Tenga en cuenta que gran parte de la magia está oculta dentro de UlyssesSPDataLoaderAdapter. Se utiliza en segundo plano, envolviendo su objeto DataLoader original, pero debe ser consciente de su existencia por si se encuentra con algún problema. También inyecta automáticamente el shift_labels correcto en el diccionario del batch, antes de que el batch se fragmente entre los ranks participantes.

Ahora, la única pieza que falta para empezar a usar ALST/UlyssesSP es agregar la pérdida entre los rangos mediante un all_gather diferenciable para obtener los gradientes correctos. El siguiente código lo hace, excluyendo además cualquier token enmascarado con -100, para obtener el promedio correcto:

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

Tenga en cuenta que los modelos manejan automáticamente shift_labels cuando está presente en el lote. La pasada hacia adelante del modelo recibe shift_labels a través de **kwargs y lo pasa a la función de pérdida, que calcula correctamente la pérdida para el paralelismo de secuencias. Si utiliza Liger Kernel, también maneja shift_labels de manera transparente y calcula la pérdida de una forma muy eficiente en memoria. Liger se recomienda encarecidamente para longitudes de secuencia largas, ya que libera memoria de GPU al utilizar operaciones fusionadas (por ejemplo, el cálculo fusionado de logit-pérdida que nunca materializa el tensor completo de logits en memoria).

Si quieres ver lo que HF Accelerate hizo a nivel interno, lee este tutorial de integración completo.

Para un ejemplo de un bucle de entrenamiento de Accelerate con ALST/UlyssesSP habilitado, consulte examples/alst_ulysses_sequence_parallelism.

[!Advertencia]

Esta API es bastante nueva y aún se encuentra en su etapa experimental. Aunque nos esforzamos por proporcionar una API estable, algunas partes pequeñas de la API pública pueden cambiar en el futuro.

Dado que se trata de un backend de DeepSpeed, se aplica la configuración habitual de DeepSpeed, por lo que también puedes combinar el paralelismo de secuencias con el descargado de estados del optimizador y/o pesos para liberar más memoria de GPU y habilitar una longitud de secuencia aún mayor. Esta tecnología se ha probado y funciona con DeepSpeed ZeRO etapa 2 y 3.