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