[Train] Add PyTorch local mode support for multi-process training with torchrun (#56218)

This PR extends the Ray Train v2 local mode support (from #55487) to
enable users to launch multiple local mode processes using torchrun for
PyTorch distributed training. **With this new feature, users can easily
switch between torchrun and Ray Train without modifying their training
code.**

<img width="1249" height="811" alt="image"
src="https://github.com/user-attachments/assets/5d998b5e-8f58-425a-b535-d4f4d0b64a5c"
/>

### Note

Ray data on multiple processes is not supported. Might need to wait for
#55114 or similar components.

## Key Changes

### Multi-Process Local Mode Support
- **`LocalTorchController`**: New controller that detects torchrun env
variables and sets contexts accordingly
- **Torchrun Integration**: Users can now launch multiple local mode
processes using `torchrun` command
- **Environment Detection**: Automatically detects torchrun environment
variables and initializes distributed training

## Usage Example

```python
import os
import tempfile

import torch
from torch.nn import CrossEntropyLoss
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose

import ray
from ray.train import Checkpoint, CheckpointConfig, RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer
from ray.train.v2.api.config import FailureConfig
import ray.train.torch

def train_func():
    # Model, Loss, Optimizer
    model = resnet18(num_classes=10)
    model.conv1 = torch.nn.Conv2d(
        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    )
    # [1] Prepare model.
    model = ray.train.torch.prepare_model(model)
    criterion = CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=0.001)

    # Data
    transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
    data_dir = os.path.join(tempfile.gettempdir(), "data")
    train_data = FashionMNIST(root=data_dir, train=True, download=True, transform=transform)
    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
    # [2] Prepare dataloader.
    train_loader = ray.train.torch.prepare_data_loader(train_loader)

    # Training
    for epoch in range(10):
        if ray.train.get_context().get_world_size() > 1:
            train_loader.sampler.set_epoch(epoch)

        for images, labels in train_loader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # [3] Report metrics and checkpoint.
        metrics = {"loss": loss.item(), "epoch": epoch}
        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            torch.save(
                model.state_dict(),
                os.path.join(temp_checkpoint_dir, "model.pt")
            )
            ray.train.report(
                metrics,
                checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
            )
        if ray.train.get_context().get_world_rank() == 0:
            print(metrics)

# Configuration for local mode
use_gpu = True
scaling_config = ScalingConfig(num_workers=0, use_gpu=use_gpu)  # Local mode
run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=1))

# Note: Ray Data not supported with multiple processes in local mode
# For multi-process training, use PyTorch DataLoader as shown above

# Initialize the Trainer
trainer = TorchTrainer(
    train_loop_per_worker=train_func,
    scaling_config=scaling_config,
    run_config=run_config,
)

# Train the model
result = trainer.fit()
```

### Running Options:

```bash
# Option 1: Single process local mode
RAY_TRAIN_V2_ENABLED=1 python test.py

# Option 2: Multi-process local mode with torchrun
RAY_TRAIN_V2_ENABLED=1 torchrun --standalone --nnodes=1 --nproc-per-node=4 test.py

# Option 3: Switch to distributed Ray Train (change num_workers=4)
# Same training code works across all modes!
```

---------

Signed-off-by: xgui <xgui@anyscale.com>
Signed-off-by: Xinyuan <43737116+xinyuangui2@users.noreply.github.com>
Co-authored-by: matthewdeng <matthew.j.deng@gmail.com>

Author: xinyuangui2

python/ray/train/v2/_internal/execution/local_mode/__init__.py

@@ -0,0 +1,92 @@
+import logging
+import os
+from typing import Callable
+
+import torch
+import torch.distributed as dist
+
+from ray.train import Result
+from ray.train.v2._internal.execution.local_mode.utils import LocalController
+from ray.train.v2._internal.execution.train_fn_utils import (
+    LocalTrainFnUtils,
+    get_train_fn_utils,
+    set_train_fn_utils,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def has_torchrun_env() -> bool:
+    """Return True if this process has torch.distributed env vars set.
+
+    For torch.distributed.init_process_group with init_method="env://", these variables are required:
+    - RANK: The rank of the current process
+    - LOCAL_RANK: The local rank of the current process
+    - WORLD_SIZE: Total number of processes participating in the job
+    - LOCAL_WORLD_SIZE: Total number of processes participating in the job on the current node
+    - MASTER_ADDR: The IP address or hostname of the master node (rank 0)
+    - MASTER_PORT: A free port on the master node for communication
+
+    """
+    torch_dist_required_vars = {
+        "RANK",
+        "LOCAL_RANK",
+        "WORLD_SIZE",
+        "LOCAL_WORLD_SIZE",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+    }
+
+    return torch_dist_required_vars.issubset(os.environ.keys())
+
+
+class LocalTorchController(LocalController):
+    def _set_train_fn_utils(self) -> None:
+        world_size = 1
+        global_rank = 0
+        local_rank = 0
+        nproc_per_node = 1
+        node_rank = 0
+        if has_torchrun_env():
+            assert not dist.is_initialized(), "torch.distributed is already initialized"
+            torch.distributed.init_process_group(
+                backend="nccl" if torch.cuda.is_available() else "gloo"
+            )
+            world_size = torch.distributed.get_world_size()
+            global_rank = torch.distributed.get_rank()
+            local_rank = int(os.environ["LOCAL_RANK"])
+            if torch.cuda.is_available():
+                torch.cuda.set_device(local_rank)
+            nproc_per_node = int(os.environ.get("LOCAL_WORLD_SIZE"))
+            node_rank = global_rank // nproc_per_node
+
+        if world_size != 1:
+            assert (
+                self.datasets is None or len(self.datasets) == 0
+            ), "Ray Data is not supported in local mode with multiple workers."
+        set_train_fn_utils(
+            LocalTrainFnUtils(
+                experiment_name=self.experiment_name,
+                world_size=world_size,
+                world_rank=global_rank,
+                local_rank=local_rank,
+                local_world_size=nproc_per_node,
+                node_rank=node_rank,
+                dataset_shards=self.datasets,
+            )
+        )
+
+    def run(self, train_func: Callable[[], None]) -> Result:
+        self._set_train_fn_utils()
+        train_func()
+        train_fn_utils = get_train_fn_utils()
+        assert isinstance(train_fn_utils, LocalTrainFnUtils)
+        result = Result(
+            metrics=train_fn_utils._get_last_metrics(),
+            checkpoint=train_fn_utils.get_checkpoint(),
+            path=None,
+            error=None,
+        )
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        return result

python/ray/train/v2/_internal/execution/local_mode/torch.py

@@ -166,9 +166,19 @@ def __init__(
         self,
         experiment_name: str,
         dataset_shards: Optional[Dict[str, DataIterator]] = None,
+        world_size: int = 1,
+        world_rank: int = 0,
+        local_rank: int = 0,
+        local_world_size: int = 1,
+        node_rank: int = 0,
     ):
         self._context = LocalTrainContext(
             experiment_name=experiment_name,
+            world_size=world_size,
+            world_rank=world_rank,
+            local_rank=local_rank,
+            local_world_size=local_world_size,
+            node_rank=node_rank,
         )
         self._dataset_shards = dataset_shards
         self._last_metrics = None

…/_internal/execution/local_mode_utils.py …/_internal/execution/local_mode/utils.pypython/ray/train/v2/_internal/execution/local_mode_utils.py renamed to python/ray/train/v2/_internal/execution/local_mode/utils.py python/ray/train/v2/_internal/execution/local_mode_utils.py renamed to python/ray/train/v2/_internal/execution/local_mode/utils.py

@@ -253,27 +253,36 @@ class LocalTrainContext(TrainContext):
     def __init__(
         self,
         experiment_name: str,
+        world_size: int = 1,
+        world_rank: int = 0,
+        local_rank: int = 0,
+        local_world_size: int = 1,
+        node_rank: int = 0,
     ):
         self.experiment_name = experiment_name
+        self.world_size = world_size
+        self.world_rank = world_rank
+        self.local_rank = local_rank
+        self.local_world_size = local_world_size
+        self.node_rank = node_rank
 
     def get_experiment_name(self) -> str:
         return self.experiment_name
 
     def get_world_size(self) -> int:
-        return 1
+        return self.world_size
 
     def get_world_rank(self) -> int:
-        return 0
+        return self.world_rank
 
     def get_local_rank(self) -> int:
-        return 0
+        return self.local_rank
 
     def get_local_world_size(self) -> int:
-        return 1
+        return self.local_world_size
 
     def get_node_rank(self) -> int:
-        """For local mode, we only use one node."""
-        return 0
+        return self.node_rank
 
     def get_storage(self):
         raise NotImplementedError("Local storage context not yet implemented. ")

python/ray/train/v2/_internal/execution/train_fn_utils.py

@@ -46,7 +46,7 @@
 from ray.train.v2._internal.execution.context import TrainRunContext
 from ray.train.v2._internal.execution.controller import TrainController
 from ray.train.v2._internal.execution.failure_handling import create_failure_policy
-from ray.train.v2._internal.execution.local_mode_utils import LocalController
+from ray.train.v2._internal.execution.local_mode.utils import LocalController
 from ray.train.v2._internal.execution.scaling_policy import create_scaling_policy
 from ray.train.v2._internal.util import ObjectRefWrapper, construct_train_func
 from ray.train.v2.api.callback import UserCallback

python/ray/train/v2/api/context.py

@@ -1,6 +1,7 @@
 import math
+import os
 import sys
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import lightgbm
 import pandas as pd
@@ -38,6 +39,8 @@
 from ray.train.tests.lightning_test_utils import DummyDataModule, LinearModule
 from ray.train.tests.util import create_dict_checkpoint
 from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.execution.local_mode.torch import LocalTorchController
+from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
 from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
 from ray.train.v2.jax import JaxTrainer
 from ray.train.xgboost import (
@@ -522,5 +525,94 @@ def xgboost_train_fn_per_worker():
         XGBoostTrainer.get_model(result.checkpoint)
 
 
+def test_torch_distributed_variables_local_train_fn_utils():
+    """Test that torch distributed variables are correctly used to create LocalTrainFnUtils."""
+
+    # Test scenario 1: Without torch distributed environment variables
+    with patch.dict(os.environ, {}, clear=True):
+        controller = LocalTorchController("test_experiment")
+
+        def dummy_train_func():
+            train_fn_utils = get_train_fn_utils()
+            # Verify default values when no torch distributed env vars are set
+            context = train_fn_utils.get_context()
+            assert context.get_world_size() == 1
+            assert context.get_world_rank() == 0
+            assert context.get_local_rank() == 0
+            assert context.get_local_world_size() == 1
+            assert context.get_node_rank() == 0
+
+        controller.run(dummy_train_func)
+
+    # Test scenario 2: With torch distributed environment variables (CPU)
+    torch_env_vars = {
+        "RANK": "2",
+        "LOCAL_RANK": "1",
+        "WORLD_SIZE": "4",
+        "LOCAL_WORLD_SIZE": "2",
+        "MASTER_ADDR": "127.0.0.1",
+        "MASTER_PORT": "29500",
+    }
+
+    with patch.dict(os.environ, torch_env_vars, clear=True), patch(
+        "torch.distributed.is_initialized", return_value=False
+    ), patch("torch.distributed.get_world_size", return_value=4), patch(
+        "torch.distributed.get_rank", return_value=2
+    ), patch(
+        "torch.cuda.is_available", return_value=False
+    ), patch(
+        "torch.distributed.init_process_group"
+    ) as mock_init_pg:
+
+        controller = LocalTorchController("test_experiment")
+
+        def dummy_train_func():
+            train_fn_utils = get_train_fn_utils()
+            # Verify torch distributed values are correctly passed
+            context = train_fn_utils.get_context()
+            assert context.get_world_size() == 4
+            assert context.get_world_rank() == 2
+            assert context.get_local_rank() == 1
+            assert context.get_local_world_size() == 2
+            assert (
+                context.get_node_rank() == 1
+            )  # global_rank // nproc_per_node = 2 // 2 = 1
+
+        controller.run(dummy_train_func)
+
+        # Verify torch.distributed methods were called with CPU backend
+        mock_init_pg.assert_called_once_with(backend="gloo")
+
+    # Test scenario 3: With torch distributed environment variables (GPU)
+    with patch.dict(os.environ, torch_env_vars, clear=True), patch(
+        "torch.distributed.is_initialized", return_value=False
+    ), patch("torch.distributed.get_world_size", return_value=4), patch(
+        "torch.distributed.get_rank", return_value=2
+    ), patch(
+        "torch.cuda.is_available", return_value=True
+    ), patch(
+        "torch.distributed.init_process_group"
+    ) as mock_init_pg, patch(
+        "torch.cuda.set_device"
+    ) as mock_set_device:
+
+        controller = LocalTorchController("test_experiment")
+
+        def dummy_train_func():
+            train_fn_utils = get_train_fn_utils()
+            # Verify torch distributed values are correctly passed
+            context = train_fn_utils.get_context()
+            assert context.get_world_size() == 4
+            assert context.get_world_rank() == 2
+            assert context.get_local_rank() == 1
+            assert context.get_local_world_size() == 2
+            assert context.get_node_rank() == 1
+
+        controller.run(dummy_train_func)
+
+        mock_init_pg.assert_called_once_with(backend="nccl")
+        mock_set_device.assert_called_once_with(1)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/v2/api/data_parallel_trainer.py

@@ -2,6 +2,7 @@
 
 from ray.train import Checkpoint, DataConfig
 from ray.train.trainer import GenDataset
+from ray.train.v2._internal.execution.local_mode.torch import LocalTorchController
 from ray.train.v2.api.config import RunConfig, ScalingConfig
 from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
 from ray.util import PublicAPI
@@ -213,3 +214,9 @@ def __init__(
             resume_from_checkpoint=resume_from_checkpoint,
             metadata=metadata,
         )
+
+    def _get_local_controller(self) -> LocalTorchController:
+        return LocalTorchController(
+            experiment_name=self.run_config.name,
+            datasets=self.datasets,
+        )