feat: add convenient selection of data parallelism (#335)

* feat: add convenient selection of data parallelism * fix spelling Co-authored-by: Aaron Berdy <[email protected]> * Rename distribution to 'data_parallel' Co-authored-by: Aaron Berdy <[email protected]>
amazon-braket · Apr 26, 2022 · 4ed61b3 · 4ed61b3
1 parent bbeab35
commit 4ed61b3
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 0 deletions.
diff --git a/src/braket/aws/aws_quantum_job.py b/src/braket/aws/aws_quantum_job.py
@@ -73,6 +73,7 @@ def create(
         hyperparameters: Dict[str, Any] = None,
         input_data: Union[str, Dict, S3DataSourceConfig] = None,
         instance_config: InstanceConfig = None,
+        distribution: str = None,
         stopping_condition: StoppingCondition = None,
         output_data_config: OutputDataConfig = None,
         copy_checkpoints_from_job: str = None,
@@ -134,6 +135,10 @@ def create(
                 to execute the job. Default: InstanceConfig(instanceType='ml.m5.large',
                 instanceCount=1, volumeSizeInGB=30).
 
+            distribution (str): A str that specifies how the job should be distributed. If set to
+                "data_parallel", the hyperparameters for the job will be set to use data parallelism
+                features for PyTorch or TensorFlow. Default: None.
+
             stopping_condition (StoppingCondition): The maximum length of time, in seconds,
                 and the maximum number of tasks that a job can run before being forcefully stopped.
                 Default: StoppingCondition(maxRuntimeInSeconds=5 * 24 * 60 * 60).
@@ -181,6 +186,7 @@ def create(
             hyperparameters=hyperparameters,
             input_data=input_data,
             instance_config=instance_config,
+            distribution=distribution,
             stopping_condition=stopping_condition,
             output_data_config=output_data_config,
             copy_checkpoints_from_job=copy_checkpoints_from_job,

diff --git a/src/braket/jobs/quantum_job_creation.py b/src/braket/jobs/quantum_job_creation.py
@@ -44,6 +44,7 @@ def prepare_quantum_job(
     hyperparameters: Dict[str, Any] = None,
     input_data: Union[str, Dict, S3DataSourceConfig] = None,
     instance_config: InstanceConfig = None,
+    distribution: str = None,
     stopping_condition: StoppingCondition = None,
     output_data_config: OutputDataConfig = None,
     copy_checkpoints_from_job: str = None,
@@ -98,6 +99,10 @@ def prepare_quantum_job(
             to execute the job. Default: InstanceConfig(instanceType='ml.m5.large',
             instanceCount=1, volumeSizeInGB=30, volumeKmsKey=None).
 
+        distribution (str): A str that specifies how the job should be distributed. If set to
+            "data_parallel", the hyperparameters for the job will be set to use data parallelism
+            features for PyTorch or TensorFlow. Default: None.
+
         stopping_condition (StoppingCondition): The maximum length of time, in seconds,
             and the maximum number of tasks that a job can run before being forcefully stopped.
             Default: StoppingCondition(maxRuntimeInSeconds=5 * 24 * 60 * 60).
@@ -191,6 +196,12 @@ def prepare_quantum_job(
             "s3Uri"
         ]
         aws_session.copy_s3_directory(checkpoints_to_copy, checkpoint_config.s3Uri)
+    if distribution == "data_parallel":
+        distributed_hyperparams = {
+            "sagemaker_distributed_dataparallel_enabled": "true",
+            "sagemaker_instance_type": instance_config.instanceType,
+        }
+        hyperparameters.update(distributed_hyperparams)
 
     create_job_kwargs = {
         "jobName": job_name,

diff --git a/test/unit_tests/braket/aws/test_aws_quantum_job.py b/test/unit_tests/braket/aws/test_aws_quantum_job.py
@@ -506,6 +506,7 @@ def prepare_job_args(aws_session, device_arn):
         "hyperparameters": Mock(),
         "input_data": Mock(),
         "instance_config": Mock(),
+        "distribution": Mock(),
         "stopping_condition": Mock(),
         "output_data_config": Mock(),
         "copy_checkpoints_from_job": Mock(),