FSDP2

Fully Sharded Data Parallel (FSDP2) distribui o modelo, os gradientes e os estados do otimizador entre as GPUs. Antes do cálculo, cada GPU reúne um conjunto completo de parâmetros de todos os shards e, em seguida, libera a memória. Essa distribuição permite treinar modelos maiores do que a memória de uma única GPU, ao custo de mais comunicação do que no DDP. Use FSDP quando o seu modelo ou os estados do otimizador não couberem em uma única GPU.

                      ┌─────────────────┐
                      │  training data  │
                      └────────┬────────┘
            ┌──────────────────┼──────────────────┐
            │ shard 0          │ shard 1          │ shard 2
            ▼                  ▼                  ▼
     ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
     │  param      │    │  param      │    │  param      │
     │  shard 0    │    │  shard 1    │    │  shard 2    │
     │  GPU 0      │    │  GPU 1      │    │  GPU 2      │
     └──────┬──────┘    └──────┬──────┘    └──────┬──────┘
            │                  │                  │
            └──────── all-gather (params) ────────┘
                               │
                    full params on each GPU
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
         forward             forward             forward
            │                  │                  │
            └───── reduce-scatter (grads) ────────┘
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
     grad shard 0       grad shard 1       grad shard 2
     optim shard 0      optim shard 1      optim shard 2
        step               step               step

Estratégias de sharding

O FSDP2 controla o sharding com ~TrainingArguments.fsdp_config. Defina fsdp=True para habilitar o FSDP e defina reshard_after_forward na configuração do FSDP para escolher o compromisso entre memória e throughput.

auto_wrap_policy controla como os módulos são agrupados em unidades FSDP. O valor padrão é "TRANSFORMER_BASED_WRAP", que agrupa as camadas de transformer do modelo. Sem agrupamento ("NO_WRAP"), o modelo inteiro se torna uma única unidade FSDP e você perde o benefício de memória do sharding.

Configurar FSDP

Esses campos controlam como o FSDP2 envolve, fragmenta e carrega o modelo. reshard_after_forward e auto_wrap_policy são abordados em Estratégias de fragmentação.

• cpu_offload transfere parâmetros e gradientes para a CPU quando não estão em uso, para economizar memória de GPU.

• transformer_layer_cls_to_wrap define a camada de transformer a ser encapsulada em uma unidade FSDP quando auto_wrap_policy é "TRANSFORMER_BASED_WRAP". Cada unidade gerencia suas próprias operações de gather e scatter. Somente os parâmetros da unidade atual são agregados durante a passagem de forward. Os parâmetros das unidades anteriores são liberados para economizar memória.

  Envolver apenas o modelo de nível superior não gera economia de memória de GPU. Envolver cada camada Linear individual torna a comunicação entre unidades muito cara. Deixe este campo vazio e o FSDP lerá o valor da definição do modelo.

• min_num_params define o número mínimo de parâmetros por módulo para o encapsulamento baseado em tamanho. Ele é usado apenas quando auto_wrap_policy é "SIZE_BASED_WRAP".

• state_dict_type controla o formato do checkpoint. O padrão é "FULL_STATE_DICT" para um único checkpoint compatível com Transformers. Use "SHARDED_STATE_DICT" para um arquivo de checkpoint por rank, o que é mais rápido para modelos grandes. Checkpoints fragmentados só podem ser carregados de volta no FSDP, então salve um "FULL_STATE_DICT" para o checkpoint final que você deseja compartilhar ou carregar fora do FSDP.

• cpu_ram_efficient_loading carrega o checkpoint do disco apenas no rank 0. As outras GPUs inicializam um modelo vazio e recebem os pesos por broadcast, evitando que múltiplos processos carreguem um modelo grande na RAM da CPU.

• activation_checkpointing recalcula as ativações durante a passagem reversa em vez de armazená-las. Use isso em vez de gradient checkpointing em TrainingArguments. Definir ambos gera um erro.

Configure o treinamento com FSDP usando um arquivo de configuração do Accelerate ou um arquivo de configuração do FSDP passado para fsdp_config.

accelerate launch train.py

{
  "version": 2,
  "reshard_after_forward": true,
  "cpu_offload": false,
  "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  "state_dict_type": "FULL_STATE_DICT",
  "cpu_ram_efficient_loading": true,
  "activation_checkpointing": true
}

from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)

Para modelos de mistura de especialistas, fsdp_size pode ser combinado com tp_size e ep_size para dividir os especialistas em uma dimensão da malha e o restante na outra. Consulte paralelismo de especialistas.

Avaliar com paralelismo de contexto

O paralelismo de contexto (CP) com FSDP2 é configurado por meio do ParallelismConfig do Accelerate, especificamente pelo parâmetro cp_size. Na inicialização, passe tanto --use_parallelism_config quanto --parallelism_config_cp_size. Um parallelism_config não vazio no YAML do Accelerate também ativa use_parallelism_config. Um --parallelism_config_cp_size isolado é ignorado e mantém cp_size=1.

No FSDP2, o CP divide a sequência entre os ranks, de modo que cada GPU mantém uma fatia do contexto (allgather ou alltoall). Isso difere do Ulysses sequence parallelism, que é executado por meio do DeepSpeed ALST e usa sp_size na mesma configuração. Escolha entre CP ou Ulysses SP. Não defina cp_size e sp_size acima de 1 simultaneamente.

evaluate() usa o mesmo contexto de paralelismo de contexto que train() por meio de prediction_step(). Cada rank executa seu fragmento de sequência enquanto esse contexto está ativo. eval_loss é a perda correta (não multiplicada por cp_size). O early stopping e load_best_model_at_end usam essa métrica.

Os logits permanecem fragmentados ao longo da sequência, e nada os reorganiza em previsões de comprimento total. Previsões e compute_metrics sobre logits não são suportados sob CP.

Reutilize seu lançamento existente de FSDP2 + cp_size quando precisar de eval_loss.

Se você criar uma subclasse de prediction_step() e contar os tokens manualmente, conte num_items_in_batch antes de entrar no contexto de paralelismo de contexto. Enquanto esse contexto estiver ativo, os buffers de sequência são fragmentados no local entre os ranks. Contar depois subestima o lote. O Trainer padrão já conta primeiro, e você só precisa disso se sobrescrever prediction_step.

Próximas etapas

• Consulte DDP para treinamento paralelo de dados quando o seu modelo couber em uma única GPU.
• Consulte DeepSpeed para otimização ZeRO e offloading para NVMe.
• Para FSDP em TPUs com PyTorch/XLA, defina xla, xla_fsdp_settings e xla_fsdp_grad_ckpt em ~TrainingArguments.fsdp_config.
• Leia o capítulo sobre FSDP do The Ultra-Scale Playbook para obter mais informações sobre como o FSDP funciona.