|
def generate_ds_config(ds_bf16, train_batch_size, nvme_offload_dir): |
|
''' |
|
DeepSpeed configuration |
|
https://huggingface.co/docs/transformers/main_classes/deepspeed |
|
''' |
|
|
|
if nvme_offload_dir: |
|
ds_config = { |
|
"fp16": { |
|
"enabled": not ds_bf16, |
|
}, |
|
"bf16": { |
|
"enabled": ds_bf16, |
|
}, |
|
"zero_optimization": { |
|
"stage": 3, |
|
"offload_param": { |
|
"device": "nvme", |
|
"nvme_path": nvme_offload_dir, |
|
"pin_memory": True, |
|
"buffer_count": 5, |
|
"buffer_size": 1e9, |
|
"max_in_cpu": 1e9 |
|
}, |
|
"overlap_comm": True, |
|
"reduce_bucket_size": "auto", |
|
"contiguous_gradients": True, |
|
"sub_group_size": 1e8, |
|
"stage3_prefetch_bucket_size": "auto", |
|
"stage3_param_persistence_threshold": "auto", |
|
"stage3_max_live_parameters": "auto", |
|
"stage3_max_reuse_distance": "auto", |
|
}, |
|
"aio": { |
|
"block_size": 262144, |
|
"queue_depth": 32, |
|
"thread_count": 1, |
|
"single_submit": False, |
|
"overlap_events": True |
|
}, |
|
"steps_per_print": 2000, |
|
"train_batch_size": train_batch_size, |
|
"train_micro_batch_size_per_gpu": 1, |
|
"wall_clock_breakdown": False |
|
} |
|
else: |
|
ds_config = { |
|
"fp16": { |
|
"enabled": not ds_bf16, |
|
}, |
|
"bf16": { |
|
"enabled": ds_bf16, |
|
}, |
|
"zero_optimization": { |
|
"stage": 3, |
|
"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", |
|
"stage3_max_live_parameters": "auto", |
|
"stage3_max_reuse_distance": "auto", |
|
}, |
|
"steps_per_print": 2000, |
|
"train_batch_size": train_batch_size, |
|
"train_micro_batch_size_per_gpu": 1, |
|
"wall_clock_breakdown": False |
|
} |
|
|
|
return ds_config |
|
|