DeepSpeed

DeepSpeed implementa tudo o que é descrito no artigo ZeRO. Algumas das otimizações mais notáveis são:

1. Particionamento do estado do otimizador (ZeRO estágio 1)
2. Particionamento de gradientes (ZeRO estágio 2)
3. Particionamento de parâmetros (ZeRO estágio 3)
4. Tratamento personalizado de treinamento com precisão mista
5. Uma variedade de otimizadores rápidos baseados em extensões CUDA
6. ZeRO-Offload para CPU e Disco/NVMe
7. Particionamento hierárquico dos parâmetros do modelo (ZeRO++)

ZeRO-Offload tem seu próprio artigo dedicado: ZeRO-Offload: Democratizing Billion-Scale Model Training. E o suporte a NVMe é descrito no artigo ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

O DeepSpeed ZeRO-2 é usado principalmente apenas para treinamento, pois seus recursos não são úteis para inferência.

O DeepSpeed ZeRO-3 também pode ser usado para inferência, pois permite carregar modelos enormes em várias GPUs, o que não seria possível em uma única GPU.

O Accelerate integra o DeepSpeed por meio de 2 opções:

1. Integração dos recursos do DeepSpeed por meio da especificação de um deepspeed config file no accelerate config. Basta fornecer seu arquivo de configuração personalizado ou usar nosso modelo. A maior parte deste documento é focada nesse recurso. Isso suporta todos os recursos principais do DeepSpeed e oferece ao usuário muita flexibilidade. O usuário pode precisar alterar algumas linhas de código, dependendo da configuração.
2. Integração via deepspeed_plugin. Isso suporta um subconjunto dos recursos do DeepSpeed e usa opções padrão para o restante das configurações. O usuário não precisa alterar nenhum código e é adequado para aqueles que estão satisfeitos com a maioria das configurações padrão do DeepSpeed.

O que está integrado?

Treinamento:

1. O Accelerate integra todos os recursos do DeepSpeed ZeRO. Isso inclui todos os estágios do ZeRO 1, 2 e 3, bem como ZeRO-Offload, ZeRO-Infinity (que pode descarregar para disco/NVMe) e ZeRO++. Abaixo está uma breve descrição do Data Parallelism usando ZeRO - Zero Redundancy Optimizer, juntamente com o diagrama deste post de blog

(Fonte: link)

a. Etapa 1 : Divide os estados do otimizador entre os workers/GPUs paralelos de dados

b. Etapa 2: Divide os estados do otimizador + gradientes entre os workers/GPUs paralelos de dados

c. Etapa 3: Divide os estados do otimizador + gradientes + parâmetros do modelo entre os workers/GPUs paralelos de dados

d. Optimizer Offload: Move os gradientes + estados do otimizador para a CPU/Disco, construindo sobre o ZERO Stage 2

e. Param Offload: Move os parâmetros do modelo para a CPU/Disco, construindo sobre o ZERO Stage 3

f. Particionamento Hierárquico: Permite treinamento eficiente em múltiplos nós com treinamento paralelo de dados entre nós e sharding ZeRO-3 dentro de um nó, construído sobre o ZeRO Stage 3.

Inferência:

1. O DeepSpeed ZeRO Inference suporta o ZeRO stage 3 com o ZeRO-Infinity. Ele usa o mesmo protocolo ZeRO do treinamento, mas não utiliza um otimizador nem um lr scheduler, e apenas o stage 3 é relevante. Para mais detalhes, veja: ZeRO Inference.

Como funciona?

Pré-requisitos: Instale a versão >=0.6.5 do DeepSpeed. Consulte os detalhes da instalação do DeepSpeed para mais informações.

Primeiro, veremos a integração de fácil uso via accelerate config. Em seguida, a integração mais flexível e com mais recursos do deepspeed config file.

Plugin DeepSpeed do Accelerate

Na sua máquina(s), basta executar:

accelerate config

e responda às perguntas feitas. Será perguntado se você deseja usar um arquivo de configuração para o DeepSpeed, ao que você deve responder não. Em seguida, responda às perguntas a seguir para gerar uma configuração básica do DeepSpeed. Isso gerará um arquivo de configuração que será usado automaticamente para definir corretamente as opções padrão ao fazer

accelerate launch my_script.py --args_to_my_script

Por exemplo, é assim que você executaria o exemplo de NLP examples/nlp_example.py (a partir da raiz do repositório) com o DeepSpeed Plugin:

Exemplo do Plugin DeepSpeed ZeRO Stage-2

compute_environment: LOCAL_MACHINE
deepspeed_config:
 gradient_accumulation_steps: 1
 gradient_clipping: 1.0
 offload_optimizer_device: none
 offload_param_device: none
 zero3_init_flag: true
 zero_stage: 2
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

accelerate launch examples/nlp_example.py --mixed_precision fp16

Exemplo do plugin DeepSpeed com ZeRO Stage-3 e offload para CPU

compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

accelerate launch examples/nlp_example.py --mixed_precision fp16

Atualmente, o Accelerate suporta as seguintes configurações por meio da CLI:

`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
`gradient_clipping`: Enable gradient clipping with value.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
`offload_optimizer_nvme_path`: Decides Nvme Path to offload optimizer states. If unspecified, will default to 'none'.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
`offload_param_nvme_path`: Decides Nvme Path to offload parameters. If unspecified, will default to 'none'.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training.
`deepspeed_moe_layer_cls_names`: Comma-separated list of transformer Mixture-of-Experts (MoE) layer class names (case-sensitive) to wrap ,e.g, `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ...
`deepspeed_hostfile`: DeepSpeed hostfile for configuring multi-node compute resources.
`deepspeed_exclusion_filter`: DeepSpeed exclusion filter string when using mutli-node setup.
`deepspeed_inclusion_filter`: DeepSpeed inclusion filter string when using mutli-node setup.
`deepspeed_multinode_launcher`: DeepSpeed multi-node launcher to use, e.g. `pdsh`, `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5). If unspecified, will default to `pdsh`.
`deepspeed_config_file`: path to the DeepSpeed config file in `json` format. See the next section for more details on this.

Para poder ajustar mais opções, você precisará usar um arquivo de configuração do DeepSpeed.

Arquivo de Configuração do DeepSpeed

Na sua máquina(s), basta executar:

accelerate config

e responder às perguntas feitas. Ele perguntará se você deseja usar um arquivo de configuração para deepspeed, ao que você responde sim e fornece o caminho para o arquivo de configuração do deepspeed. Isso gerará um arquivo de configuração que será usado automaticamente para definir corretamente as opções padrão ao fazer

accelerate launch my_script.py --args_to_my_script

Por exemplo, é assim que você executaria o exemplo de NLP examples/by_feature/deepspeed_with_config_support.py (a partir da raiz do repositório) com o Arquivo de Configuração do DeepSpeed:

Exemplo de arquivo de configuração do DeepSpeed para o ZeRO Stage-2

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/deepspeed_config_templates/zero_stage2_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

com o conteúdo de zero_stage2_config.json sendo:

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

accelerate launch examples/by_feature/deepspeed_with_config_support.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 3 \
--with_tracking \
--report_to "wandb"\

Exemplo de arquivo de configuração do DeepSpeed com ZeRO Stage-3 e offload para CPU

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/deepspeed_config_templates/zero_stage3_offload_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

com o conteúdo de zero_stage3_offload_config.json sendo:

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": "auto"
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

accelerate launch examples/by_feature/deepspeed_with_config_support.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 32 \
--per_device_eval_batch_size 32 \
--num_train_epochs 3 \
--with_tracking \
--report_to "wandb"\

Exemplo de configuração do ZeRO++ Você pode usar os recursos do ZeRO++ usando os parâmetros de configuração apropriados. Observe que o ZeRO++ é uma extensão para o ZeRO Stage 3. Veja como o arquivo de configuração pode ser modificado, a partir do tutorial do ZeRO++ do DeepSpeed:

{
    "zero_optimization": {
        "stage": 3,
        "reduce_bucket_size": "auto",

        "zero_quantized_weights": true,
        "zero_hpz_partition_size": 8,
        "zero_quantized_gradients": true,

        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

Para particionamento hierárquico, o tamanho da partição zero_hpz_partition_size deve idealmente ser definido como o número de GPUs por nó. (Por exemplo, o arquivo de configuração acima pressupõe 8 GPUs por nó)

Alterações importantes no código ao usar o arquivo de configuração do DeepSpeed

1. Otimizadores e agendadores do DeepSpeed. Para mais informações sobre esses componentes, consulte a documentação de DeepSpeed Optimizers e DeepSpeed Schedulers. Veremos as alterações necessárias no código ao usar esses componentes.

   a. DS Optim + DS Scheduler: o caso em que as chaves optimizer e scheduler estão presentes no arquivo de configuração do DeepSpeed. Nessa situação, elas serão usadas e o usuário deve utilizar accelerate.utils.DummyOptim e accelerate.utils.DummyScheduler para substituir os otimizadores e agendadores do PyTorch/personalizados em seu código. Abaixo está o trecho de examples/by_feature/deepspeed_with_config_support.py mostrando isso:

   Copiado

    # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer
    optimizer_cls = (
        torch.optim.AdamW
        if accelerator.state.deepspeed_plugin is None
        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)
   
    # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        lr_scheduler = get_scheduler(
            name=args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=args.num_warmup_steps,
            num_training_steps=args.max_train_steps,
        )
    else:
        lr_scheduler = DummyScheduler(
            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
        )

   b. Otimizador personalizado + Agendador personalizado: o caso em que as chaves optimizer e scheduler estão ausentes no arquivo de configuração do DeepSpeed. Nessa situação, nenhuma alteração de código é necessária por parte do usuário e esse é o caso ao usar a integração via DeepSpeed Plugin. No exemplo acima, podemos ver que o código permanece inalterado se as chaves optimizer e scheduler estiverem ausentes no arquivo de configuração do DeepSpeed.

   c. Otimizador personalizado + Agendador do DS: o caso em que apenas a chave scheduler está presente no arquivo de configuração do DeepSpeed. Nessa situação, o usuário precisa usar accelerate.utils.DummyScheduler para substituir o agendador do PyTorch/personalizado em seu código.

   d. DS Optim + Custom Scheduler: o caso em que apenas a chave optimizer está presente no arquivo de configuração do DeepSpeed. Isso resultará em um erro, pois você só pode usar o DS Scheduler ao usar o DS Optim.

2. Observe os valores auto nos arquivos de configuração do DeepSpeed do exemplo acima. Esses são tratados automaticamente pelo método prepare com base no modelo, nos dataloaders, no otimizador fictício e nos agendadores fictícios fornecidos ao método prepare. Somente os campos auto especificados nos exemplos acima são tratados pelo método prepare e os demais devem ser especificados explicitamente pelo usuário.

Os valores auto são calculados da seguinte forma:

• reduce_bucket_size: hidden_size * hidden_size
• stage3_prefetch_bucket_size: int(0.9 * hidden_size * hidden_size)
• stage3_param_persistence_threshold: 10 * hidden_size

Para que o recurso auto funcione para essas 3 entradas de configuração - o Accelerate usará model.config.hidden_size ou max(model.config.hidden_sizes) como hidden_size. Se nenhuma dessas opções estiver disponível, a inicialização falhará e você terá que definir manualmente essas 3 entradas de configuração. Lembre-se de que as duas primeiras entradas de configuração são os buffers de comunicação - quanto maiores forem, mais eficientes serão as comunicações, e quanto maiores forem, mais memória de GPU elas consumirão, portanto, é um compromisso de desempenho ajustável.

Observações ao usar o arquivo de configuração do DeepSpeed

Abaixo está um script de exemplo usando deepspeed_config_file em diferentes cenários.

Código test.py:

from accelerate import Accelerator
from accelerate.state import AcceleratorState


def main():
    accelerator = Accelerator()
    accelerator.print(f"{AcceleratorState()}")


if __name__ == "__main__":
    main()

Cenário 1: arquivo de configuração do accelerate manipulado manualmente, contendo deepspeed_config_file junto com outras entradas.

1. Conteúdo da configuração do accelerate:

command_file: null
commands: null
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: 'cpu'
  offload_param_device: 'cpu'
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
  deepspeed_config_file: 'ds_config.json'
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
gpu_ids: null
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_name: null
tpu_zone: null
use_cpu: false

2. ds_config.json:

{
    "bf16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": false,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_clipping": 1.0,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": 10,
    "steps_per_print": 2000000
}

3. Saída de accelerate launch test.py:

ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored:
['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device',
'zero3_save_16bit_model', 'mixed_precision'].
Please specify them appropriately in the DeepSpeed config file.
If you are using an accelerate config file, remove other config variables mentioned in the above specified list.
The easiest method is to create a new config following the questionnaire via `accelerate config`.
It will only ask for the necessary config variables when using `deepspeed_config_file`.

Cenário 2: Use a solução do erro para criar um novo config do accelerate e verifique que nenhum erro de ambiguidade é mais lançado.

1. Execute accelerate config:

$ accelerate config
-------------------------------------------------------------------------------------------------------------------------------
In which compute environment are you running?
This machine
-------------------------------------------------------------------------------------------------------------------------------
Which type of machine are you using?
multi-GPU
How many different machines will you use (use more than 1 for multi-node training)? [1]:
Do you wish to optimize your script with torch dynamo?[yes/NO]:
Do you want to use DeepSpeed? [yes/NO]: yes
Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes
Please enter the path to the json DeepSpeed config file: ds_config.json
Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
How many GPU(s) should be used for distributed training? [1]:4
accelerate configuration saved at ds_config_sample.yaml

2. Conteúdo da configuração do accelerate:

compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_config_file: ds_config.json
  zero3_init_flag: true
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false

3. Saída de accelerate launch test.py:

Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: bf16
ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}

Cenário 3: Defina os argumentos do comando accelerate launch relacionados ao DeepSpeed como "auto" no arquivo de configuração do DeepSpeed e verifique se tudo funciona como esperado.

1. Novo ds_config.json com "auto" para os argumentos de comando do accelerate launch DeepSpeed:

{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": "auto",
        "stage3_gather_16bit_weights_on_model_save": "auto",
        "offload_optimizer": {
            "device": "auto"
        },
        "offload_param": {
            "device": "auto"
        }
    },
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "steps_per_print": 2000000
}

2. Saída de accelerate launch --mixed_precision="fp16" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device="cpu" --offload_optimizer_device="nvme" --zero3_save_16bit_model="true" test.py:

Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: fp16
ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}

Nota:

1. Os valores "auto" restantes são tratados na chamada de accelerator.prepare(), conforme explicado no ponto 2 de Important code changes when using DeepSpeed Config File.
2. Somente quando gradient_accumulation_steps é auto, o valor passado ao criar o objeto Accelerator via Accelerator(gradient_accumulation_steps=k) será utilizado. Ao usar o DeepSpeed Plugin, o valor dele será utilizado e ele sobrescreverá o valor passado ao criar o objeto Accelerator.

Salvando e carregando

1. O salvamento e o carregamento de modelos não mudam para o ZeRO Stage-1 e Stage-2.

2. sob o ZeRO Stage-3, state_dict contém apenas os placeholders, pois os pesos do modelo são particionados entre várias GPUs. O ZeRO Stage-3 tem 2 opções:

   a. Salvar os pesos do modelo de 16 bits inteiros para carregá-los diretamente mais tarde usando model.load_state_dict(torch.load(pytorch_model.bin)). Para isso, defina zero_optimization.stage3_gather_16bit_weights_on_model_save como True no arquivo de configuração do DeepSpeed ou defina zero3_save_16bit_model como True no Plugin do DeepSpeed. Observe que essa opção requer a consolidação dos pesos em uma única GPU, o que pode ser lento e exigir muita memória, portanto, use esse recurso apenas quando necessário. Abaixo está o trecho de examples/by_feature/deepspeed_with_config_support.py mostrando isso:

   Copiado

   unwrapped_model = accelerator.unwrap_model(model)
   
   # New Code #
   # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
   # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
   # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
   # For Zero Stages 1 and 2, models are saved as usual in the output directory.
   # The model name saved is `pytorch_model.bin`
   unwrapped_model.save_pretrained(
       args.output_dir,
       is_main_process=accelerator.is_main_process,
       save_function=accelerator.save,
       state_dict=accelerator.get_state_dict(model),
   )

   b. Para obter pesos de 32 bits, primeiro salve o modelo usando model.save_checkpoint(). Abaixo está o trecho de examples/by_feature/deepspeed_with_config_support.py mostrando isso:

   Copiado

   success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
   status_msg = f"checkpointing: PATH={PATH}, ckpt_id={ckpt_id}"
   if success:
       logging.info(f"Success {status_msg}")
   else:
       logging.warning(f"Failure {status_msg}")

   Isso criará as partições do modelo e do otimizador ZeRO, juntamente com o script zero_to_fp32.py, no diretório do checkpoint. Você pode usar esse script para fazer a consolidação offline. Ele não requer arquivos de configuração nem GPUs. Aqui está um exemplo de uso:

   Copiado

   $ cd /path/to/checkpoint_dir
   $ ./zero_to_fp32.py . pytorch_model.bin
   Processing zero checkpoint at global_step1
   Detected checkpoint of type zero stage 3, world_size: 2
   Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)

   Para obter um modelo de 32 bits para salvar/infere, você pode executar:

   Copiado

   from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
   
   unwrapped_model = accelerator.unwrap_model(model)
   fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)

   Se você estiver interessado apenas no state_dict, pode fazer o seguinte:

   Copiado

   from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
   
   state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)

   Observe que todas essas funções requerem ~2x de memória (RAM geral) do tamanho do checkpoint final.

ZeRO Inference

O DeepSpeed ZeRO Inference suporta o ZeRO stage 3 com o ZeRO-Infinity. Ele usa o mesmo protocolo ZeRO do treinamento, mas não utiliza um otimizador nem um lr scheduler, e apenas o stage 3 é relevante. Com a integração do accelerate, você só precisa preparar o modelo e o dataloader, conforme mostrado abaixo:

model, eval_dataloader = accelerator.prepare(model, eval_dataloader)

Algumas ressalvas a serem consideradas

1. A integração atual não suporta Pipeline Parallelism do DeepSpeed.
2. A integração atual não suporta mpu, limitando o paralelismo de tensor que é suportado no Megatron-LM.
3. A integração atual não suporta múltiplos modelos.

Multi-node DeepSpeed

O DeepSpeed suporta inferência e treinamento em vários nós por meio de uma variedade de diferentes lançadores. Você pode especificar um lançador diferente definindo a configuração deepspeed_multinode_launcher na CLI ou no arquivo de configuração do DeepSpeed.

Atualmente, o accelerate suporta a passagem de configuração para os seguintes lançadores multi-nó do DeepSpeed: pdsh (padrão), standard, openmpi, mvapich, mpich, slurm, nossh (requer DeepSpeed >= 0.14.5).

Leia a documentação do DeepSpeed para mais informações sobre os diferentes lançadores. Por padrão, o DeepSpeed tentará usar SSH sem senha da máquina principal para os outros nós para executar o comando do lançador. Nessa configuração, o comando accelerate launch precisa ser executado apenas no nó principal. Se estiver usando o lançador nossh, será necessário executar o comando accelerate launch em cada nó usando a configuração copiada.

Recursos do DeepSpeed

A documentação sobre os aspectos internos relacionados ao DeepSpeed pode ser encontrada aqui.

• GitHub do projeto
• Documentação de uso
• Documentação da API
• Posts de blog

Artigos:

• ZeRO: Otimizações de memória para o treinamento de modelos com trilhões de parâmetros
• ZeRO-Offload: Democratizando o Treinamento de Modelos em Escala de Bilhões
• ZeRO-Infinity: Rompendo a barreira de memória de GPU para aprendizado profundo em escala extrema
• ZeRO++: Comunicação Coletiva Extremamente Eficiente para Treinamento de Modelos Gigantes

Por fim, lembre-se de que o Accelerate apenas integra o DeepSpeed, portanto, se você tiver qualquer problema ou dúvida em relação ao uso do DeepSpeed, abra uma issue no GitHub do DeepSpeed.

Para aqueles interessados nas semelhanças e diferenças entre FSDP e DeepSpeed, confira o guia de conceitos aqui!