Accelerate documentation

DeepSpeed

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

DeepSpeed

DeepSpeed는 ZeRO 논문에 설명된 모든 내용을 구현합니다. 주요 최적화 사항은 다음과 같습니다:

  1. 옵티마이저 상태 파티셔닝 (ZeRO 단계 1)
  2. Gradient partitioning (ZeRO stage 2)
  3. 파라미터 분할 (ZeRO stage 3)
  4. 사용자 정의 혼합 정밀도 학습 처리
  5. fast CUDA-extension-based optimizers
  6. CPU 및 Disk/NVMe로의 ZeRO-Offload
  7. ZeRO++의 모델 파라미터 계층적 분할

ZeRO-Offload에는 전용 논문이 있습니다: ZeRO-Offload: Democratizing Billion-Scale Model Training. NVMe 지원은 논문 ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning에 설명되어 있습니다.

DeepSpeed ZeRO-2는 추론에는 기능이 유용하지 않으므로 주로 학습에만 사용됩니다.

DeepSpeed ZeRO-3는 단일 GPU에서는 불가능한 대규모 모델을 여러 GPU에 로드할 수 있게 해주므로 추론에도 사용할 수 있습니다.

Accelerate는 DeepSpeed를 두 가지 옵션을 통해 통합합니다:

  1. accelerate config에서 deepspeed config file을 지정하여 DeepSpeed 기능을 통합합니다. 사용자 정의 config 파일을 제공하거나 템플릿을 사용하면 됩니다. 이 문서의 대부분은 이 기능에 초점을 맞추고 있습니다. 이 기능은 DeepSpeed의 모든 핵심 기능을 지원하며 사용자에게 많은 유연성을 제공합니다. config에 따라 코드의 몇 줄을 변경해야 할 수 있습니다.
  2. deepspeed_plugin을 통한 통합. DeepSpeed 기능의 일부만 지원하며, 나머지 구성은 기본 옵션을 사용합니다. 사용자는 코드를 변경할 필요가 없으며, DeepSpeed의 대부분의 기본 설정이 적합한 사용자에게 좋습니다.

무엇이 통합되었나요?

학습:

  1. Accelerate는 DeepSpeed ZeRO의 모든 기능을 통합합니다. 여기에는 ZeRO 단계 1, 2, 3을 비롯해 ZeRO-Offload, ZeRO-Infinity(디스크/NVMe로 오프로드 가능), ZeRO++가 모두 포함됩니다. 아래는 ZeRO - Zero Redundancy Optimizer를 활용한 Data Parallelism에 대한 간단한 설명과 이 블로그 게시물의 다이어그램입니다. ZeRO Data Parallelism

(출처: 링크)

a. Stage 1 : 옵티마이저 상태를 데이터 병렬 워커/GPU에 걸쳐 샤드합니다

b. Stage 2 : 옵티마이저 상태와 그라디언트를 데이터 병렬 워커/GPU에 걸쳐 샤딩

c. Stage 3: 옵티마이저 상태 + 그래디언트 + 모델 매개변수를 데이터 병렬 워커/GPU에 걸쳐 샤드

d. Optimizer Offload: ZERO Stage 2를 기반으로 그래디언트와 옵티마이저 상태를 CPU/디스크로 오프로드합니다

e. Param Offload: ZERO Stage 3를 기반으로 모델 파라미터를 CPU/디스크로 오프로드합니다

f. 계층적 파티셔닝: 노드 간 데이터 병렬 학습과 노드 내 ZeRO-3 샤딩을 통해 효율적인 멀티노드 학습을 가능하게 하며, ZeRO Stage 3를 기반으로 구축됩니다.

Note: With respect to Disk Offload, the disk should be an NVME for decent speed but it technically works on any Disk

추론:

  1. DeepSpeed ZeRO Inference는 ZeRO-Infinity를 사용하여 ZeRO stage 3를 지원합니다. 학습과 동일한 ZeRO 프로토콜을 사용하지만, 옵티마이저와 lr 스케줄러는 사용하지 않으며 stage 3만 관련이 있습니다. 자세한 내용은 ZeRO Inference를 참조하세요.

동작 원리

사전 요구 사항: DeepSpeed 버전 >=0.6.5를 설치하세요. 자세한 내용은 DeepSpeed 설치 정보를 참고하세요.

먼저 accelerate config를 통한 간편한 통합을 살펴봅니다. 이후 더 유연하고 기능이 풍부한 deepspeed config file 통합을 다룹니다.

Accelerate DeepSpeed 플러그인

컴퓨터에서 다음을 실행하세요:

accelerate config

그리고 질문한 내용에 답합니다. DeepSpeed 설정 파일을 사용할지 묻는 질문에 '아니요'라고 답해야 합니다. 그런 다음 다음 질문에 답하여 기본 DeepSpeed 설정을 생성하세요. 이 과정을 통해 doing 시 기본 옵션을 올바르게 설정하는 데 자동으로 사용되는 설정 파일이 생성됩니다.

accelerate launch my_script.py --args_to_my_script

예를 들어, DeepSpeed Plugin을 사용하여 저장소 루트에서 NLP 예제 examples/nlp_example.py를 실행하는 방법은 다음과 같습니다:

ZeRO Stage-2 DeepSpeed 플러그인 예제

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

CPU Offload를 사용한 ZeRO Stage-3 DeepSpeed 플러그인 예제

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

현재 Accelerate는 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.

더 많은 옵션을 조정하려면 DeepSpeed 설정 파일을 사용해야 합니다.

DeepSpeed 설정 파일

컴퓨터에서 다음을 실행하세요:

accelerate config

그리고 질문한 항목에 답합니다. deepspeed 설정 파일을 사용할지 묻는데, yes라고 답하고 deepspeed 설정 파일의 경로를 제공합니다. 이 과정을 통해 설정 파일이 생성되며, 다음 작업을 수행할 때 기본 옵션을 적절히 설정하기 위해 자동으로 사용됩니다.

accelerate launch my_script.py --args_to_my_script

예를 들어, DeepSpeed Config File을 사용하여 NLP 예제 examples/by_feature/deepspeed_with_config_support.py를 실행하는 방법(레포지토리의 루트에서)은 다음과 같습니다:

ZeRO Stage-2 DeepSpeed 구성 파일 예시

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

zero_stage2_config.json의 내용은 다음과 같습니다:

{
    "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"\

CPU offload를 사용한 ZeRO Stage-3 DeepSpeed 설정 파일 예시

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

zero_stage3_offload_config.json의 내용은 다음과 같습니다:

{
    "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"\

ZeRO++ 설정 예시 ZeRO++의 기능을 사용하려면 적절한 설정 매개변수를 사용하십시오. ZeRO++는 ZeRO Stage 3의 확장입니다. DeepSpeed의 ZeRO++ 튜토리얼에서 설정 파일을 수정하는 방법은 다음과 같습니다:

{
    "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
    }
}

계층적 분할의 경우, 분할 크기 zero_hpz_partition_size는 이상적으로 노드당 GPU 수로 설정해야 합니다. (예: 위의 설정 파일은 노드당 8개의 GPU를 가정합니다)

DeepSpeed Config File 사용 시 주요 코드 변경 사항

  1. DeepSpeed Optimizers와 Schedulers. 이들에 대한 자세한 내용은 DeepSpeed Optimizers 및 DeepSpeed Schedulers 문서를 참고하세요. 이들을 사용할 때 코드에서 필요한 변경 사항을 살펴보겠습니다.

    a. DS Optim + DS Scheduler: DeepSpeed 설정 파일에 optimizer와 scheduler 키가 모두 존재하는 경우입니다. 이 경우 해당 설정이 사용되며, 사용자는 코드에서 PyTorch/커스텀 옵티마이저와 스케줄러를 accelerate.utils.DummyOptim과 accelerate.utils.DummyScheduler로 대체해야 합니다. 아래는 이를 보여주는 examples/by_feature/deepspeed_with_config_support.py의 코드 조각입니다:

     # 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. Custom Optim + Custom Scheduler: DeepSpeed 설정 파일에 optimizer 및 scheduler 키가 모두 없는 경우입니다. 이 경우 사용자는 코드 변경 없이 DeepSpeed Plugin을 통한 통합을 사용할 수 있습니다. 위 예시에서 DeepSpeed 설정 파일에 optimizer 및 scheduler 키가 없으면 코드가 변경되지 않음을 확인할 수 있습니다.

    c. Custom Optim + DS Scheduler: DeepSpeed 설정 파일에 scheduler 키만 존재하는 경우입니다. 이 경우 사용자는 코드에서 PyTorch/Custom 스케줄러를 accelerate.utils.DummyScheduler로 대체해야 합니다.

    d. DS Optim + Custom Scheduler: DeepSpeed 설정 파일에 optimizer 키만 존재하는 경우입니다. DS Optim을 사용하지 않으면 DS Scheduler를 사용할 수 없으므로 오류가 발생합니다.

  2. 위 예시의 DeepSpeed 설정 파일에서 auto 값을 확인해 보세요. 이러한 값은 prepare 메서드에 전달된 모델, 데이터로더, 더미 옵티마이저 및 더미 스케줄러를 기반으로 prepare 메서드가 자동으로 처리합니다. 위 예시에 명시된 auto 필드만 prepare 메서드에 의해 처리되며, 나머지는 사용자가 명시적으로 지정해야 합니다.

auto 값은 다음과 같이 계산됩니다:

  • 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

이 3개의 설정 항목에 대해 auto 기능이 작동하려면, Accelerate는 hidden_size로 model.config.hidden_size 또는 max(model.config.hidden_sizes)를 사용합니다. 이 둘 중 어느 것도 사용할 수 없으면 실행이 실패하며, 이 3개의 설정 항목을 수동으로 설정해야 합니다. 첫 번째와 두 번째 설정 항목은 통신 버퍼임을 기억하세요. 버퍼가 클수록 통신 효율이 높아지지만, 그만큼 GPU 메모리 사용량도 증가하므로 조정 가능한 성능 트레이드오프입니다.

DeepSpeed 설정 파일 사용 시 참고 사항

아래는 deepspeed_config_file를 다양한 시나리오에서 사용하는 샘플 스크립트입니다.

test.py 코드:

from accelerate import Accelerator
from accelerate.state import AcceleratorState


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


if __name__ == "__main__":
    main()

시나리오 1: deepspeed_config_file와 다른 항목들이 포함된 accelerate 설정 파일을 수동으로 수정한 경우.

  1. 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
  1. 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
}
  1. 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`.

시나리오 2: 오류 해결 방법을 사용하여 새로운 accelerate config를 생성하고 이제 모호성 오류가 발생하지 않는지 확인합니다.

  1. 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
  1. 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
  1. 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}}

시나리오 3: DeepSpeed 구성 파일에서 accelerate launch 명령 인자 중 DeepSpeed 관련 항목을 "auto"로 설정하고, 모든 것이 예상대로 작동하는지 확인합니다.

  1. accelerate launch DeepSpeed 명령 인수의 "auto"를 사용하는 새로운 ds_config.json:
{
    "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
}
  1. 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}}

참고:

  1. Remaining "auto" values are handled in accelerator.prepare() call as explained in point 2 of Important code changes when using DeepSpeed Config File.
  2. gradient_accumulation_steps가 auto인 경우에만, Accelerator(gradient_accumulation_steps=k)를 통해 Accelerator 객체를 생성할 때 전달된 값이 사용됩니다. DeepSpeed Plugin을 사용하는 경우, 해당 플러그인의 값이 사용되며 Accelerator 객체를 생성할 때 전달된 값을 덮어씁니다.

저장 및 로드

  1. ZeRO Stage-1과 Stage-2의 경우 모델 저장 및 로딩 방식은 변경되지 않습니다.

  2. ZeRO Stage-3에서는 모델 가중치가 여러 GPU에 분할되어 있기 때문에 state_dict에는 플레이스홀더만 포함됩니다. ZeRO Stage-3에는 두 가지 옵션이 있습니다:

    a. 16bit 모델 가중치를 전체로 저장하여 나중에 model.load_state_dict(torch.load(pytorch_model.bin))를 사용하여 직접 로드합니다. 이를 위해 DeepSpeed Config 파일에서 zero_optimization.stage3_gather_16bit_weights_on_model_save를 True로 설정하거나 DeepSpeed Plugin에서 zero3_save_16bit_model를 True로 설정합니다. 이 옵션은 가중치를 하나의 GPU에 통합해야 하므로 느리고 메모리 요구량이 높을 수 있으므로 필요한 경우에만 이 기능을 사용하세요. 아래는 이를 보여주는 examples/by_feature/deepspeed_with_config_support.py의 스니펫입니다:

    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. 32비트 가중치를 얻으려면 먼저 model.save_checkpoint()를 사용하여 모델을 저장합니다. 아래는 이를 보여주는 examples/by_feature/deepspeed_with_config_support.py의 코드 조각입니다:

    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}")

    이 작업을 수행하면 체크포인트 디렉터리에 ZeRO 모델 및 옵티마이저 파티션과 함께 zero_to_fp32.py 스크립트가 생성됩니다. 이 스크립트를 사용하여 오프라인 통합을 수행할 수 있습니다. 설정 파일이나 GPU가 필요하지 않습니다. 사용 예시는 다음과 같습니다:

    $ 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)

    저장/추론을 위한 32bit 모델을 얻으려면 다음을 수행할 수 있습니다:

    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)

    state_dict에만 관심이 있다면 다음을 수행할 수 있습니다:

    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)

    이 모든 함수는 최종 체크포인트 크기의 약 2배의 메모리(일반 RAM)가 필요하다는 점에 유의하세요.

ZeRO 추론

DeepSpeed ZeRO Inference는 ZeRO-Infinity를 사용하여 ZeRO stage 3를 지원합니다. 학습과 동일한 ZeRO 프로토콜을 사용하지만, 옵티마이저와 lr 스케줄러는 사용하지 않으며 stage 3만 관련이 있습니다. accelerate 통합을 사용하면 아래에 표시된 대로 모델과 dataloader만 준비하면 됩니다:

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

주의해야 할 몇 가지 사항

  1. 현재 통합은 DeepSpeed의 Pipeline Parallelism을 지원하지 않습니다.
  2. 현재 통합은 mpu를 지원하지 않아 Megatron-LM에서 지원하는 텐서 병렬성의 범위가 제한됩니다.
  3. 현재 통합은 여러 모델을 지원하지 않습니다.

멀티 노드 DeepSpeed

DeepSpeed는 다양한 런처를 통해 멀티 노드 추론 및 학습을 지원합니다. CLI 또는 DeepSpeed 설정 파일에서 deepspeed_multinode_launcher 설정을 지정하여 다른 런처를 사용할 수 있습니다.

현재 Accelerate는 다음 DeepSpeed 멀티노드 런처에 대한 설정 전달을 지원합니다: pdsh(기본값), standard, openmpi, mvapich, mpich, slurm, nossh(DeepSpeed >= 0.14.5 필요).

다양한 런처에 대한 자세한 내용은 DeepSpeed 문서를 참고하세요. 기본적으로 DeepSpeed는 메인 머신 노드에서 다른 노드로 패스워드 없는 SSH를 사용하여 런처 명령을 실행하려고 합니다. 이 구성에서는 accelerate launch 명령을 메인 노드에서만 실행하면 됩니다. nossh 런처를 사용하는 경우, 복사된 구성을 사용하여 모든 노드에서 accelerate launch 명령을 실행해야 합니다.

DeepSpeed 리소스

deepspeed와 관련된 내부 구조에 대한 문서는 여기에서 확인할 수 있습니다.

논문:

마지막으로, Accelerate는 DeepSpeed만 통합한다는 점을 기억해 주세요. 따라서 DeepSpeed 사용과 관련하여 문제나 질문이 있다면 DeepSpeed GitHub에 이슈를 등록해 주세요.

FSDP와 DeepSpeed의 유사점과 차이점에 관심이 있으시다면, 여기의 개념 가이드를 확인해 보세요!

Update on GitHub