Example: Migrating Local PyTorch Training Code to ModelArts

Overview

This guide shows how to turn a PyTorch training script used for local debugging into a format that works with ModelArts training jobs. It covers rebuilding data paths, setting up the runtime environment, and adapting the code for single-node multi-PU or multi-node multi-PU setups. It also includes examples of moving data from GPUs to Ascend NPUs, along with code samples before and after the changes.

Debug the code locally before uploading it to Huawei Cloud for training. This prevents hard-to-explain issues. Follow this order: local testing, joint debugging of the training job, and then the final training job.

Data and Output Path Reconstruction

Context
Generally, local training scripts directly read and write local disk paths. ModelArts training jobs run in containers, and data and model output need to be forwarded by OBS. Therefore, the first step of the reconstruction is to replace the local path in the script with the data_url (training input) and train_url (training output) parameters that can be injected by the platform. When creating a training job, you need to click Add under Environment Variable, set the input parameter name to data_url, and select an OBS directory as the data storage path. For the output, set the output parameter name to train_url, and similarly point it to an OBS directory.

Code Example

Pre-migration code (setting the data and model to local paths)

import argparse
import torch
from torchvision import datasets, transforms
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=64)
args = parser.parse_args()
# Local data path
train_data = datasets.MNIST('/home/user/data/mnist', train=True,
                             download=False, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size)
#... Training process...
# Local path for saving models
torch.save(model.state_dict(), '/home/user/output/model.pth')

After migration (parameterized data_url/train_url, with automatic mapping between OBS and local container paths)
```
import argparse
import os
import torch
from torchvision import datasets, transforms
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=64)
# ModelArts downloads the OBS data corresponding to data_url to the local container path before passing it in
parser.add_argument('--data_url', type=str, default='/home/work/data')
# ModelArts automatically uploads the contents of the local container directory corresponding to train_url back to OBS
parser.add_argument('--train_url', type=str, default='/home/work/output')
args = parser.parse_args()
train_data = datasets.MNIST(args.data_url, train=True,
                             download=False, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size)
#... Training process...
os.makedirs(args.train_url, exist_ok=True)
torch.save(model.state_dict(), os.path.join(args.train_url, 'model.pth'))
```
To create a directory or file in your training script, use the local_path directory listed in inputs or outputs. This avoids write failures from permission issues or path isolation.

If you use the ModelArts SDK to debug in a notebook before submitting for remote training, the output configuration follows the same convention: local_path is the local directory in the notebook, where the training script should save the output model or other data; obs_path is the OBS directory, to which the SDK automatically uploads the model files from local_path. The example code in this section is written using PyTorch as an example. The overall workflow remains identical across different AI frameworks. You only need to modify the value of the framework_type parameter, with no need to rewrite the migration logic when switching to another framework.

Runtime Environment (Image) Reconstruction

Selection Suggestions
If the local PyTorch/CUDA version matches the ModelArts preset image, you are advised to reuse the preset image. For example, select PyTorch and pytorch_1.8.0-cuda_10.2-py_3.7-ubuntu_18.04-x86_64 from the engine and version drop-down list. These images have been fully verified and many common installation packages have been preset. You need to create custom images only when the preset images cannot meet special requirements.
Custom Image Build Specifications
If the dependency is complex and you need to create a custom image, follow these guidelines: Keep the container image size under 15 GB. Use an open-source official image, like an official PyTorch image. Build containers in layers. Each layer should not exceed 1 GB or contain more than 100,000 files. Build layers that do not change often first, such as the OS, CUDA driver, Python, and PyTorch. If your training data and code change frequently, do not include them in the container image to avoid rebuilding it often.

If you need to migrate the local conda environment to a container (the container can meet the environment isolation requirements, and you are not advised to create multiple conda envs in the container), you can use package migration.
```
# Create a conda environment named pytorch based on the base environment to be migrated on the local host or cloud host.
conda create --name pytorch --clone base
pip install conda-pack
# Pack pytorch env to generate pytorch.tar.gz.
conda pack -n pytorch -o pytorch.tar.gz
```
Training Job Configuration Using a Custom Image
When using a custom image, the image address, code directory, and log path of the training job point to SWR and OBS, respectively. For example, the image address is set to swr.cn-north-4.myhuaweicloud.com/deep-learning/pytorch:1.8.1-cuda11.1, and the code directory is set to the directory where the boot script file is stored in OBS, the training code will be automatically downloaded to the ${MA_JOB_DIR}/demo-code directory of the training container.

Single-Node Multi-PU/Multi-Node Multi-PU Distributed Adaptation

Strategy
If you are currently using single-PU training locally, you are advised to adopt DistributedDataParallel (DDP) rather than DataParallel (DP) when scaling up to distributed training on ModelArts. DDP spawns multiple processes for computation, significantly increasing hardware resource utilization and enabling true distributed computing based on torch.distributed. Code difference between DP and DDP:
```
import torch
class Net(torch.nn.Module):
    pass
model = Net().cuda()
### DataParallel Begin ###
model = torch.nn.DataParallel(Net().cuda())
### DataParallel End ###
```

DDP Adaptation Code Example

The following example demonstrates adapting the official ResNet-18 classification task on the CIFAR-10 dataset. The training code includes three types of input parameters: basic training parameters, distributed parameters, and data-related parameters. The distributed parameters are automatically passed in by the platform and do not need to be manually defined.

Before Adaptation (Single-PU Training, Core Entry Point)

import torch
from torch import nn, optim
def main():
    model = ResNet18()
    model.cuda()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    for epoch in range(epochs):
        for data, label in train_loader:
            data, label = data.cuda(), label.cuda()
            optimizer.zero_grad()
            loss = nn.CrossEntropyLoss()(model(data), label)
            loss.backward()
            optimizer.step()

After Adaptation (Added process group initialization, Sampler, and DDP wrapper; marked as distributed adaptation points)

import argparse
import torch
import torch.distributed as dist
from torch import nn, optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
parser = argparse.ArgumentParser()
parser.add_argument('--init_method', default=None, help='tcp_port')
parser.add_argument('--rank', type=int, default=0, help='index of current task')
parser.add_argument('--world_size', type=int, default=1, help='total number of tasks')
args, unknown = parser.parse_known_args()
def main():
    ### Distributed adaptation: initialize process group ###
    dist.init_process_group(backend='nccl', init_method=args.init_method,
                             rank=args.rank, world_size=args.world_size)
    torch.cuda.set_device(args.rank % torch.cuda.device_count())
    ### End of distributed adaptation ###
    model = ResNet18().cuda()
    ### Distributed adaptation: wrap model with DDP ###
    model = torch.nn.parallel.DistributedDataParallel(model)
    ### End of distributed adaptation ###
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    ### Distributed adaptation: split data with DistributedSampler ###
    train_sampler = DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset, batch_size=64, sampler=train_sampler)
    ### End of distributed adaptation ###
    for epoch in range(epochs):
        train_sampler.set_epoch(epoch)   # Distributed adaptation: Ensure different shuffling per epoch
        for data, label in train_loader:
            data, label = data.cuda(), label.cuda()
            optimizer.zero_grad()
            loss = nn.CrossEntropyLoss()(model(data), label)
            loss.backward()
            optimizer.step()

The code above supports multi-node distributed training while remaining compatible with both CPU and GPU distributed training environments. You can easily switch back to single-node single-PU training mode by simply commenting out the marked distributed adaptation points. The init_method parameter contains the master node's IP address and port, which is automatically passed in by the platform without requiring user input. When running on a single-node multi-PU setup, you need to manually specify the world_size and rank hyperparameters when creating the training job. When running on a multi-node setup (where the number of compute nodes is greater than 1), you do not need to set them; the platform automatically injects the world_size and rank hyperparameters.

Storage Optimization for Large Datasets
For distributed training scenarios involving large datasets, you are advised to first transfer the dataset to an OBS bucket using the obsutil tool, and then migrate it to Scalable File Service (SFS) to improve I/O performance. Example command:
```
# Transfer code from OBS to SFS
./obsutil cp obs://your_bucket/YOLOX/ /mnt/sfs_turbo/code/ -f -r
```

Migrating Training from GPU to Ascend NPU

If you are using Ascend NPU compute on ModelArts, you need to complete code adaptation at the hardware level in addition to the cloud migration steps mentioned above. Note that NPUs and GPUs differ in structural architecture. Therefore, the migration process is not a direct 1-to-1 replacement; while torch.cuda and torch.npu can be used as substitutes at the expression layer, underlying differences remain in actual operator dispatching, memory management, and collective communications.

Installing Ascend Extension for PyTorch (torch_npu)
PyTorch does not natively support Ascend backends out of the box (it only directly supports CUDA and AMD ROCm). Consequently, native PyTorch GPU training code cannot run directly on Ascend devices. PyTorch 2.1 introduced a plugin mechanism for new hardware adaptation. By installing the Ascend Extension for PyTorch plugin, you can run PyTorch code directly on NPU devices. Installation verification:
```
python3 -c "import torch;import torch_npu;print(torch_npu.npu.is_available())"
```
Automatic Migration (Recommended for Simple Scenarios)
If your code does not use advanced GPU capabilities (such as custom C++/CUDA operators or direct GPU memory manipulation), you can try automatic migration first. This only requires adding two lines of code after importing torch in your main training script:
- Before migration:
```
import torch
import torch.nn as nn
#... Normal GPU training code...
```
- After migration:
```
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu   # Automatically maps CUDA APIs to NPU
import torch.nn as nn
# ... No modification needed in training code; torch.cuda calls are automatically converted to torch.npu equivalents ...
```

Manual Migration (When Automatic Migration Fails or Involves Advanced GPU Features)

If the training throws errors after automatic migration, you must manually replace CUDA APIs with NPU APIs one by one. The core modification areas include device assignment, CUDA API replacement, and distributed communication backend switching:

Device assignment comparison

# Before migration
device = torch.device('cuda:{}'.format(args.gpu))
torch.cuda.set_device(args.gpu)
# After migration
device = torch.device('npu:{}'.format(args.gpu))
torch_npu.npu.set_device(args.gpu)

Common CUDA API replacements

# Before migration
torch.cuda.is_available()
model.cuda(args.gpu)
images = images.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# After migration
torch_npu.npu.is_available()
model.npu(args.gpu)
images = images.npu(args.gpu, non_blocking=True)
target = target.npu(args.gpu, non_blocking=True)

Distributed communication backend switching (in multi-PU scenarios, you must switch the communication backend in addition to the single-PU modifications):

# Before migration (GPU uses nccl)
dist.init_process_group(backend='nccl', init_method="tcp://127.0.0.1:port",
                         ......, rank=args.rank)
# After migration (NPU uses hccl)
dist.init_process_group(backend='hccl', init_method="tcp://127.0.0.1:port",
                         ......, rank=args.rank)

Note: The Ascend NPU platform does not support the torch.nn.DataParallel API. If this API is present in your training script, you must manually convert it to torch.nn.parallel.DistributedDataParallel for multi-PU training.

Precision and Performance Verification
After completing the migration,you are advised to use the msprobe tool to compare differences between the baseline (GPU/CPU) environment and the Ascend environment. It provides features such as precision pre-checks, precision comparisons, and gradient monitoring. When troubleshooting precision issues, this tool automatically sets the dropout probability parameter (p) in APIs like torch.nn.functional.dropout and torch.nn.Dropout to 0 to eliminate randomness. Regarding performance, PyTorch executes operations at the operator level (OP-based) on Ascend AI processors. The general principle for performance optimization is to reduce host operator dispatch latency and device operator execution time. The MA-Advisor performance auto-diagnosis tool provided by ModelArts can automatically scan profiling data and generate tuning recommendations, offering a measured performance boost of approximately 10% to 30%.

If you encounter unresolved errors during automatic migration or manual adaptation, you can search for clues on the Ascend Community Forum or in PyTorch Issues on Gitee. If the issue persists, you can submit a ticket through the Huawei Cloud ModelArts portal to seek expert support.

Migration Checklist

The training script has parameterized local paths into data_url and train_url for input and output.
You have verified that the image version (preset or custom) is compatible with local PyTorch and CUDA versions.
You have planned the OBS-to-SFS data transfer route for large dataset scenarios.
You have completed DDP adaptation for distributed scenarios, distinguishing the hyperparameter injection methods between single-node multi-PU and multi-node multi-PU setups.
If targeting NPUs, you have verified the torch_npu installation and attempted automatic migration using transfer_to_npu first.
For scenarios where automatic migration fails, you have manually replaced device interfaces, communication backends, and other dependencies item by item.
You have completed precision comparison using tools such as msprobe, and performed performance diagnosis using MA-Advisor.

Parent topic: Training Example Code

Previous topic: Training Example Code

Next topic: Distributed Model Training