Parâmetros

• context_parallel_config (ContextParallelConfig, opcional) — Configuração para paralelismo de contexto.
• tensor_parallel_config (TensorParallelConfig, opcional) — Configuração para paralelismo de tensor.

Parâmetros

• ring_degree (int, opcional, padrão 1) — Número de dispositivos a serem usados para Ring Attention. A sequência é dividida entre os dispositivos. Cada dispositivo calcula a atenção entre seus blocos locais de Q e KV, que são passados sequencialmente ao redor do anel. Menor uso de memória (mantém apenas 1/N de KV por vez), sobrepõe o cálculo à comunicação, mas requer N iterações para ver todos os tokens. Ideal para sequências longas com memória/largura de banda limitadas. Número de dispositivos a serem usados para ring attention dentro de uma região de paralelismo de contexto. Deve ser um divisor do número total de dispositivos na malha de paralelismo de contexto.
• ulysses_degree (int, opcional, padrão 1) — Número de dispositivos a serem usados para Ulysses Attention. A sequência é dividida entre os dispositivos. Cada dispositivo calcula o QKV local e, em seguida, faz o all-gather de todos os blocos KV para calcular a atenção completa em uma única passagem. Maior uso de memória (armazena todos os KV), requer comunicação all-to-all de alta largura de banda, mas com menor latência. Ideal para sequências moderadas com boa largura de banda de interconexão.
• convert_to_fp32 (bool, opcional, padrão True) — Se deve converter a saída e o LSE para float32 para estabilidade numérica do ring attention.
• rotate_method (str, opcional, padrão "allgather") — Método a ser usado para rotacionar os estados de chave/valor entre dispositivos no ring attention. Atualmente, apenas "allgather" é suportado.
• ulysses_anything (bool, opcional, padrão False) — Se deve ativar o modo “Ulysses Anything”, que suporta comprimentos de sequência e números de cabeças arbitrários que não são divisíveis uniformemente por ulysses_degree. Quando ativado, ulysses_degree deve ser maior que 1 e ring_degree deve ser 1.
• ring_anything (bool, opcional, padrão False) — Se deve ativar o modo “Ring Anything”, que suporta comprimentos de sequência arbitrários. Quando ativado, ring_degree deve ser maior que 1 e ulysses_degree deve ser 1.
• mesh (torch.distributed.device_mesh.DeviceMesh, opcional) — Um device mesh personalizado para usar em paralelismo de contexto. Se fornecido, este mesh será usado em vez de criar um novo. Isso é útil ao combinar paralelismo de contexto com outras estratégias de paralelismo (p. ex., FSDP, paralelismo de tensor) que compartilham o mesmo device mesh. O mesh deve ter as dimensões “ring” e “ulysses”. Use tamanho 1 para dimensões que não estão sendo usadas (p. ex., mesh_shape=(2, 1, 4) com mesh_dim_names=("ring", "ulysses", "fsdp") para apenas ring attention com FSDP).

Parâmetros

• tp_degree (int, padrão 1) — Número de dispositivos para dividir. Deve ser um divisor do número de cabeças de atenção (e dimensões ocultas do FFN) do modelo que está sendo paralelizado.
• mesh (torch.distributed.device_mesh.DeviceMesh, opcional) — Um device mesh personalizado para usar. Se fornecido, tp_degree é inferido a partir de mesh.size() e o argumento é ignorado. Útil ao combinar TP com outras estratégias de paralelismo (por exemplo, CP) que compartilham o mesmo mesh.

Parâmetros

• model (torch.nn.Module) — O modelo a ser fragmentado no local.
• config (TensorParallelConfig) — The tensor-parallel config. Its device mesh must already be set up with config.setup(...).
• tp_plan (dict) — Um mapeamento plano de globs de nomes de módulos para um estilo "colwise"/"rowwise" (ou uma variante empacotada), geralmente o _tp_plan do modelo.
• weights_already_sharded (bool, padrão False) — Indica se os parâmetros planejados já são shards de DTensor, como ocorre após um carregamento streaming de from_pretrained. Se True, apenas os forward hooks são registrados. Isso é passado explicitamente em vez de ser detectado, porque um parâmetro planejado ausente do checkpoint ainda seria um meta tensor e faria a detecção indicar “não shardado” para um modelo que, na verdade, está parcialmente shardado.