deepspeed (#288)

* deepspeed * shard * full param deepspeed works by this commit * offload optimizer & documentation * format & fix save deepspeed weight * format & update save_checkpoint * update pipfile * update pipfile * zero init for transformers * add some new config * fix bug * min 1e6 * update deepspeed config * Update requirements.txt * remove duplicate code * throw warning when compile w/ deepspeed * black * integrate deepspeed into wrap_model_distributed * remove unuse code * style * fix bug * fix bug * max token len to 16k * deepspeed save lora * update get optimizer * fix check disk * comment out offload CPU * Pipfile.lock * Update requirements.txt * make black * add default * minor fix * minor fix * minor fix * fix val loader * potential val loader fix * update * lock * Update requirements.txt * improve model saving for deepspeed * solved INFLIGHT problem * update doc * deepspeed default push to hub by cpu * Revert "improve model saving for deepspeed" This reverts commit 62fc9c5. * remove unuse code * Update requirements.txt * deepspeed==0.11.1 * Update requirements.txt * temp fix for deepspeed slow gen * style * style * fix --------- Co-authored-by: haqishen <haqishen@gmail.com> Co-authored-by: Philipp Singer <killver@gmail.com> Co-authored-by: psinger <psinger@users.noreply.github.com>
h2oai · Oct 24, 2023 · 67d3a3c · 67d3a3c
1 parent 08475e3
commit 67d3a3c
Show file tree

Hide file tree

Showing 15 changed files with 814 additions and 494 deletions.
diff --git a/Pipfile b/Pipfile
@@ -50,6 +50,7 @@ tiktoken = "==0.5.1"
 hf-transfer = "==0.1.3"
 peft = "==0.5.0"
 azure-storage-file-datalake = ">=12.12.0"
+deepspeed = "==0.11.1"
 keyring = "==24.2.0"
 
 [dev-packages]

diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/documentation/docs/tooltips/experiments/_deepspeed-offload-optimizer.mdx b/documentation/docs/tooltips/experiments/_deepspeed-offload-optimizer.mdx
@@ -0,0 +1 @@
+Whether to offload optimizer to cpu for saving more GPU ram during training. Note that turn on offload_optimizer would further slow down training.
diff --git a/documentation/docs/tooltips/experiments/_deepspeed-reduce-bucket-size.mdx b/documentation/docs/tooltips/experiments/_deepspeed-reduce-bucket-size.mdx
@@ -0,0 +1 @@
+Number of elements reduced/allreduced at a time. Limits the memory required for the allgather for large model sizes. Smaller values use less memory, but slow down training.
diff --git a/documentation/docs/tooltips/experiments/_deepspeed-stage3-max-live-parameters.mdx b/documentation/docs/tooltips/experiments/_deepspeed-stage3-max-live-parameters.mdx
@@ -0,0 +1 @@
+The maximum number of parameters resident per GPU before releasing. Smaller values use less memory, but slow down training.
diff --git a/documentation/docs/tooltips/experiments/_deepspeed-stage3-max-reuse-distance.mdx b/documentation/docs/tooltips/experiments/_deepspeed-stage3-max-reuse-distance.mdx
@@ -0,0 +1 @@
+Do not release a parameter if it will be reused within this threshold of parameters. Smaller values use less memory, but slow down training.
diff --git a/...ion/docs/tooltips/experiments/_deepspeed-stage3-param-persistence-threshold.mdx b/...ion/docs/tooltips/experiments/_deepspeed-stage3-param-persistence-threshold.mdx
@@ -0,0 +1 @@
+Do not partition parameters smaller than this threshold. Smaller values use less memory, but can greatly increase communication and slow down training. (especially latency-bound messages).
diff --git a/documentation/docs/tooltips/experiments/_deepspeed-stage3-prefetch-bucket-size.mdx b/documentation/docs/tooltips/experiments/_deepspeed-stage3-prefetch-bucket-size.mdx
@@ -0,0 +1 @@
+Maximum number of parameter elements to fetch ahead of use. Smaller values use less memory, but slow down training..
diff --git a/documentation/docs/tooltips/experiments/_use-deepspeed.mdx b/documentation/docs/tooltips/experiments/_use-deepspeed.mdx
@@ -0,0 +1 @@
+Whether to use deepspeed for saving GPU ram during training. Note that turning on DeepSpeed can slow down training.
diff --git a/llm_studio/app_utils/sections/experiment.py b/llm_studio/app_utils/sections/experiment.py
@@ -1680,7 +1680,9 @@ async def experiment_push_to_huggingface_dialog(q: Q, error: str = ""):
         num_running_queued = len(
             experiments[experiments["status"].isin(["queued", "running"])]
         )
-        if num_running_queued > 0:
+        experiment_path = q.client["experiment/display/experiment_path"]
+        cfg = load_config_yaml(os.path.join(experiment_path, "cfg.yaml"))
+        if num_running_queued > 0 or cfg.environment.use_deepspeed:
             default_device = "cpu"
 
         try:

diff --git a/llm_studio/app_utils/utils.py b/llm_studio/app_utils/utils.py
@@ -114,36 +114,41 @@ def start_process(
     env = {**os.environ, **env_vars}
 
     if num_gpus == 0:
-        p = subprocess.Popen(
-            [
-                "python",
-                "train_wave.py",
-                "-Y",
-                config_name,
-                "-Q",
-                ",".join([str(x) for x in process_queue]),
-            ],
-            env=env,
-        )
+        cmd = [
+            "python",
+            "train_wave.py",
+            "-Y",
+            config_name,
+        ]
     # Do not delete for debug purposes
     # elif num_gpus == 1:
-    #     p = subprocess.Popen(
-    #         [
-    #             "env",
-    #             f"CUDA_VISIBLE_DEVICES={','.join(gpu_list)}",
-    #             "python",
-    #             "-u",
-    #             "train_wave.py",
-    #             "-P",
-    #             config_name,
-    #             "-Q",
-    #             ",".join([str(x) for x in process_queue]),
-    #         ]
-    #     )
+    #     cmd = [
+    #         "env",
+    #         f"CUDA_VISIBLE_DEVICES={','.join(gpu_list)}",
+    #         "python",
+    #         "-u",
+    #         "train_wave.py",
+    #         "-P",
+    #         config_name,
+    #     ]
     else:
         free_port = find_free_port()
-        p = subprocess.Popen(
-            [
+        if cfg.environment.use_deepspeed:
+            logger.info("Starting deepspeed...")
+            cmd = [
+                "env",
+                "deepspeed",
+                "--include",
+                f"localhost:{','.join(gpu_list)}",
+                "--master_port",
+                f"{str(free_port)}",
+                "train_wave.py",
+                "-Y",
+                config_name,
+            ]
+        else:
+            logger.info("Starting torchrun...")
+            cmd = [
                 "env",
                 f"CUDA_VISIBLE_DEVICES={','.join(gpu_list)}",
                 "torchrun",
@@ -152,11 +157,17 @@ def start_process(
                 "train_wave.py",
                 "-Y",
                 config_name,
-                "-Q",
-                ",".join([str(x) for x in process_queue]),
-            ],
-            env=env,
-        )
+            ]
+
+    if len(process_queue) > 0:
+        cmd.append("-Q")
+        cmd.append(",".join([str(x) for x in process_queue]))
+
+    p = subprocess.Popen(
+        cmd,
+        env=env,
+    )
+
     logger.info(f"Percentage of RAM memory used: {psutil.virtual_memory().percent}")
 
     return p

diff --git a/llm_studio/python_configs/text_causal_language_modeling_config.py b/llm_studio/python_configs/text_causal_language_modeling_config.py
@@ -227,9 +227,9 @@ class ConfigNLPCausalLMTokenizer(DefaultConfig):
 
     def __post_init__(self):
         super().__post_init__()
-        self._possible_values["max_length_prompt"] = (32, 8192, 32)
-        self._possible_values["max_length_answer"] = (32, 8192, 32)
-        self._possible_values["max_length"] = (32, 8192, 32)
+        self._possible_values["max_length_prompt"] = (32, 1024 * 16, 32)
+        self._possible_values["max_length_answer"] = (32, 1024 * 16, 32)
+        self._possible_values["max_length"] = (32, 1024 * 16, 32)
         self._possible_values["padding_quantile"] = (0, 1, 0.01)
         self._padding_side = "left"
 
@@ -343,6 +343,13 @@ class ConfigNLPCausalLMEnvironment(DefaultConfig):
 
     compile_model: bool = False
     use_fsdp: bool = False
+    use_deepspeed: bool = False
+    deepspeed_reduce_bucket_size: int = 1e6
+    deepspeed_stage3_prefetch_bucket_size: int = 1e6
+    deepspeed_stage3_param_persistence_threshold: int = 1e6
+    #     deepspeed_offload_optimizer: bool = False
+    #     deepspeed_stage3_max_live_parameters: int = 1e9
+    #     deepspeed_stage3_max_reuse_distance: int = 1e9
 
     find_unused_parameters: bool = False
     trust_remote_code: bool = True
@@ -376,6 +383,37 @@ def __post_init__(self):
 
         self._possible_values["number_of_workers"] = (1, multiprocessing.cpu_count(), 1)
         self._possible_values["seed"] = possible_values.Number(step=1, min=-1)
+        self._possible_values["deepspeed_reduce_bucket_size"] = possible_values.Number(
+            step=1, min=1e6
+        )
+        self._possible_values[
+            "deepspeed_stage3_prefetch_bucket_size"
+        ] = possible_values.Number(step=1, min=1e6)
+        self._possible_values[
+            "deepspeed_stage3_param_persistence_threshold"
+        ] = possible_values.Number(step=1, min=1e6)
+        self._possible_values[
+            "deepspeed_stage3_max_live_parameters"
+        ] = possible_values.Number(step=1, min=1e6)
+        self._possible_values[
+            "deepspeed_stage3_max_reuse_distance"
+        ] = possible_values.Number(step=1, min=1e6)
+        self._nesting.add(
+            [
+                "deepspeed_reduce_bucket_size",
+                "deepspeed_stage3_prefetch_bucket_size",
+                "deepspeed_stage3_param_persistence_threshold",
+                # "deepspeed_offload_optimizer",
+            ],
+            [Dependency(key="use_deepspeed", value=False, is_set=False)],
+        )
+        # self._nesting.add(
+        #     [
+        #         "deepspeed_stage3_max_live_parameters",
+        #         "deepspeed_stage3_max_reuse_distance",
+        #     ],
+        #     [Dependency(key="deepspeed_offload_optimizer", value=False, is_set=False)],  # noqa: E501
+        # )
 
 
 @dataclass