forestcalled/text-generation-webui
0
1def generate_ds_config(ds_bf16, train_batch_size, nvme_offload_dir):2 '''3 DeepSpeed configuration4 https://huggingface.co/docs/transformers/main_classes/deepspeed5 '''6 7 if nvme_offload_dir:8 ds_config = {9 "fp16": {10 "enabled": not ds_bf16,11 },12 "bf16": {13 "enabled": ds_bf16,14 },15 "zero_optimization": {16 "stage": 3,17 "offload_param": {18 "device": "nvme",19 "nvme_path": nvme_offload_dir,20 "pin_memory": True,21 "buffer_count": 5,22 "buffer_size": 1e9,23 "max_in_cpu": 1e924 },25 "overlap_comm": True,26 "reduce_bucket_size": "auto",27 "contiguous_gradients": True,28 "sub_group_size": 1e8,29 "stage3_prefetch_bucket_size": "auto",30 "stage3_param_persistence_threshold": "auto",31 "stage3_max_live_parameters": "auto",32 "stage3_max_reuse_distance": "auto",33 },34 "aio": {35 "block_size": 262144,36 "queue_depth": 32,37 "thread_count": 1,38 "single_submit": False,39 "overlap_events": True40 },41 "steps_per_print": 2000,42 "train_batch_size": train_batch_size,43 "train_micro_batch_size_per_gpu": 1,44 "wall_clock_breakdown": False45 }46 else:47 ds_config = {48 "fp16": {49 "enabled": not ds_bf16,50 },51 "bf16": {52 "enabled": ds_bf16,53 },54 "zero_optimization": {55 "stage": 3,56 "offload_param": {57 "device": "cpu",58 "pin_memory": True59 },60 "overlap_comm": True,61 "contiguous_gradients": True,62 "reduce_bucket_size": "auto",63 "stage3_prefetch_bucket_size": "auto",64 "stage3_param_persistence_threshold": "auto",65 "stage3_max_live_parameters": "auto",66 "stage3_max_reuse_distance": "auto",67 },68 "steps_per_print": 2000,69 "train_batch_size": train_batch_size,70 "train_micro_batch_size_per_gpu": 1,71 "wall_clock_breakdown": False72 }73 74 return ds_config75