Feature: Support new distribution mechanism for PT-XLA (#241)

* Support new distribution mechanism for PT-XLA

* test: Adding test to check new PT-XLA distribution mechanism

Author: Lokiiiiii

src/sagemaker_pytorch_container/training.py

@@ -23,6 +23,7 @@
 LAUNCH_SMDATAPARALLEL_ENV_NAME = 'sagemaker_distributed_dataparallel_enabled'
 LAUNCH_MPI_ENV_NAME = 'sagemaker_mpi_enabled'
 LAUNCH_PYTORCH_DDP_ENV_NAME = "sagemaker_pytorch_ddp_enabled"
+LAUNCH_PYTORCH_XLA_ENV_NAME = "sagemaker_pytorch_xla_multi_worker_enabled"
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +61,10 @@ def train(training_environment):
     smdataparallel_enabled = training_environment.additional_framework_parameters.get(
         LAUNCH_SMDATAPARALLEL_ENV_NAME, False
     )
+
+    pytorch_xla_enabled = training_environment.additional_framework_parameters.get(
+        LAUNCH_PYTORCH_XLA_ENV_NAME, False
+    )
     # default scenario
     runner_type = runner.ProcessRunnerType
 
@@ -72,6 +77,9 @@ def train(training_environment):
         elif smdataparallel_enabled:
             runner_type = runner.SMDataParallelRunnerType
             logger.info('Invoking SMDataParallel')
+        elif pytorch_xla_enabled:
+            runner_type = runner.PyTorchXLARunnerType
+            logger.info('Invoking PT-XLA Runner')
     logger.info('Invoking user training script.')
     try:
         entry_point.run(uri=training_environment.module_dir,

test/container/1.10.2/Dockerfile.dlc.trcomp

@@ -0,0 +1,6 @@
+ARG region
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/huggingface-pytorch-trcomp-training:1.10.2-transformers4.17.0-gpu-py38-cu113-ubuntu20.04
+
+COPY dist/sagemaker_pytorch_training-*.tar.gz /sagemaker_pytorch_training.tar.gz
+RUN pip install --upgrade --no-cache-dir /sagemaker_pytorch_training.tar.gz && \
+    rm /sagemaker_pytorch_training.tar.gz

test/unit/test_train.py

@@ -23,7 +23,7 @@
 from mock import MagicMock, PropertyMock
 from mock import patch
 
-from sagemaker_pytorch_container.training import main, train, _dns_lookup, MASTER_PORT
+from sagemaker_pytorch_container.training import main, train, _dns_lookup, LAUNCH_PYTORCH_XLA_ENV_NAME, MASTER_PORT
 
 
 @pytest.fixture(name='training_env')
@@ -106,6 +106,22 @@ def test_train_pytorch_ddp(run_module, training_env):
     )
 
 
+@patch("sagemaker_training.entry_point.run")
+@patch('socket.gethostbyname', MagicMock())
+def test_train_pytorch_xla_distributed(run_module, training_env):
+    training_env.additional_framework_parameters[LAUNCH_PYTORCH_XLA_ENV_NAME] = True
+
+    train(training_env)
+    run_module.assert_called_with(
+        uri=training_env.module_dir,
+        user_entry_point=training_env.user_entry_point,
+        args=training_env.to_cmd_args(),
+        env_vars=training_env.to_env_vars(),
+        capture_error=True,
+        runner_type=runner.PyTorchXLARunnerType,
+    )
+
+
 @patch('sagemaker_training.entry_point.run', MagicMock())
 @patch('socket.gethostbyname', MagicMock())
 def test_environment(training_env):