FSDP vs DeepSpeed

O Accelerate oferece flexibilidade em frameworks de treinamento, ao integrar duas ferramentas extremamente poderosas para treinamento distribuído, a saber, Pytorch FSDP e Microsoft DeepSpeed. O objetivo deste tutorial é traçar paralelos, bem como destacar possíveis diferenças, para capacitar o usuário a alternar sem problemas entre esses dois frameworks.

Para alternar entre os frameworks, recomendamos executar o código accelerate launch passando o arquivo de configuração correto com --config_file, ou passando diretamente os argumentos respectivos para FSDP e DeepSpeed .

Exemplos de configurações do Accelerate podem ser encontrados aqui para DeepSpeed e FSDP, ou no example zoo sob “Launch Configurations”

Este tutorial é apenas para cenários de nó único, multi-GPU.

Configurando Funcionalidades

Os tensores do modelo são divididos em diferentes GPUs em uma tentativa de escalar os tamanhos dos modelos; isso é chamado de sharding no FSDP e de partitioning no DeepSpeed. O sharding do FSDP e os estágios do DeepSpeed ZeRO (partitioning) são configurados por --fsdp_sharding_strategy e --zero_stage, respectivamente. Em particular, o FULL_SHARD do FSDP corresponde ao estágio 3 do DeepSpeed ZeRO; veja este mapeamento abrangente entre as configurações de sharding do FSDP e do DeepSpeed ZeRO. A tabela abaixo resume e agrupa configurações semelhantes:

Para descrições detalhadas do acima, consulte a documentação de lançamento do Accelerate.

Para acessar outras configurações do DeepSpeed, como as configurações de precisão mista, você precisa passar um --deepspeed_config_file, consulte a documentação.

O DeepSpeed também pode ser configurado via DeepSpeedPlugin, por exemplo, DeepSpeedPlugin.zero_stage é equivalente a --zero_stage, e DeepSpeedPlugin.hf_ds_config pode ser usado para passar --deepeed_config_file.

O FSDP também pode ser configurado via FullyShardedDataParallelPlugin, por exemplo, FullyShardedDataParallelPlugin.sharding_strategy é equivalente a --fsdp_sharding_strategy.

Checkpointing

Observe que, embora o FSDP possa ser configurado via --fsdp_state_dict_type para salvar checkpoints completos / fragmentados.

Para o DeepSpeed Zero3, é possível passar --zero3_save_16bit_model true, o que consolida convenientemente o modelo em um único rank e salva; isso é o equivalente no FSDP de fsdp_state_dict_type: FULL_STATE_DICT.

Para modelos grandes, consolidar o modelo em um único rank pode ser muito lento.

Para checkpoints mais rápidos, para FSDP use fsdp_state_dict_type: SHARDED_STATE_DICT, e para DeepSpeed Zero3 use o script zero_to_fp32.py para converter checkpoints fragmentados posteriormente.

Offloading

O FSDP permite apenas offload all-or-nothing (ou seja, ou offload parâmetros, gradientes e otimizador, ou mantém todos na GPU), mas o DeepSpeed pode offload parâmetros e otimizador de forma diferente. Além disso, o DeepSpeed também suporta offloading para NVME.

Pré-carregamento

O FSDP permite duas configurações de prefetching, --fsdp_forward_prefetch e --fsdp_backward_prefetch, para melhorar a sobreposição de comunicações e computação, a um custo de memória extra, consulte a documentação do FSDP. Para o DeepSpeed, o prefetching será ativado quando necessário, e sua ativação depende de certos hiperparâmetros, como stage3_param_persistence_threshold, stage3_max_reuse_distance, etc., que podem ser configurados para o Zero3; o accelerate pode definir esses hiperparâmetros automaticamente se você não os definir explicitamente no arquivo de configuração do deepspeed.

Para o FSDP, defina fsdp_backward_prefetch: BACKWARD_PRE para melhorar o throughput, caso a memória permita.

Carregamento de Modelos

Embora o FSDP exija um --fsdp_cpu_ram_efficient_loading true explícito para ativar o carregamento eficiente do modelo, o transformers ativará o recurso semelhante sempre que o DeepSpeed Zero3 for usado.

Para o FSDP, sempre que --fsdp_cpu_ram_efficient_loading true for definido, o accelerate definirá automaticamente sync_module_states como true. Para o carregamento eficiente de RAM, os pesos serão carregados apenas em um único rank, e, portanto, é necessário que sync_module_states seja usado para transmitir os pesos para os outros ranks.

Model

O FSDP requer um --fsdp_auto_wrap_policy explícito para que o algoritmo decida como agendar as operações de all-gather e reduce-scatter. Mas, para o DeepSpeed, isso é transparente para o usuário.

Para o FSDP, basta definir fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP. Nas versões mais recentes do transformers, fazemos o possível para determinar o fsdp_transformer_layer_cls_to_wrap adequado para os modelos de transformers do HF. No entanto, se você receber um erro relacionado a isso, especifique-o.

Parâmetros de Invocação

O FSDP requer a flag explícita --fsdp_use_orig_params ao usar torch.compile, consulte a documentação do PyTorch. Para o DeepSpeed, isso é transparente para o usuário.

Para o FSDP, ao usar torch.compile, defina fsdp_use_orig_params: True.

Treinamento

Deepspeed requer as flags explícitas --gradient_accumulation_steps e --gradient_clipping. Para o FSDP, isso é transparente para o usuário.

Ao usar o DeepSpeed, defina gradient_accumulation_steps: "auto" e gradient_clipping: "auto" para recuperar automaticamente os valores definidos no Accelerator ou no TrainingArguments (se estiver usando o transformers).

Sobre as diferenças no tratamento da precisão dos dados

Para discutir como a precisão dos dados é tratada no FSDP e no DeepSpeed, é instrutivo primeiro apresentar uma visão geral de como os parâmetros do modelo são processados nesses frameworks. Antes que os parâmetros do modelo/otimizador sejam distribuídos entre as GPUs, a preparação dos parâmetros envolve primeiro “achatar” (flatten) os parâmetros em torch.Tensor unidimensionais. A implementação do FSDP/DeepSpeed difere em relação ao dtype no qual esses parâmetros “achatados” são armazenados, e isso tem implicações para a forma como o torch.Optimizer aloca seus dtypes. A tabela abaixo descreve os processos de ambos os frameworks; a coluna “Local” indica o processo que ocorre em nível de GPU individual, portanto, qualquer sobrecarga de memória causada pelo upcasting deve ser entendida como amortizada pelo número de GPUs utilizadas.

Como regra geral, para um treinamento estável com precisão mista automática, todos os parâmetros treináveis devem estar em torch.float32.

Portanto, ao usar o DeepSpeed com um pequeno número de GPUs, esteja ciente de possíveis sobrecargas de memória significativas devido ao upcasting durante a preparação.

Com o FSDP, na ausência de precisão mista, é possível operar o torch.Optimizer em baixa precisão torch_dtype, o que pode ser útil ao usar um pequeno número de GPUs.

Com precisão mista, FSDP e DeepSpeed farão o upcast na etapa de preparação do modelo (cf. tabela acima). Mas observe que o FSDP salvará os checkpoints na precisão upcasted; o DeepSpeed ainda pode salvar checkpoints de baixa precisão se --zero3_save_16bit_model for especificado.

Para esclarecer a tabela acima, considere os exemplos concretos abaixo; o passo prévio e o passo real do otimizador foram combinados para brevidade. Com FSDP, é possível operar nos dois modos mostrados abaixo, mas o DeepSpeed só pode operar em um.