feature: add data parallelism support (#454) (#511) (#495)

Author: ChaiBapchya

src/sagemaker/debugger/__init__.py

@@ -29,7 +29,7 @@
 from sagemaker.debugger.metrics_config import (  # noqa: F401
     DataloaderProfilingConfig,
     DetailedProfilingConfig,
-    HerringProfilingConfig,
+    SMDataParallelProfilingConfig,
     HorovodProfilingConfig,
     PythonProfilingConfig,
 )

src/sagemaker/debugger/framework_profile.py

@@ -16,7 +16,7 @@
 from sagemaker.debugger.metrics_config import (
     DetailedProfilingConfig,
     DataloaderProfilingConfig,
-    HerringProfilingConfig,
+    SMDataParallelProfilingConfig,
     HorovodProfilingConfig,
     PythonProfilingConfig,
 )
@@ -33,7 +33,7 @@
     DataloaderProfilingConfig,
     PythonProfilingConfig,
     HorovodProfilingConfig,
-    HerringProfilingConfig,
+    SMDataParallelProfilingConfig,
 ]
 
 
@@ -53,7 +53,7 @@ def __init__(
         dataloader_profiling_config=None,
         python_profiling_config=None,
         horovod_profiling_config=None,
-        herring_profiling_config=None,
+        smdataparallel_profiling_config=None,
         start_step=None,
         num_steps=None,
         start_unix_time=None,
@@ -88,8 +88,8 @@ def __init__(
                 collected by the Python profiler (cProfile or Pyinstrument).
             horovod_profiling_config (HorovodProfilingConfig): The configuration for metrics
                 collected by horovod when using horovod for distributed training.
-            herring_profiling_config (HerringProfilingConfig): The configuration for metrics
-                collected by herring when using herring for distributed training.
+            smdataparallel_profiling_config (SMDataParallelProfilingConfig): The configuration for
+                metrics collected by SageMaker Distributed training.
             start_step (int): The step at which to start profiling.
             num_steps (int): The number of steps to profile.
             start_unix_time (int): The UNIX time at which to start profiling.
@@ -108,7 +108,7 @@ def __init__(
             dataloader_profiling_config,
             python_profiling_config,
             horovod_profiling_config,
-            herring_profiling_config,
+            smdataparallel_profiling_config,
         )
 
         use_one_config_for_all_metrics = (

src/sagemaker/debugger/metrics_config.py

@@ -18,8 +18,8 @@
     DATALOADER_PROFILING_START_STEP_DEFAULT,
     DETAILED_PROFILING_CONFIG_NAME,
     DETAILED_PROFILING_START_STEP_DEFAULT,
-    HERRING_PROFILING_CONFIG_NAME,
-    HERRING_PROFILING_START_STEP_DEFAULT,
+    SMDATAPARALLEL_PROFILING_CONFIG_NAME,
+    SMDATAPARALLEL_PROFILING_START_STEP_DEFAULT,
     HOROVOD_PROFILING_CONFIG_NAME,
     HOROVOD_PROFILING_START_STEP_DEFAULT,
     PROFILING_NUM_STEPS_DEFAULT,
@@ -298,6 +298,9 @@ def __init__(
             start_step = PYTHON_PROFILING_START_STEP_DEFAULT
             num_steps = PYTHON_PROFILING_NUM_STEPS_DEFAULT
 
+        if profile_default_steps:
+            cprofile_timer = cProfileTimer.DEFAULT
+
         super().__init__(
             PYTHON_PROFILING_CONFIG_NAME, start_step, num_steps, start_unix_time, duration
         )
@@ -367,8 +370,8 @@ def __init__(
         )
 
 
-class HerringProfilingConfig(MetricsConfigBase):
-    """Configuration for metrics collected by herring when using herring for distributed training.
+class SMDataParallelProfilingConfig(MetricsConfigBase):
+    """Configuration for metrics collected by SageMaker Distributed training.
 
     By default, profile step 15 of training.
     """
@@ -381,9 +384,9 @@ def __init__(
         duration=None,
         profile_default_steps=False,
     ):
-        """If profile_default_steps is set to True or none of the range fields are specified,
-        use the default config for herring profiling. Otherwise, profile according to the
-        specified range fields.
+        """If profile_default_steps is set to True or none of the range fields are specified, use
+        the default profiling config for SageMaker Distributed training. Otherwise, profile
+        according to the specified range fields.
 
         Args:
             start_step (int): The step at which to start profiling.
@@ -396,9 +399,9 @@ def __init__(
             profile_default_steps, bool
         ), ErrorMessages.INVALID_PROFILE_DEFAULT_STEPS.value
         if profile_default_steps or start_step is num_steps is start_unix_time is duration is None:
-            start_step = HERRING_PROFILING_START_STEP_DEFAULT
+            start_step = SMDATAPARALLEL_PROFILING_START_STEP_DEFAULT
             num_steps = PROFILING_NUM_STEPS_DEFAULT
 
         super().__init__(
-            HERRING_PROFILING_CONFIG_NAME, start_step, num_steps, start_unix_time, duration
+            SMDATAPARALLEL_PROFILING_CONFIG_NAME, start_step, num_steps, start_unix_time, duration
         )

src/sagemaker/debugger/profiler_constants.py

@@ -22,13 +22,13 @@
 DATALOADER_PROFILING_CONFIG_NAME = "DataloaderProfilingConfig"
 PYTHON_PROFILING_CONFIG_NAME = "PythonProfilingConfig"
 HOROVOD_PROFILING_CONFIG_NAME = "HorovodProfilingConfig"
-HERRING_PROFILING_CONFIG_NAME = "HerringProfilingConfig"
+SMDATAPARALLEL_PROFILING_CONFIG_NAME = "SMDataParallelProfilingConfig"
 
 DETAILED_PROFILING_START_STEP_DEFAULT = 5
 DATALOADER_PROFILING_START_STEP_DEFAULT = 7
 PYTHON_PROFILING_START_STEP_DEFAULT = 9
 HOROVOD_PROFILING_START_STEP_DEFAULT = 13
-HERRING_PROFILING_START_STEP_DEFAULT = 15
+SMDATAPARALLEL_PROFILING_START_STEP_DEFAULT = 15
 PROFILING_NUM_STEPS_DEFAULT = 1
 START_STEP_DEFAULT = 0
 PYTHON_PROFILING_NUM_STEPS_DEFAULT = 3

src/sagemaker/debugger/utils.py

@@ -118,3 +118,4 @@ class cProfileTimer(Enum):
     TOTAL_TIME = "total_time"
     CPU_TIME = "cpu_time"
     OFF_CPU_TIME = "off_cpu_time"
+    DEFAULT = "default"

src/sagemaker/estimator.py

@@ -1903,6 +1903,8 @@ class Framework(EstimatorBase):
 
     LAUNCH_PS_ENV_NAME = "sagemaker_parameter_server_enabled"
     LAUNCH_MPI_ENV_NAME = "sagemaker_mpi_enabled"
+    LAUNCH_SM_DDP_ENV_NAME = "sagemaker_distributed_dataparallel_enabled"
+    INSTANCE_TYPE = "sagemaker_instance_type"
     MPI_NUM_PROCESSES_PER_HOST = "sagemaker_mpi_num_of_processes_per_host"
     MPI_CUSTOM_MPI_OPTIONS = "sagemaker_mpi_custom_mpi_options"
     CONTAINER_CODE_CHANNEL_SOURCEDIR_PATH = "/opt/ml/input/data/code/sourcedir.tar.gz"

src/sagemaker/fw_utils.py

@@ -50,6 +50,12 @@
 
 DEBUGGER_UNSUPPORTED_REGIONS = ("us-iso-east-1",)
 SINGLE_GPU_INSTANCE_TYPES = ("ml.p2.xlarge", "ml.p3.2xlarge")
+SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES = ("ml.p3.16xlarge", "ml.p3dn.24xlarge", "local_gpu")
+SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS = {
+    "tensorflow": ["2.3.0", "2.3.1"],
+    "pytorch": ["1.6.0"],
+}
+SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]
 
 
 def validate_source_dir(script, directory):
@@ -255,9 +261,8 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
             .. code:: python
 
                 {
-                    'parameter_server':
-                    {
-                        'enabled': True
+                    "parameter_server": {
+                        "enabled": True
                     }
                 }
 
@@ -279,6 +284,154 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
         logger.warning(PARAMETER_SERVER_MULTI_GPU_WARNING)
 
 
+def validate_smdistributed(
+    instance_type, framework_name, framework_version, py_version, distribution, image_uri=None
+):
+    """Check if smdistributed strategy is correctly invoked by the user.
+
+    Currently, two strategies are supported: `dataparallel` or `modelparallel`.
+    Validate if the user requested strategy is supported.
+
+    Currently, only one strategy can be specified at a time. Validate if the user has requested
+    more than one strategy simultaneously.
+
+    Validate if the smdistributed dict arg is syntactically correct.
+
+    Additionally, perform strategy-specific validations.
+
+    Args:
+        instance_type (str): A string representing the type of training instance selected.
+        framework_name (str): A string representing the name of framework selected.
+        framework_version (str): A string representing the framework version selected.
+        py_version (str): A string representing the python version selected.
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) For example:
+
+            .. code:: python
+
+                {
+                    "smdistributed": {
+                        "dataparallel": {
+                            "enabled": True
+                        }
+                    }
+                }
+        image_uri (str): A string representing a Docker image URI.
+
+    Raises:
+        ValueError: if distribution dictionary isn't correctly formatted or
+            multiple strategies are requested simultaneously or
+            an unsupported strategy is requested or
+            strategy-specific inputs are incorrect/unsupported
+    """
+    if "smdistributed" not in distribution:
+        # Distribution strategy other than smdistributed is selected
+        return
+
+    # distribution contains smdistributed
+    smdistributed = distribution["smdistributed"]
+    if not isinstance(smdistributed, dict):
+        raise ValueError("smdistributed strategy requires a dictionary")
+
+    if len(smdistributed) > 1:
+        # more than 1 smdistributed strategy requested by the user
+        err_msg = (
+            "Cannot use more than 1 smdistributed strategy. \n"
+            "Choose one of the following supported strategies:"
+            f"{SMDISTRIBUTED_SUPPORTED_STRATEGIES}"
+        )
+        raise ValueError(err_msg)
+
+    # validate if smdistributed strategy is supported
+    # currently this for loop essentially checks for only 1 key
+    for strategy in smdistributed:
+        if strategy not in SMDISTRIBUTED_SUPPORTED_STRATEGIES:
+            err_msg = (
+                f"Invalid smdistributed strategy provided: {strategy} \n"
+                f"Supported strategies: {SMDISTRIBUTED_SUPPORTED_STRATEGIES}"
+            )
+            raise ValueError(err_msg)
+
+    # smdataparallel-specific input validation
+    if "dataparallel" in smdistributed:
+        _validate_smdataparallel_args(
+            instance_type, framework_name, framework_version, py_version, distribution, image_uri
+        )
+
+
+def _validate_smdataparallel_args(
+    instance_type, framework_name, framework_version, py_version, distribution, image_uri=None
+):
+    """Check if request is using unsupported arguments.
+
+    Validate if user specifies a supported instance type, framework version, and python
+    version.
+
+    Args:
+        instance_type (str): A string representing the type of training instance selected. Ex: `ml.p3.16xlarge`
+        framework_name (str): A string representing the name of framework selected. Ex: `tensorflow`
+        framework_version (str): A string representing the framework version selected. Ex: `2.3.1`
+        py_version (str): A string representing the python version selected. Ex: `py3`
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) Ex:
+
+            .. code:: python
+
+                {
+                    "smdistributed": {
+                        "dataparallel": {
+                            "enabled": True
+                        }
+                    }
+                }
+        image_uri (str): A string representing a Docker image URI.
+
+    Raises:
+        ValueError: if
+            (`instance_type` is not in SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES or
+            `py_version` is not python3 or
+            `framework_version` is not in SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSION
+    """
+    smdataparallel_enabled = (
+        distribution.get("smdistributed").get("dataparallel").get("enabled", False)
+    )
+
+    if not smdataparallel_enabled:
+        return
+
+    is_instance_type_supported = instance_type in SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES
+
+    err_msg = ""
+
+    if not is_instance_type_supported:
+        # instance_type is required
+        err_msg += (
+            f"Provided instance_type {instance_type} is not supported by smdataparallel.\n"
+            "Please specify one of the supported instance types:"
+            f"{SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES}\n"
+        )
+
+    if not image_uri:
+        # ignore framework_version & py_version if image_uri is set
+        # in case image_uri is not set, then both are mandatory
+        supported = SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS[framework_name]
+        if framework_version not in supported:
+            err_msg += (
+                f"Provided framework_version {framework_version} is not supported by"
+                " smdataparallel.\n"
+                f"Please specify one of the supported framework versions: {supported} \n"
+            )
+
+        if "py3" not in py_version:
+            err_msg += (
+                f"Provided py_version {py_version} is not supported by smdataparallel.\n"
+                "Please specify py_version=py3"
+            )
+
+    if err_msg:
+        raise ValueError(err_msg)
+
+
 def python_deprecation_warning(framework, latest_supported_version):
     """
     Args: