Quick start: Distributed training with DLC

3.1 Create a dataset

A dataset stores your training code, data, and model artifacts. This section explains how to use a dataset based on Alibaba Cloud Object Storage Service (OSS).

1. In the left-side navigation pane of the PAI console, choose Datasets > Custom Dataset > Create datasets.

2. Configure the following key parameters and keep the default values for the others.

   • Name: For example, enter dataset_mnist.

   • Storage Type: Select OSS.

   • OSS Path: Click the folder icon , select a bucket, and create a new directory, such as dlc_mnist.

     If you have not activated OSS or do not have a bucket available in the current region, follow these steps to activate OSS and create a bucket:

     (Optional) Activate OSS and create a bucket

     1. Activate OSS.

     2. Log on to the OSS console, click Create Bucket, enter a Bucket name, and select the same Region as PAI. Keep the default values for other parameters and click Create.

   Click Confirm to create the dataset.

3. Upload the training code and data.

   1. Download the provided training script: mnist_train.py. To simplify the process, the script automatically downloads the training data to the dataSet directory of your dataset during runtime.

      For your own projects, you can pre-upload your code and training data to a PAI dataset.

      Single-node, single-GPU training code example: mnist_train.py

      import torch
      import torch.nn as nn
      import torch.nn.functional as F
      import torch.optim as optim
      from torch.utils.data import DataLoader
      from torchvision import datasets, transforms
      from torch.utils.tensorboard import SummaryWriter
      # Hyperparameters
      batch_size = 64  # Batch size for each training iteration
      learning_rate = 0.01  # Learning rate
      num_epochs = 20  # Number of training epochs
      # Check if a GPU is available
      device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
      # Data preprocessing
      transform = transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize((0.5,), (0.5,))
      ])
      train_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=True, download=True, transform=transform)
      val_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=False, download=False, transform=transform)
      train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
      val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
      # Define a simple neural network
      class SimpleCNN(nn.Module):
          def __init__(self):
              super(SimpleCNN, self).__init__()
              # First convolutional layer: 1 input channel (grayscale image), 10 output channels, 5x5 kernel
              self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
              # Second convolutional layer: 10 input channels, 20 output channels, 3x3 kernel
              self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
              # Fully connected layer: Input is 20 * 5 * 5 (feature map size after convolution and pooling), output is 128
              self.fc1 = nn.Linear(20 * 5 * 5, 128)
              # Output layer: 128 -> 10 (for 10 digit classes)
              self.fc2 = nn.Linear(128, 10)
          def forward(self, x):
              # Input x shape: [batch, 1, 28, 28]
              x = F.max_pool2d(F.relu(self.conv1(x)), 2)  # [batch, 10, 12, 12]
              x = F.max_pool2d(F.relu(self.conv2(x)), 2)  # [batch, 20, 5, 5]
              x = x.view(-1, 20 * 5 * 5)  # Flatten to [batch, 500]
              x = F.relu(self.fc1(x))      # [batch, 128]
              x = self.fc2(x)              # [batch, 10]
              return x
      # Instantiate the model and move it to the GPU if available
      model = SimpleCNN().to(device)
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(model.parameters(), lr=learning_rate)
      # Create a TensorBoard SummaryWriter to visualize the model training process
      writer = SummaryWriter('/mnt/data/output/runs/mnist_experiment')
      # Variable to save the model with the highest accuracy
      best_val_accuracy = 0.0
      # Train the model and log loss and accuracy
      for epoch in range(num_epochs):
          model.train()
          for batch_idx, (data, target) in enumerate(train_loader):
              data, target = data.to(device), target.to(device)  # Move data and target to the GPU
              # Zero the gradients
              optimizer.zero_grad()
              # Forward pass
              output = model(data)
              # Calculate loss
              loss = criterion(output, target)
              # Backward pass
              loss.backward()
              # Update parameters
              optimizer.step()
              # Log training loss to TensorBoard
              if batch_idx % 100 == 0:  # Log every 100 batches
                  writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + batch_idx)
                  print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
          # Validate the model and log validation loss and accuracy
          model.eval()
          val_loss = 0
          correct = 0
          with torch.no_grad():  # Do not calculate gradients
              for data, target in val_loader:
                  data, target = data.to(device), target.to(device)  # Move data and target to the GPU
                  output = model(data)
                  val_loss += criterion(output, target).item()  # Accumulate validation loss
                  pred = output.argmax(dim=1, keepdim=True)  # Get predicted label
                  correct += pred.eq(target.view_as(pred)).sum().item()  # Accumulate number of correct predictions
          val_loss /= len(val_loader)  # Calculate average validation loss
          val_accuracy = 100. * correct / len(val_loader.dataset)  # Calculate validation accuracy
          print(f'Validation Loss: {val_loss:.4f}, Accuracy: {correct}/{len(val_loader.dataset)} ({val_accuracy:.0f}%)')
          # Log validation loss and accuracy to TensorBoard
          writer.add_scalar('Loss/validation', val_loss, epoch)
          writer.add_scalar('Accuracy/validation', val_accuracy, epoch)
          # Save the model with the highest validation accuracy
          if val_accuracy > best_val_accuracy:
              best_val_accuracy = val_accuracy
              torch.save(model.state_dict(), '/mnt/data/output/best_model.pth')
              print(f'Model saved with accuracy: {best_val_accuracy:.2f}%')
      # Close the SummaryWriter
      writer.close()
      print('Training complete. writer.close()')

   2. Upload the code. On the dataset details page, click View Data to go to the OSS console. Then, click Upload Object >Select Files > Upload Object to upload the training script to OSS.

3.2 Create a DLC job

1. In the left-side navigation pane of the PAI console, choose Deep Learning Containers (DLC) > Create Job.

2. Configure the following key parameters and keep the default values for the others. For a full list of parameters, see Create a training job.

   • Image Configuration: Select Image Address, and then enter the appropriate image address for your Region.

     In the top navigation bar of the console, use the Region selector to confirm your current region, such as China (Hangzhou).

     Region	Image address
     China (Beijing)	dsw-registry-vpc.cn-beijing.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
     China (Shanghai)	dsw-registry-vpc.cn-shanghai.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
     China (Hangzhou)	dsw-registry-vpc.cn-hangzhou.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
     Other	Find your region ID and replace <Region ID> in the image address to get the full URL: dsw-registry-vpc.<Region ID>.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04

     This image uses the same verified environment from the Get Started with DSW Interactive Modeling topic. When modeling with PAI, a typical workflow is to first verify the environment and develop code in DSW, and then use DLC for training.

   • Dataset Mount: Select Custom Dataset and choose the dataset that you created in the previous step. The default Mount Path is /mnt/data.

   • Startup Command: python /mnt/data/mnist_train.py

     The startup command is the same as when you run the code in DSW or locally. However, because mnist_train.py  is now mounted to /mnt/data/, you only need to change the code path to /mnt/data/mnist_train.py. 

   • Source: Select Public Resources, and then select ecs.gn7i-c8g1.2xlarge for the Resource Type.

     If this instance type is out of stock, you can choose another GPU instance type.

   Click OK to create the job. The job takes about 15 minutes to run. During execution, you can click Logs to view the training progress.

   The training log shows that the MNIST model reaches a validation accuracy of 9886/10000 (approximately 99%) at Epoch 18. When the job status changes to Succeeded, the log displays the final output: Training complete. writer.close().

   After the job completes, DLC saves the best model checkpoint and TensorBoard log to the output path of the mounted dataset.

   The dlc_mnist/ directory in your OSS bucket now contains a dataSet/ subdirectory for the training data, an output/ subdirectory for the model artifacts, and the mnist_train.py script.

3.3 (Optional) view in TensorBoard

You can use the TensorBoard visualization tool to view the loss curve and better understand the training process.

1. On the details page of the DLC job, click the Tensorboard tab, and then click Create TensorBoard.

2. Set Configuration Type to By Task. For Summary Path, enter the path where the SummaryWriter saves data in the training code: /mnt/data/output/runs/. Click OK to start.

   This corresponds to the code snippet: writer = SummaryWriter('/mnt/data/output/runs/mnist_experiment')

3. Click View TensorBoard to see the train_loss curve (training set loss) and the validation_loss curve (validation set loss).

   

   (Optional) Adjust hyperparameters to improve model performance

   You can evaluate your model's performance by observing the trends in the loss values:

   • The train_loss and validation_loss still show a downward trend before the training is complete (underfitting).

     You can increase num_epochs (the number of training epochs) or slightly increase the learning_rate and then retrain the model to better fit the training data.

   • The train_loss continuously decreases, but the validation_loss starts to increase before the training is complete (overfitting).

     You can decrease num_epochs or slightly decrease the learning_rate and then retrain the model to prevent overfitting.

   • Both train_loss and validation_loss are stable before the training is complete (good fit).

     When the model is in this state, you can proceed to the next steps.

   Due to length constraints, this article cannot provide extensive details on parameter tuning. You can enroll in the Alibaba Cloud Large Model ACP course to learn about key parameters in tuning commands and how to use loss curves to decide whether to continue tuning.

3.4 Deploy the trained model

For more information, see V. Use EAS to deploy model services.

4.1 Prepare the dataset and script

If you have already created a dataset in the 3. Single-node, single-GPU training section, simply download mnist_train_distributed.py and upload it. Otherwise, first 3.1 Create a dataset, and then upload the code.

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
        self.fc1 = nn.Linear(20 * 5 * 5, 128)
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 20 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
def main():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    local_rank = int(os.environ["LOCAL_RANK"])
    dist.init_process_group(backend='nccl')
    torch.cuda.set_device(local_rank)
    device = torch.device('cuda', local_rank)
    batch_size = 64
    learning_rate = 0.01
    num_epochs = 20
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    # Only the main process (rank=0) needs to download; other processes must wait for it.
    # Make processes with rank != 0 wait at the barrier.
    if rank != 0:
        dist.barrier()
    # All processes create the dataset.
    # However, only the rank 0 process actually performs the download.
    train_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=True, download=(rank == 0), transform=transform)
    # After the rank 0 process finishes downloading, it also reaches the barrier, releasing all processes.
    if rank == 0:
        dist.barrier()
    # At this point, all processes are synchronized and can continue.
    val_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=False, download=False, transform=transform)
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
    model = SimpleCNN().to(device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
    criterion = nn.CrossEntropyLoss().to(device)
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
    if rank == 0:
        writer = SummaryWriter('/mnt/data/output_distributed/runs/mnist_experiment')
    best_val_accuracy = 0.0
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 100 == 0:
                # Each rank and local_rank prints its own loss.
                print(f"Rank: {rank}, Local_Rank: {local_rank} -- Train Epoch: {epoch} "
                      f"[{batch_idx * len(data) * world_size}/{len(train_loader.dataset)} "
                      f"({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}")
                if rank == 0:
                    writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + batch_idx)
        # Validation
        model.eval()
        val_loss = 0
        correct = 0
        total = 0
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
                output = model(data)
                val_loss += criterion(output, target).item() * data.size(0)
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()
                total += target.size(0)
        val_loss_tensor = torch.tensor([val_loss], dtype=torch.float32, device=device)
        correct_tensor = torch.tensor([correct], dtype=torch.float32, device=device)
        total_tensor = torch.tensor([total], dtype=torch.float32, device=device)
        dist.all_reduce(val_loss_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(correct_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(total_tensor, op=dist.ReduceOp.SUM)
        val_loss = val_loss_tensor.item() / total_tensor.item()
        val_accuracy = 100. * correct_tensor.item() / total_tensor.item()
        if rank == 0:
            print(f'Validation Loss: {val_loss:.4f}, Accuracy: {int(correct_tensor.item())}/{int(total_tensor.item())} ({val_accuracy:.0f}%)')
            writer.add_scalar('Loss/validation', val_loss, epoch)
            writer.add_scalar('Accuracy/validation', val_accuracy, epoch)
            if val_accuracy > best_val_accuracy:
                best_val_accuracy = val_accuracy
                torch.save(model.module.state_dict(), '/mnt/data/output_distributed/best_model.pth')
                print(f'Model saved with accuracy: {best_val_accuracy:.2f}%')
    if rank == 0:
        writer.close()
    dist.destroy_process_group()
    if rank == 0:
        print('Training complete. writer.close()')
if __name__ == "__main__":
    main()

4.2 Create a distributed DLC job

1. In the left-side navigation pane of the PAI console, choose Deep Learning Containers (DLC) > Create Job.

2. Configure the following key parameters and keep the default values for the others. For a full list of parameters, see Create a training job.

   • Image Configuration: Select Image Address, and then enter the appropriate image address for your Region.

     Region	Image address
     China (Beijing)	dsw-registry-vpc.cn-beijing.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
     China (Shanghai)	dsw-registry-vpc.cn-shanghai.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
     China (Hangzhou)	dsw-registry-vpc.cn-hangzhou.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
     Other	Find your region ID and replace <Region ID> in the image address to get the full URL: dsw-registry-vpc.<Region ID>.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04

     This image uses the same verified environment from the Get Started with DSW Interactive Modeling topic. When modeling with PAI, a typical workflow is to first verify the environment and develop code in DSW, and then use DLC for training.

   • Dataset Mount: Select Custom Dataset and choose the dataset that you created. The default Mount Path is /mnt/data.

   • Startup Command: torchrun --nproc_per_node=1 --nnodes=${WORLD_SIZE} --node_rank=${RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} /mnt/data/mnist_train_distributed.py

     DLC automatically injects common environment variables, such as MASTER_ADDR and WORLD_SIZE, which you can access with the $VARIABLE_NAME syntax.

   • Source: Select Public Resources, set Number of Nodes to 2, and select ecs.gn7i-c8g1.2xlarge for the Resource Type.

     If this instance type is out of stock, you can choose another GPU instance type.

   Click Confirm to create the job. The job takes about 10 minutes to run. During execution, you can view the training Log for both instances on the Overview page.

   After the job completes, DLC saves the best model checkpoint and TensorBoard log to the output_distributed path of the mounted dataset.

4.3 (Optional) view in TensorBoard

You can use the TensorBoard visualization tool to view the loss curve and better understand the training process.

1. On the details page of the DLC job, click the TensorBoard tab, and then click Create TensorBoard.

2. Set Configuration Type to By Task. For Summary Path, enter the path where the training code saves the summary files: /mnt/data/output_distributed/runs. Click OK to start.

   This corresponds to the code snippet: writer = SummaryWriter('/mnt/data/output_distributed/runs/mnist_experiment')

3. Click View TensorBoard to see the train_loss curve (training set loss) and the validation_loss curve (validation set loss).

   

   (Optional) Adjust hyperparameters to improve model performance

   You can evaluate your model's performance by observing the trends in the loss values:

   • Before training ends, train_loss and validation_loss are still decreasing (underfitting).

     You can increase num_epochs (the number of training epochs) or slightly increase the learning_rate and then retrain the model to better fit the training data.

   • The train_loss continuously decreases, while the validation_loss starts to increase before training ends (overfitting).

     You can decrease num_epochs or slightly decrease the learning_rate and then retrain the model to prevent overfitting.

   • Both train_loss and validation_loss are stable before the training ends (good fit).

     When the model is in this state, you can proceed to the next steps.

   Due to length constraints, this article cannot provide extensive details on parameter tuning. You can enroll in the Alibaba Cloud Large Model ACP course to learn about key parameters in tuning commands and how to use loss curves to decide whether to continue tuning.

4.4 Deploy the trained model

For more information, see V. Deploying model services using EAS.