[Deepspeed Wav2vec2] integration (#11638)

* wip * wip - but working with https://github.com/microsoft/DeepSpeed/pull/1044 * cleanup * workaround * working 5/8 modes * solve fp32 distributed zero3 * style * sync * sync * rework * deprecation * cleanup * https://github.com/microsoft/DeepSpeed/pull/1044 pr was merged * clean up * add a guide * more prose * more prose * fix * more prose * sub_group_size was too big * Apply suggestions from code review Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> * refactor * bug fix * make the true check explicit * new deepspeed release Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
2021-06-08 12:32:03 -07:00
parent 32290d87f6
commit 11d86d3de4
11 changed files with 496 additions and 64 deletions
--- a/examples/research_projects/wav2vec2/README.md
+++ b/examples/research_projects/wav2vec2/README.md
@@ -127,3 +127,60 @@ logs references and predictions. Using the Buckwalter format, text is also logge

 `--max_duration_in_seconds="15"` filters out examples whose audio is longer than the specified limit,
 which helps with capping GPU memory usage.
+
+
+### DeepSpeed Integration
+
+To learn how to deploy Deepspeed Integration please refer to [this guide](https://huggingface.co/transformers/master/main_classes/deepspeed.html#deepspeed-trainer-integration).
+
+But to get started quickly all you need is to install:
+```
+pip install deepspeed
+```
+and then use the default configuration files in this directory:
+
+* `ds_config_wav2vec2_zero2.json`
+* `ds_config_wav2vec2_zero3.json`
+
+Here are examples of how you can use DeepSpeed:
+
+(edit the value for `--num_gpus` to match the number of GPUs you have)
+
+ZeRO-2:
+
+```
+PYTHONPATH=../../../src deepspeed --num_gpus 2 \
+run_asr.py \
+--output_dir=output_dir --num_train_epochs=2 --per_device_train_batch_size=2 \
+--per_device_eval_batch_size=2 --evaluation_strategy=steps --save_steps=500 --eval_steps=100 \
+--logging_steps=5 --learning_rate=5e-4 --warmup_steps=3000 \
+--model_name_or_path=patrickvonplaten/wav2vec2_tiny_random_robust \
+--dataset_name=patrickvonplaten/librispeech_asr_dummy --dataset_config_name=clean \
+--train_split_name=validation --validation_split_name=validation --orthography=timit \
+--preprocessing_num_workers=1 --group_by_length --freeze_feature_extractor --verbose_logging \
+--deepspeed ds_config_wav2vec2_zero2.json
+```
+
+For ZeRO-2 with more than 1 gpu you need to use (which is already in the example configuration file):
+```
+    "zero_optimization": {
+        ...
+        "find_unused_parameters": true,
+        ...
+    }
+```
+
+ZeRO-3:
+
+```
+PYTHONPATH=../../../src deepspeed --num_gpus 2 \
+run_asr.py \
+--output_dir=output_dir --num_train_epochs=2 --per_device_train_batch_size=2 \
+--per_device_eval_batch_size=2 --evaluation_strategy=steps --save_steps=500 --eval_steps=100 \
+--logging_steps=5 --learning_rate=5e-4 --warmup_steps=3000 \
+--model_name_or_path=patrickvonplaten/wav2vec2_tiny_random_robust \
+--dataset_name=patrickvonplaten/librispeech_asr_dummy --dataset_config_name=clean \
+--train_split_name=validation --validation_split_name=validation --orthography=timit \
+--preprocessing_num_workers=1 --group_by_length --freeze_feature_extractor --verbose_logging \
+--deepspeed ds_config_wav2vec2_zero3.json
+```
--- a/examples/research_projects/wav2vec2/ds_config_wav2vec2_zero2.json
+++ b/examples/research_projects/wav2vec2/ds_config_wav2vec2_zero2.json
@@ -0,0 +1,51 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "betas": "auto",
+            "eps": "auto",
+            "weight_decay": "auto"
+        }
+    },
+
+    "scheduler": {
+        "type": "WarmupLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto"
+        }
+    },
+
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "find_unused_parameters": true,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
--- a/examples/research_projects/wav2vec2/ds_config_wav2vec2_zero3.json
+++ b/examples/research_projects/wav2vec2/ds_config_wav2vec2_zero3.json
@@ -0,0 +1,57 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "betas": "auto",
+            "eps": "auto",
+            "weight_decay": "auto"
+        }
+    },
+
+    "scheduler": {
+        "type": "WarmupLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto"
+        }
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_fp16_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
--- a/examples/research_projects/wav2vec2/test_wav2vec2_deepspeed.py
+++ b/examples/research_projects/wav2vec2/test_wav2vec2_deepspeed.py
@@ -0,0 +1,185 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
+# hack it in for now:
+import sys
+from pathlib import Path
+
+
+git_repo_path = Path(__file__).resolve().parents[3] / "src"
+sys.path.insert(1, str(git_repo_path))
+
+import dataclasses  # noqa
+import io  # noqa
+import json  # noqa
+import os  # noqa
+import unittest  # noqa
+from copy import deepcopy  # noqa
+
+from parameterized import parameterized  # noqa
+from transformers import TrainingArguments, is_torch_available  # noqa
+from transformers.deepspeed import is_deepspeed_available  # noqa
+from transformers.file_utils import WEIGHTS_NAME  # noqa
+from transformers.testing_utils import (  # noqa
+    CaptureLogger,
+    ExtendSysPath,
+    TestCasePlus,
+    execute_subprocess_async,
+    get_gpu_count,
+    mockenv_context,
+    require_deepspeed,
+    require_torch_gpu,
+    require_torch_multi_gpu,
+    slow,
+)
+from transformers.trainer_utils import set_seed  # noqa
+
+
+set_seed(42)
+
+WAV2VEC2_TINY = "patrickvonplaten/wav2vec2_tiny_random_robust"
+
+
+ZERO2 = "zero2"
+ZERO3 = "zero3"
+stages = [ZERO2, ZERO3]
+
+
+@slow
+@require_deepspeed
+@require_torch_gpu
+class TestDeepSpeedWav2Vec2(TestCasePlus):
+    @parameterized.expand(stages)
+    def test_fp32_non_distributed(self, stage):
+        self.run_and_check(
+            stage=stage,
+            distributed=False,
+            fp16=False,
+        )
+
+    @require_torch_multi_gpu
+    @parameterized.expand(stages)
+    def test_fp32_distributed(self, stage):
+        self.run_and_check(
+            stage=stage,
+            distributed=True,
+            fp16=False,
+        )
+
+    @parameterized.expand(stages)
+    def test_fp16_non_distributed(self, stage):
+        self.run_and_check(
+            stage=stage,
+            distributed=False,
+            fp16=True,
+        )
+
+    @require_torch_multi_gpu
+    @parameterized.expand(stages)
+    def test_fp16_distributed(self, stage):
+        self.run_and_check(
+            stage=stage,
+            distributed=True,
+            fp16=True,
+        )
+
+    def do_checks(self, output_dir):
+        # XXX: run_asr is premature and doesn't save any results
+        # so all we check for now is that the process didn't fail
+        pass
+
+    # XXX: need to do better validation beyond just that the run was successful
+    def run_and_check(
+        self,
+        stage,
+        model_name: str = WAV2VEC2_TINY,
+        eval_steps: int = 10,
+        distributed: bool = True,
+        quality_checks: bool = True,
+        fp16: bool = True,
+    ):
+
+        output_dir = self.run_trainer(
+            stage=stage,
+            model_name=model_name,
+            eval_steps=eval_steps,
+            num_train_epochs=1,
+            distributed=distributed,
+            fp16=fp16,
+        )
+
+        self.do_checks(output_dir)
+
+        return output_dir
+
+    def run_trainer(
+        self,
+        stage: str,
+        model_name: str,
+        eval_steps: int = 10,
+        num_train_epochs: int = 1,
+        distributed: bool = True,
+        fp16: bool = True,
+    ):
+
+        output_dir = self.get_auto_remove_tmp_dir("./xxx", after=False)
+        args = f"""
+            --model_name_or_path {model_name}
+            --dataset_name patrickvonplaten/librispeech_asr_dummy
+            --dataset_config_name clean
+            --train_split_name validation
+            --validation_split_name validation
+            --output_dir {output_dir}
+            --num_train_epochs {str(num_train_epochs)}
+            --per_device_train_batch_size 2
+            --per_device_eval_batch_size 2
+            --evaluation_strategy steps
+            --learning_rate 5e-4
+            --warmup_steps 8
+            --orthography timit
+            --preprocessing_num_workers 1
+            --group_by_length
+            --freeze_feature_extractor
+            --report_to none
+            --logging_steps 0
+            --save_steps 0
+            --eval_steps {eval_steps}
+            --report_to none
+        """.split()
+
+        if fp16:
+            args.extend(["--fp16"])
+
+        # currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
+        # hence the separate config files
+        ds_args = f"--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json".split()
+        script = [f"{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py"]
+        launcher = self.get_launcher(distributed)
+
+        cmd = launcher + script + args + ds_args
+        # keep for quick debug
+        # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
+        execute_subprocess_async(cmd, env=self.get_env())
+
+        return output_dir
+
+    def get_launcher(self, distributed=False):
+        # 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
+        # - it won't be able to handle that
+        # 2. for now testing with just 2 gpus max (since some quality tests may give different
+        # results with mode gpus because we use very little data)
+        num_gpus = min(2, get_gpu_count()) if distributed else 1
+        return f"deepspeed --num_nodes 1 --num_gpus {num_gpus}".split()