구성 파일
{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},
//fp16과 공존할 수 없음
"bf16": {
"enabled": "auto"
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto",
"freeze_step": 400,
"cuda_aware": false,
"comm_backend_name": "nccl"
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": [0|1|2|3],
"offload_optimizer": {
"device": "[cpu|nvme]",
"pin_memory": true,
"nvme_path": "/local_nvme",
"ratio": 0.3,
"buffer_count": 4,
"fast_init": false
},
"offload_param": {
"device": "[cpu|nvme]",
"nvme_path": "/local_nvme",
"pin_memory": true,
"buffer_count": 5,
"buffer_size": 1e8,
"max_in_cpu": 1e9
},
"allgather_partitions": true,
"allgather_bucket_size": 2e8,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 2e8,
"contiguous_gradients": true,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e6,
"sub_group_size": 1e12,
"elastic_checkpoint": true,
"stage3_gather_16bit_weights_on_model_save": true,
"ignore_unused_parameters": true,
"round_robin_gradients": true,
"zero_hpz_partition_size": 1,
"zero_quantized_weights": true,
"zero_quantized_gradients": true
},
"tensorboard": {
"enabled": true,
"output_path": "output/ds_logs/",
"job_name": "train_bert"
},
"wandb": {
"enabled": true,
"group": "my_group",
"team": "my_team",
"project": "my_project"
},
//그라데이션 폭발을 방지하기 위한 그라데이션 트리밍
"gradient_clipping": 1.0,
//배치 크기 설정 트레인_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation_steps * torch.distributed.get_world_size()
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"gradient_accumulation_steps": "auto",
//인쇄 로그 실행
"steps_per_print": 100,
"wall_clock_breakdown": false,
"dump_state": true
}
매개변수 Auto
위의 구성에서 많은 매개변수가 '자동'으로 설정되어 있는 것을 볼 수 있는데, 이는 두 가지 경우로 나뉩니다:
trainer = Trainer( model=model, tokenizer=tokenizer, args=training_args, **data_module ) trainer.train()ds_config = get_train_ds_config( offload=args.offload, stage=args.zero_stage, enable_tensorboard=args.enable_tensorboard, tb_path=args.tensorboard_path, tb_name="sft", ) ds_config["train_micro_batch_size_per_gpu"] = args.per_device_train_batch_size ds_config["train_batch_size"] = ( args.per_device_train_batch_size * torch.distributed.get_world_size() * args.gradient_accumulation_steps ) model, optimizer, _, lr_scheduler = deepspeed.initialize( model=model, optimizer=optimizer, args=args, config=ds_config, lr_scheduler=lr_scheduler, dist_init_required=True, ) for epoch in range(args.num_train_epochs): model.train() for step, batch in enumerate(train_dataloader): start = time.time() batch = to_device(batch, device) outputs = model(**batch, use_cache=False) loss = outputs.loss if args.print_loss: print( f"Epoch: {epoch}, Step: {step}, Rank: {torch.distributed.get_rank()}, loss = {loss}" ) model.backward(loss) model.step()위의 두 가지 방법의 차이점은 두 번째 방법은 매개 변수를 수동으로 설정해야 한다는 것입니다.
제로 최적화에 대한 이해
"zero_optimization": {
"stage": [0|1|2|3],
"offload_optimizer": {
"device": "cpu",
"pin_memory": true,
},
//CPU에 로드된 모델 매개변수
"offload_param": {
"device": "cpu",
"pin_memory": true,
},
"allgather_partitions": true,
"allgather_bucket_size": 2e8,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 2e8,
"round_robin_gradients": true,
"contiguous_gradients": true,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e6,
"sub_group_size": 1e12,
"elastic_checkpoint": true,
"stage3_gather_16bit_weights_on_model_save": true,
"ignore_unused_parameters": true,
"zero_hpz_partition_size": 1,
"zero_quantized_weights": true,
"zero_quantized_gradients": true
}
ZeRO의 본질:
그래픽 카드 간 데이터 통신의 시간 소모적인 특성으로 인해 일반적으로 데이터는 대량으로 전송됩니다. 따라서 기본적으로 이 매개변수가 조절하는 것은 통신 대역폭과 버킷 크기 간의 절충안입니다.
ZeRO-2 매개변수 분석
- ZeRO-2의 경우: "overlap_comm": true GPU 메모리 사용량을 증가시켜 전체 감소 지연 시간을 줄입니다. overlap_comm은
allgather_bucket_size및 reduce_bucket_size 값의 4.5배를 사용합니다. 따라서 5e8로 설정하면 9GB의 메모리 공간이 필요합니다. 따라서 GPU 메모리가 8GB 이하인 경우 OOM 오류를 방지하려면 이 매개변수를 약 2e8로 줄여야 하며, 이 경우 3.6GB가 필요합니다. GPU가 훨씬 더 큰 경우 OOM이 발생하기 시작하면 동일한 작업을 수행해야 할 수도 있습니다. 버퍼 크기를 줄이면 통신 속도가 느려지는 대신 GPU 메모리가 늘어납니다. 버퍼 크기가 작을수록 통신 속도가 느려지고 GPU가 다른 작업을 위해 더 많은 메모리를 사용할 수 있게 됩니다. 따라서 배치 크기가 큰 것이 중요하다면 트레이닝 시간을 약간 늦추는 것이 좋은 절충안일 수 있습니다. - 라운드_로빈_그레디언트는 CPU 오프로드를 병렬화할 수 있습니다.
ZeRO-3 매개변수 분석
- ,
allgather_partitions, 및 reduce_scatter 구성 파라미터는 ZeRO-3에서는 사용되지 않습니다.
allgather_bucket_size: 1e9
stage3_max_live_parameters: 1e9
OOM 문제가 발생하면 stage3_max_reuse_distance 줄이세요. 메모리는 공유하므로 스택이 쌓이지 않고 총 2GB를 사용하므로 성능에 미치는 영향은 미미합니다.
ZeRO-3의 경우: 다음 구성 값은 모델의 숨겨진 크기에 따라 다르며 수동으로 구성할 필요가 없습니다:
reduce_bucket_size: 감소량
stage3_prefetch_버킷_크기: 프리페치 매개변수를 위한 고정 버퍼의 크기, 통신에 반비례합니다.
stage3_param_지속성_임계값: 통신에 반비례하는 이 임계값보다 작은 매개변수를 나누지 마세요.
sub_group_size 기본값은 1e9입니다. 다음과 같은 경우 기본값을 변경해야 할 수 있습니다:
최적화 단계에서 OOM 발생: 임시 버퍼의 메모리 사용률을 줄이기 위해 sub_group_size를 줄임 최적화 단계에서 시간이 오래 걸림: 데이터 버퍼 증가로 인해 대역폭 사용률을 개선하기 위해 sub_group_size를 늘림.
ZeRO-3의 구성을 조정하여 ZeRO-2의 성능에 가깝게 만들 수 있습니다:
- 를 가장 큰 매개변수보다 큰 숫자(예:
stage3_max_live_parameters)로 설정합니다. 이렇게 하면 파라미터가 GPU에 유지됩니다. - ZeRO-2에는 이 옵션이 없으므로 offload_params를 끄면 모델 파라미터가 CPU에 로드됩니다.
변경하지 않고
stage3_max_reuse_distance꺼도 성능이 크게 향상될 수 있습니다.- 를 가장 큰 매개변수보다 큰 숫자(예:
ZeRO-2 일반적인 구성
{
"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": 3e-5,
"betas": [0.8, 0.999],
"eps": 1e-8,
"weight_decay": 3e-7
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 3e-5,
"warmup_num_steps": 500
}
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"allgather_partitions": true,
"allgather_bucket_size": 2e8,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 2e8,
"contiguous_gradients": true
},
"steps_per_print": 2000,
"wall_clock_breakdown": false
}
ZeRO-3일반적인 구성
{
"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": 3e-5,
"betas": [0.8, 0.999],
"eps": 1e-8,
"weight_decay": 3e-7
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 3e-5,
"warmup_num_steps": 500
}
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
// 이 기능을 끄고 CPU에 매개변수를 로드하지 않도록 하세요!
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": auto,
"stage3_prefetch_bucket_size": auto,
"stage3_param_persistence_threshold": auto,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"steps_per_print": 2000,
"wall_clock_breakdown": false
}
모범 사례
따라서 성능과 GPU 사용량의 균형을 맞추기 위해 다음 단계를 따를 수 있습니다.
먼저 배치 크기를 1로 설정합니다.
- Ampere 이상의 GPU에서는 bf16을, 구형 GPU 아키텍처에서는 fp16을 사용합니다.
- 활성화
--gradient_checkpointing 1또는 직접model.gradient_checkpointing_enable()- OOM이 발생하면 다음 단계를 수행합니다. - 먼저 ZeRO 2단계를 시도하고 OOM이 발생하면 다음 단계를 수행하세요.
- ZeRO 2단계 + offload_optimizer를 사용해 보세요 - OOM이 발생하면 다음 단계를 수행하세요.
- ZeRO 3단계로 전환 - OOM이 발생하면 다음 단계를 수행하세요.
- 오프로드_파람을 CPU에 활성화 - OOM이 발생하면 다음 단계를 수행합니다.
- 오프로드_옵티마이저를 CPU에 활성화 - OOM이 발생하면 다음 단계를 수행합니다.
- 여전히 배치 크기 1에 익숙하지 않다면 먼저 다양한 기본값을 확인하고 가능한 한 줄여보세요. 예를 들어, 생성 기능을 사용하면서 넓은 검색 번들을 사용하지 않는다면 메모리를 많이 차지하므로 더 작게 설정하세요.
