Главная Обзор Помощь
Вход
zhengjl
/
text-generation-webui
1
0
Ответвить 0
Файлы Задачи 0 Запросы на слияние 0 Вики
Ветка: main
Ветки Метки
text-generat...
/
modules
/
deepspeed_parameters.py

deepspeed_parameters.py 2.5 KB
Постоянная ссылка История Исходник

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475 
  1. def generate_ds_config(ds_bf16, train_batch_size, nvme_offload_dir):
    2. 

  2. 

  3.     '''
    4. 

  4.     DeepSpeed configration
    5. 

  5.     https://huggingface.co/docs/transformers/main_classes/deepspeed
    6. 

  6.     '''
    7. 

  7. 

  8.     if nvme_offload_dir:
    9. 

  9.         ds_config = {
    10. 

  10.             "fp16": {
    11. 

  11.                 "enabled": not ds_bf16,
    12. 

  12.             },
    13. 

  13.             "bf16": {
    14. 

  14.                 "enabled": ds_bf16,
    15. 

  15.             },
    16. 

  16.             "zero_optimization": {
    17. 

  17.                 "stage": 3,
    18. 

  18.                 "offload_param": {
    19. 

  19.                     "device": "nvme",
    20. 

  20.                     "nvme_path": nvme_offload_dir,
    21. 

  21.                     "pin_memory": True,
    22. 

  22.                     "buffer_count": 5,
    23. 

  23.                     "buffer_size": 1e9,
    24. 

  24.                     "max_in_cpu": 1e9
    25. 

  25.                 },
    26. 

  26.                 "overlap_comm": True,
    27. 

  27.                 "reduce_bucket_size": "auto",
    28. 

  28.                 "contiguous_gradients": True,
    29. 

  29.                 "sub_group_size": 1e8,
    30. 

  30.                 "stage3_prefetch_bucket_size": "auto",
    31. 

  31.                 "stage3_param_persistence_threshold": "auto",
    32. 

  32.                 "stage3_max_live_parameters": "auto",
    33. 

  33.                 "stage3_max_reuse_distance": "auto",
    34. 

  34.             },
    35. 

  35.             "aio": {
    36. 

  36.                 "block_size": 262144,
    37. 

  37.                 "queue_depth": 32,
    38. 

  38.                 "thread_count": 1,
    39. 

  39.                 "single_submit": False,
    40. 

  40.                 "overlap_events": True
    41. 

  41.             },
    42. 

  42.             "steps_per_print": 2000,
    43. 

  43.             "train_batch_size": train_batch_size,
    44. 

  44.             "train_micro_batch_size_per_gpu": 1,
    45. 

  45.             "wall_clock_breakdown": False
    46. 

  46.         }
    47. 

  47.     else:
    48. 

  48.         ds_config = {
    49. 

  49.             "fp16": {
    50. 

  50.                 "enabled": not ds_bf16,
    51. 

  51.             },
    52. 

  52.             "bf16": {
    53. 

  53.                 "enabled": ds_bf16,
    54. 

  54.             },
    55. 

  55.             "zero_optimization": {
    56. 

  56.                 "stage": 3,
    57. 

  57.                 "offload_param": {
    58. 

  58.                     "device": "cpu",
    59. 

  59.                     "pin_memory": True
    60. 

  60.                 },
    61. 

  61.                 "overlap_comm": True,
    62. 

  62.                 "contiguous_gradients": True,
    63. 

  63.                 "reduce_bucket_size": "auto",
    64. 

  64.                 "stage3_prefetch_bucket_size": "auto",
    65. 

  65.                 "stage3_param_persistence_threshold": "auto",
    66. 

  66.                 "stage3_max_live_parameters": "auto",
    67. 

  67.                 "stage3_max_reuse_distance": "auto",
    68. 

  68.             },
    69. 

  69.             "steps_per_print": 2000,
    70. 

  70.             "train_batch_size": train_batch_size,
    71. 

  71.             "train_micro_batch_size_per_gpu": 1,
    72. 

  72.             "wall_clock_breakdown": False
    73. 

  73.         }
    74. 

  74. 

  75.     return ds_config
    76.