Context Parallel no 🤗 accelerate

Este guia abordará os fundamentos do uso de paralelismo de contexto no 🤗accelerate. Para os leitores mais curiosos, também abordaremos alguns aspectos técnicos nas seções seguintes.

Veja também o muito relacionado Guia de Paralelismo de Sequência.

Por que paralelismo de contexto?

Com o advento dos grandes modelos de linguagem e, recentemente, dos modelos de raciocínio, o comprimento da sequência tem crescido rapidamente. Isso, combinado com a complexidade de memória quadrática da atenção, gerou a necessidade de formas mais eficientes de treinar modelos com sequências longas. Com um comprimento de sequência de 128k, o requisito de memória da matriz de atenção é 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads para precisão bf16, dado uma implementação vanilla de atenção. Concedido que, com o uso de flash attention ou SDPA, que não materializam esses pesos de atenção, isso diminui drasticamente, mas o crescimento nos requisitos de memória ainda é considerável.

O paralelismo de contexto permite dividir as entradas do cálculo de atenção ao longo da dimensão da sequência e calcular a atenção em paralelo em várias GPUs. Com isso, é possível treinar modelos com sequências longas, escalando potencialmente para um comprimento de sequência de 1M+.

Como usar paralelismo de contexto?

from accelerate.utils import ParallelismConfig, TorchContextParallelConfig

+ cp_config = TorchContextParallelConfig(
+       cp_comm_strategy="alltoall", # no need to use cp_config at all, if you want to use the default "allgather"
+ )

+ parallelism_config = ParallelismConfig(
+     cp_size=8,
+     cp_handler=cp_config,  # or just cp_size=8, if you want to use the default "allgather"
+ )

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

Como em qualquer outro recurso do 🤗accelerate, você também pode habilitar o paralelismo de contexto passando as flags correspondentes para accelerate launch. Neste caso, não há diferença:

accelerate launch --parallelism-config-cp-size 8 --parallelism-config-cp-comm-strategy [allgather|alltoall] ...

Você também pode definir cp_size e cp_comm_strategy no comando accelerate config, o que os salvará no seu arquivo de configuração do accelerate, para que você não precise passá-los toda vez que executar seu script.

O paralelismo de contexto é compatível com outras estratégias de paralelismo, como paralelismo de dados, paralelismo de tensor e FSDP2. Você pode simplesmente combiná-las definindo os tamanhos de paralelismo para os valores desejados, por exemplo, --parallelism-config-dp-size 8 --parallelism-config-tp-size 2 --parallelism-config-cp-size 8. Ou você pode usar a classe ParallelismConfig para defini-las programaticamente.

O paralelismo de contexto está intimamente acoplado ao FSDP2, sobre o qual você pode saber mais na introdução ao FSDP2. Isso significa que o paralelismo de contexto só funciona se você usar FullyShardedDataParallelPlugin ou --use-fsdp com a versão definida como 2 no seu programa. Se nenhum FSDP2 for usado, um erro será lançado.

O paralelismo de contexto funciona apenas com SDPA e apenas sem máscara ou com máscara causal. Não podemos detectar isso corretamente para você, portanto, é sua responsabilidade garantir que está usando SDPA sem máscara ou com máscara causal. Se você usar qualquer outra implementação de atenção, um erro será lançado.

Após habilitar o paralelismo de contexto com os métodos mencionados acima, você pode aplicá-lo ao seu loop de treinamento. Fornecemos um wrapper fino ao redor de torch.distributed.tensor.experimental.context_parallel que você pode usar no seu loop de treinamento, o que abstrai parte da complexidade de usá-lo (mais detalhes sobre isso adiante). Para minimizar as alterações que você precisa fazer no seu loop de treinamento, fornecemos um context manager que é um noop se o paralelismo de contexto não estiver habilitado e aplica o paralelismo de contexto se estiver habilitado. Dessa forma, você pode usá-lo no seu loop de treinamento sem alterar nenhum código com base na sua configuração de paralelismo. Você pode usá-lo da seguinte forma:

for batch in dataloader:
    with accelerator.maybe_context_parallel(
        buffers=[batch["input_ids"], batch["attention_mask"]],
        buffer_seq_dims=[1, 1],
        no_restore_buffers={batch["input_ids"], batch["labels"]},
    ):
        outputs = model(**batch)
        ...

Este gerenciador de contexto precisa ser recriado a cada etapa de treinamento, conforme mostrado no exemplo acima. É crucial fazer isso.

Isso pode escalar o tamanho do contexto para um comprimento de sequência potencialmente superior a 1M. Abaixo, apresentamos a velocidade e o uso de memória do paralelismo de contexto para um tamanho de contexto de até 256k. Podemos observar que, ao dobrarmos o tamanho do contexto e o número de GPUs, conseguimos manter um uso de memória consistente, o que potencialmente permite a escalação de comprimento de contexto ilimitado.

Figura 1: Uso de memória e velocidade do paralelismo de contexto para tamanhos de contexto de até 256k.

Esses exemplos foram criados com um script que você pode encontrar na pasta de exemplos. Para executar o exemplo em 8 GPUs H100 (comprimento de sequência de 128k), você pode usar o seguinte comando:

Copiado

accelerate launch --use-fsdp --fsdp-activation-checkpointing=TRUE examples/torch_native_parallelism/nd_parallel.py --cp-size=8 --sequence-length=128000

Interface do Accelerate

O gerenciador de contexto aceita alguns argumentos, que são usados para configurar o paralelismo de contexto.

• buffers: Esta é uma lista de tensores que devem ser fragmentados ao longo da dimensão da sequência. Esses tensores geralmente são ids de entrada, rótulos e máscara de atenção.
• buffer_seq_dims: Esta é uma lista de inteiros que especifica a dimensão de sequência dos buffers, na ordem da lista buffers. Se você passar buffers=[input_ids, shift_labels] com ambos tendo a forma [batch_size, sequence_length], você passaria buffer_seq_dims=[1, 1], pois a dimensão de sequência é a segunda dimensão dos tensores. Isso é necessário para o cálculo correto das saídas do modelo.
• no_restore_buffers: A implementação de paralelismo de contexto modifica os buffers in-place, convertendo-os em torch.distributed.tensor.Dtensors. Após a saída do gerenciador de contexto, um kernel de comunicação precisaria ser lançado para restaurar os buffers ao seu estado original (geralmente all-gather). Isso leva algum tempo, portanto, é recomendado passar os mesmos tensors do argumento buffers, para evitar comunicação desnecessária, a menos que você tenha certeza de que precisa usar os buffers após a saída do gerenciador de contexto.

O paralelismo de contexto não é compatível com labels que são uma cópia de input_ids, pois os modelos do 🤗 transformers podem deslocá-los para habilitar o modelamento de linguagem causal por conta própria. Imagine o seguinte caso: labels = [l1, l2, l3, l4, … li] se aplicarmos o paralelismo de contexto, cada rank acabará com uma parte dos labels, como esta: labels_rank_0 = [l1, l2], labels_rank_1 = [l3, l4], … após o código de modelamento do transformers deslocar os labels, o resultado será: labels_rank_0 = [l2, PAD], labels_rank_1 = [l3, PAD], … onde PAD é um token de preenchimento. Isso resultará em um cálculo incorreto da perda, pois os labels não estão mais alinhados com as entradas. Por causa disso, você precisa deslocar manualmente os labels antes de passá-los para o modelo

Opções configuráveis

O Accelerate oferece apenas uma opção para configurar o paralelismo de contexto (exceto por cp_size)

• cp_comm_strategy: O método de rotação a ser usado para os shards. Recomendamos fortemente manter como "allgather", pois é muito provável que ele tenha desempenho superior ao "alltoall" na maioria dos casos.

O tamanho do paralelismo de contexto é bastante autoexplicativo: é o número de ranks entre os quais as entradas devem ser divididas em shards. A rotação de shards do paralelismo de contexto define como os shards das entradas são rotacionados entre os ranks. Abordaremos as 2 opções em mais detalhes na próxima seção.

Você pode ver um exemplo de ponta a ponta no arquivo exemplo de paralelismo ND, onde é possível treinar um modelo de 8B com comprimento de contexto de até 128k em um único nó 8xH100. Usando treinamento multi-nó, você pode escalar isso para comprimento de sequência de 1M+ em múltiplas GPUs. Você também pode combiná-lo sem problemas com outras estratégias de paralelismo para atender às suas necessidades.

Detalhes técnicos

Esta seção é bastante técnica, então, se você não precisa aprender os detalhes internos do paralelismo de contexto, pode pular e começar a construir 🚀

Vamos usar a palavra shard extensivamente nas seções a seguir, então vamos defini-la primeiro. Se chamamos um tensor de sharded ao longo da dimensão Dth, ao longo de N ranks, queremos dizer que esse tensor é dividido em N partes, onde cada parte do tensor tem a forma [..., D//N, ...].

Então, como funciona?

O paralelismo de contexto funciona ao dividir as matrizes Q, K and V ao longo da dimensão da sequência. Cada rank possui seu shard designado de Q, que chamaremos de Q_i. Essa matriz permanece apenas neste rank durante toda a computação. Da mesma forma, cada rank possui seu próprio shard de K e V, que chamaremos de K_i e V_i. Em seguida, cada rank calcula a atenção com seu próprio shard de Q_i, K_i e V_i, que chamaremos de attn_i. Durante essa computação, um kernel de comunicação é iniciado para reunir os Ks e Vs de todos os outros ranks. A primitiva de comunicação utilizada depende da opção context_parallel_shard_rotation. Dessa forma, cada rank calcula a atenção local, primeiro com Q_i, K_i e V_i, e depois com K_j e V_j de todos os outros ranks. Como cada rank mantém as matrizes Q, K and V divididas ao longo da dimensão da sequência, as matrizes resultantes são menores e podem caber em uma única GPU.

Podemos formalizar isso no seguinte pseudocódigo:

comm_kernel = {"allgather": allgather, "alltoall": alltoall}[context_parallel_shard_rotation]
Qi, Ki, Vi = shard(Q, K, V, seq_dim)
attn[i] = attn(Qi, Ki, Vi)
for j in range(context_parallel_size):
    Kj, Vj = comm_kernel()
    attn[j] = attn(Qi, Kj, Vj) # [batch, num_heads, seq_len // context_parallel_size, head_dim]

final_attn = combine(attn)

all-to-all vs all-gather

all-gather

Então, qual é a diferença entre all-to-all e all-gather? No all-gather, a comunicação é muito simples. Depois (bem, antes, pois geralmente leva mais tempo) de calcularmos a atenção local attn_i, iniciamos um all-gather para reunir todos os outros Ks e Vs de todos os outros ranks. Como essa comunicação é feita, cada rank possui todos os Ks e Vs de todos os outros ranks e pode calcular a atenção com eles sequencialmente. No cenário ideal, o all-gather termina exatamente no momento em que o cálculo de attn_i é concluído. No entanto, isso nunca acontece na prática, então o sobreposição real ideal é alcançada quando o attn_i completo é sobreposto a uma parte da comunicação, e, para iniciar o cálculo com K_j e V_j, aguardamos o all-gather terminar.

all-to-all

All-to-all, ou às vezes chamado de ring-rotation, utiliza um padrão de comunicação em anel. Após concluir o cálculo de attn_i, um all-to-all é iniciado para enviar K_i e V_i para os ranks vizinhos. Em seguida, repetimos isso context_parallel_size-1 vezes, para que cada rank veja todas as shards de K e V de todos os outros ranks uma vez. No cenário ideal, fazemos prefetch das shards K_i+1 e V_i+1 do rank vizinho e essa comunicação é exatamente sobreposta ao cálculo do nosso attn_i atual. Novamente, na realidade, essa sobreposição perfeita nunca acontece. Dada a natureza dessa abordagem, se não alcançarmos uma sobreposição perfeita, a penalidade é muito maior do que com all-gather.

Como escolher o método de rotação adequado?

Em teoria, all-to-all deveria ser a melhor escolha. No entanto, na prática, raramente é. Por isso, usamos all-gather por padrão, pois é mais provável que ele alcance um desempenho melhor. Extensos benchmarks da equipe do torchtitan também mostram que all-to-all raramente supera all-gather. Ainda assim, oferecemos ambas as opções, pois você pode encontrar uma delas melhor para o seu caso de uso.

Você pode ver diretamente esse problema na saída do profiler na imagem abaixo:

Figura 1: Em vermelho, você pode ver o tempo ocioso, enquanto aguardamos a finalização do kernel all-to-all. Destacado na primeira barra azul, você pode ver que leva ~250us para terminar, o que é repetido N-1 vezes para cada chamada de atenção, onde N é o tamanho do paralelismo de contexto.

Por que apenas FSDP2?

Suportamos paralelismo de contexto apenas com FSDP2, pois criamos uma malha conjunta de context_parallel_size e dp_shard_size para aproveitar todo o seu potencial. O funcionamento é o seguinte: fazemos o sharding do modelo na malha conjunta de tamanho cp_size*dp_shard_size, o que maximiza a economia de memória. Isso é uma espécie de “lunch grátis”, pois a comunicação do FSDP é totalmente sobreposta ao cálculo de atenção, conforme mostrado nas imagens abaixo.

Figura 2: Nos retângulos azuis (Stream 23), é possível ver que o pré-busca do shard de `FSDP` é totalmente sobreposto ao cálculo de atenção (Stream 7), enquanto nos retângulos vermelhos (Stream 24), é possível ver que o kernel de all-gather resulta em uma bolha de tempo ocioso, na qual nosso stream de computação (7) fica ocioso.

Na figura acima, você também pode notar a diferença entre all-to-all e all-gather. Enquanto no all-to-all (Figura 1), lançamos um kernel de comunicação N-1 vezes para cada chamada de atenção, no all-gather (Figura 2), lançamos um kernel de comunicação apenas uma vez. Isso resulta em um bubble maior, mas ele ocorre apenas uma vez por chamada de atenção, enquanto no all-to-all, ele ocorre N-1 vezes.

Despacho de dados em malha conjunta

Garantimos que o mesmo lote de dados seja enviado para todo o subgrupo cp, para que os resultados estejam corretos. (Isso significa que cada rank no subgrupo cp recebe o mesmo lote de dados.) No entanto, também enviamos lotes diferentes para cada rank do grupo dp_shard. Imagine da seguinte forma:

# 8 GPUS, --dp_shard_size 4, --cp_size 2
# mesh = [[0, 1], [2, 3], [4, 5], [6, 7]]
# model is sharded across the whole mesh (each GPU holds 1/8 of the model)
# GPUs 0,1 = batch 0
# GPUs 2,3 = batch 1
... and so on.