Modern frontier models such as the DeepSeek-V4 or GLM-5 series often contain hundreds of billions of parameters whose model weights alone require terabytes of memory and storage. Such models simply cannot fit within a single GPU or NPU regardless of its specifications. Training and serving LLMs at this scale requires distributing the workload among GPUs or NPUs often spanning multiple datacenter servers.

This article demonstrates how MindSpore and CANN provides native capabilities for distributed model training through a motivating example adapted from the official PyTorch YouTube tutorial.

Prerequisites

Introductory knowledge of deep learning and convolutional networks such as would be had by following through the first 7 chapters of the D2L textbook.

Hardware specifications

The motivating example in this article was developed and tested within my AI homelab based on the GMKtec NucBox K11 Mini PC paired with the OrangePi AI Studio Pro extension dock.

Learn all about my AI homelab setup from my previous blog posts.

1. Building my AI homelab, part 1
2. Building my AI homelab, part 2

Software environment and source code

The full training script can be found in my GitHub repository: DonaldKellett/my-ascend-jobs

This motivating example was validated within a containerized execution environment using a customized container image with MindSpore 2.10.0 + CANN 8.5.0 preinstalled and preconfigured. The CANN base container image is available on AscendHub.

1. The Dockerfile for the CANN base image can be found in the upstream GitHub repository: Ascend/cann-container-image
2. The layered Python packages used to build my customized MindSpore container image can be found in my GitHub repository: DonaldKellett/my-cann-container-image

The requirements.txt shown below for reference.

absl-py==2.5.0
decorator==5.3.1
matplotlib==3.11.1
mindspore==2.10.0
mlflow==3.15.1
ml-dtypes==0.4.0
numpy==1.26.4
sympy==1.14.0
tornado==6.5.8

Analyzing the source code

For simplicity, I used a monofile approach with everything defined in main.py. It’s sufficient for our use case since only a few helper classes and functions are involved.

Our main.py is roughly divided into the following sections.

1. Package imports and dependencies
2. Helper classes and functions
3. Our distributed training pipeline defined in main()
4. Invoking our main() function at the end of the script

Without diving into the details, here’s the overall structure of our distributed training workflow.

"""
Package imports and dependencies
"""
import ...
from ... import ...
"""
Helper classes and functions
"""
def transform_ds(...):
    ...
class MLflowLogging(...):
    ...
"""
Our distributed training pipeline
"""
def main():
    ...
if __name__ == '__main__':
    main()

Read on for the detailed breakdown and explanation of each component within our workflow.

Package imports and dependencies

The main Python packages and libraries involved in our workflow are listed below.

1. MindSpore: Huawei’s deep learning framework optimized for their Ascend datacenter and edge NPUs
2. MLflow: allows us to track deep learning experiments and visualize metrics such as model loss and accuracy
3. Matplotlib: used to visualize samples within our training data and output the result in PNG format

import gzip
import matplotlib.pyplot as plt
import mindspore
import mindspore.amp as amp
import mindspore.communication as communication
import mindspore.context as context
import mindspore.dataset as ds
import mindspore.dataset.transforms as transforms
import mindspore.dataset.vision as vision
import mindspore.nn as nn
import mindspore.ops as ops
import mlflow
import os
import urllib.request
from mindspore import dtype as mstype
from mindspore.train import Callback, LossMonitor, Model

Transforming our images and labels for training

We’ll train a modern variation of LeNet on the Fashion MNIST dataset. The standard image-based data transformations apply.

1. Each pixel per 28x28 grayscale image is scaled by a factor of 1/255 to ensure all values lie within the range [0,1]. This helps to stabilize training and prevent exploding gradients
2. Reorder the image dimensions so the channel dimension comes first: (channel,height,width). This is standard practice for feeding images into convolutional neural networks (CNN)
3. Image labels are one-hot encoded, a standard procedure for training simple classification models by minimizing the softmax cross-entropy loss

Our function transform_ds() receives a DataSet object and uses MindSpore’s built in data transformation methods to transform and split our dataset into batches of 64 samples per NPU. Since our workload is split across 2 NPUs, the global batch size is 64 * 2 = 128 samples per batch.

def transform_ds(dataset, batch_size):
    image_transforms = [
        vision.Resize(size=(28, 28)),
        vision.Rescale(rescale=1/255, shift=0),
        vision.HWC2CHW(),
        transforms.TypeCast(data_type=mstype.float32)
    ]
    label_transforms = [
        transforms.OneHot(num_classes=10),
        transforms.TypeCast(data_type=mstype.float32)
    ]
    dataset = dataset.map(operations=image_transforms, input_columns='image')
    dataset = dataset.map(operations=label_transforms, input_columns='label')
    dataset = dataset.batch(batch_size=batch_size, drop_remainder=False)
    return dataset

Tracking model hyperparameters and metrics with MLflow

Next, we define an MLflowLogging class based on mindspore.train.Callback. It creates an MLflow run to track model hyperparameters such as the learning rate and batch size as well as metrics such as:

1. The training loss at the end of each step
2. The validation loss at the end of each epoch
3. The validation accuracy at the end of each epoch

class MLflowLogging(Callback):
    def __init__(self, run_name, network_type, loss_fn, learning_rate, batch_size, epochs, num_shards, shard_id, weight_decay=0.0, momentum=0.0, optimizer='sgd'):
        super().__init__()
        self.run_name = run_name
        self.network_type = network_type
        self.loss_fn = loss_fn
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.epochs = epochs
        self.weight_decay = weight_decay
        self.momentum = momentum
        self.optimizer = optimizer
        self.num_shards = num_shards
        self.shard_id = shard_id
        self.run = mlflow.start_run(run_name=self.run_name)
        hyperparameters = {
            'learning_rate': self.learning_rate,
            'weight_decay': self.weight_decay,
            'momentum': self.momentum,
            'loss_fn': self.loss_fn,
            'optimizer': self.optimizer,
            'batch_size': self.batch_size,
            'network_type': self.network_type,
            'epochs': self.epochs,
            'num_shards': num_shards,
            'shard_id': shard_id
        }
        mlflow.log_params(hyperparameters)
    def on_train_step_end(self, run_context):
        cb_params = run_context.original_args()
        current_loss = cb_params.net_outputs.asnumpy().mean()
        mlflow.log_metric('train_loss', current_loss, step=cb_params.cur_step_num)
    def on_train_epoch_end(self, run_context):
        cb_params = run_context.original_args()
        if hasattr(cb_params, 'eval_results') and cb_params.eval_results:
            val_loss = cb_params.eval_results.get('loss', 0.0)
            val_accuracy = cb_params.eval_results.get('accuracy', 0.0)
            mlflow.log_metric('val_loss', val_loss, step=cb_params.cur_epoch_num)
            mlflow.log_metric('val_accuracy', val_accuracy, step=cb_params.cur_epoch_num)
    def on_train_end(self, run_context):
        mlflow.end_run()

With our helper classes and functions defined, let’s move on to the actual training pipeline defined in main().

Initializing HCCL for distributed training

Huawei Collective Communication Library (HCCL) is the communication library included within the CANN software which enables Ascend NPUs to communicate with each other via dedicated high-speed network links and protocols such as HCCS, RoCE and PCIe for distributed training and inference.

The distributed training job consists of a scheduler process and multiple worker subprocesses. Each worker corresponds to an available NPU device. The MS_ROLE environment variable distinguishes the scheduler from its workers. The scheduler is assigned the MS_SCHED role while workers carry the role MS_WORKER.

MS_ROLE = os.getenv('MS_ROLE')
print(f'Running msrun with role: {MS_ROLE}')

Next, we invoke mindspore.communication.init to initialize HCCL as our communication backend. Additionally, we use set_context() and set_auto_parallel_context() under mindspore.context to set the following options.

1. mode=context.GRAPH_MODE: set the execution mode to Graph mode
2. parallel_mode=context.ParallelMode.DATA_PARALLEL: use Data Parallel mode to distribute the training data across multiple NPUs
3. gradients_mean=True: use the mean of computed gradients on each NPU for parameter update at the end of each training step

context.set_context(mode=context.GRAPH_MODE, device_target='Ascend')
communication.init('hccl')
context.set_auto_parallel_context(
    parallel_mode=context.ParallelMode.DATA_PARALLEL,
    gradients_mean=True
)

With 2 NPUs on our node, each worker subprocess receives its own rank ID and the group size is 2. We obtain the rank ID and group size via the get_rank() and get_group_size() functions respectively. Furthermore, we define 2 additional variables.

1. MLFLOW_TRACKING_URI: the MLflow tracking server URL where we log our model hyperparameters and metrics to
2. MPLBACKEND: specify the Matplotlib AGG backend to visualize and output samples from our training data in PNG format

rank_id = communication.get_rank()
rank_size = communication.get_group_size()
MLFLOW_TRACKING_URI = os.getenv('MLFLOW_TRACKING_URI')
MPLBACKEND = os.getenv('MPLBACKEND')
print(f'Running rank {rank_id} of {rank_size}')
if rank_id == 0:
    print(f'Using MLflow tracking URI: {MLFLOW_TRACKING_URI}')
    print(f'Using Matplotlib backend: {MPLBACKEND}')

The rank_id == 0 guard ensures only the worker subprocess with rank ID 0 displays the MLflow tracking server URL and Matplotlib backend respectively so the message isn’t displayed twice.

Now we set the MLflow experiment name and use mlflow.set_experiment to set the active experiment.

experiment_name = '01-multi-npu-training'
experiment = mlflow.set_experiment(experiment_name=experiment_name)

Downloading the Fashion MNIST dataset

Here we create a dedicated directory data/fashion/ to store our training data downloaded from a local mirror. The Python standard library is sufficient for this task.

Take note of the details below.

1. We guard the entire data fetching pipeline with rank_id == 0. This ensures we download the training data exactly once and prevents concurrent file I/O operations from corrupting our training data
2. Notice the mindspore.ops.communication.barrier call right after our data fetching pipeline. This ensures all workers wait for the training data to be available before proceeding with the rest of the training pipeline, preventing individual workers from accessing incomplete or missing training data

dataset_dir = 'data/fashion/'
if rank_id == 0:
    os.makedirs(dataset_dir, exist_ok=True)
    prefix_url = 'https://assets.donaldsebleung.com/datasets/fashion-mnist'
    X_train_url = f'{prefix_url}/train-images-idx3-ubyte.gz'
    y_train_url = f'{prefix_url}/train-labels-idx1-ubyte.gz'
    X_test_url = f'{prefix_url}/t10k-images-idx3-ubyte.gz'
    y_test_url = f'{prefix_url}/t10k-labels-idx1-ubyte.gz'
    X_train_path = os.path.join(dataset_dir, 'train-images-idx3-ubyte')
    y_train_path = os.path.join(dataset_dir, 'train-labels-idx1-ubyte')
    X_test_path = os.path.join(dataset_dir, 't10k-images-idx3-ubyte')
    y_test_path = os.path.join(dataset_dir, 't10k-labels-idx1-ubyte')
    with urllib.request.urlopen(X_train_url) as response:
        with open(X_train_path, 'wb') as out_file:
            data_gzip = response.read()
            data = gzip.decompress(data_gzip)
            out_file.write(data)
    with urllib.request.urlopen(y_train_url) as response:
        with open(y_train_path, 'wb') as out_file:
            data_gzip = response.read()
            data = gzip.decompress(data_gzip)
            out_file.write(data)
    with urllib.request.urlopen(X_test_url) as response:
        with open(X_test_path, 'wb') as out_file:
            data_gzip = response.read()
            data = gzip.decompress(data_gzip)
            out_file.write(data)
    with urllib.request.urlopen(y_test_url) as response:
        with open(y_test_path, 'wb') as out_file:
            data_gzip = response.read()
            data = gzip.decompress(data_gzip)
            out_file.write(data)
if rank_size > 1:
    # All workers should wait for the full dataset to be downloaded before proceeding
    ops.communication.barrier()

Visualizing our dataset with Matplotlib

We use matplotlib.pyplot to display 10 random training samples as usual. 2 details worth noting:

1. We guard against concurrent file I/O with rank_id == 0 again, ensuring the PNG visualization is rendered and written to disk exactly once
2. Instead of invoking fig.show() to display the visualization directly within our Jupyter notebook, here we call plt.savefig() to save the visualization as a PNG file instead followed by plt.close().

if rank_id == 0:
    visualization_ds = ds.FashionMnistDataset(dataset_dir=dataset_dir, usage='train', shuffle=True)
    visualization_ds_samples = visualization_ds.batch(batch_size=10)
    X_samples, y_samples = next(visualization_ds_samples.create_tuple_iterator())
    X_samples, y_samples = X_samples.asnumpy(), y_samples.asnumpy()
    LABELS = [
        'T-shirt/top',
        'Trouser',
        'Pullover',
        'Dress',
        'Coat',
        'Sandal',
        'Shirt',
        'Sneaker',
        'Bag',
        'Ankle boot'
    ]
    fig, axes = plt.subplots(2, 5)
    for axis_idx, axis in enumerate(axes.flatten()):
        axis.set_title(LABELS[y_samples[axis_idx]])
        axis.imshow(X_samples[axis_idx])
        axis.axis('off')
    fig.tight_layout()
    plt.savefig(
        '00-mindspore-hccl-lenet-fashion-mnist-samples.png',
        dpi=300,
        bbox_inches='tight',
        transparent=False,
        format='png'
    )
    plt.close(fig)

Sharding and transforming our dataset

Here we instantiate our mindspore.dataset.FashionMnistDataset instances using the available training and validation data. For the training set, we further specify the following parameters to split the data between both NPUs.

1. num_shards: the group size of our distributed training job
2. shard_id: the rank ID assigned to this worker

train_ds = ds.FashionMnistDataset(
    dataset_dir=dataset_dir,
    usage='train',
    shuffle=True,
    num_shards=rank_size,
    shard_id=rank_id
)
test_ds = ds.FashionMnistDataset(
    dataset_dir=dataset_dir,
    usage='test',
    shuffle=True
)

Our per-NPU batch size is 64 and the global batch size is 64 * 2 = 128. Use the transform_ds() function we defined earlier to apply the necessary data transformations for model training and evaluation.

batch_size = 64
train_ds = transform_ds(
    dataset=train_ds,
    batch_size=batch_size
)
test_ds = transform_ds(
    dataset=test_ds,
    batch_size=batch_size
)

Defining our neural network, loss function and optimizer

We’ll use a modernized variation of LeNet for our neural network with the following improvements over the original implementation.

1. Replace sigmoid with ReLU activation to prevent vanishing gradients
2. Max-pooling often performs better in CNNs compared to average pooling

net = nn.SequentialCell([
    nn.Conv2d(1, 6, kernel_size=5, pad_mode='valid'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(6, 16, kernel_size=5, pad_mode='valid'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Dense(256, 120, activation='relu'),
    nn.Dense(120, 84, activation='relu'),
    nn.Dense(84, 10)
])

As usual for simple classification algorithms, we’ll use softmax cross-entropy for our loss function and minibatch SGD for our optimizer. Furthermore, let’s wrap our neural network with mindspore.amp.auto_mixed_precision to handle type casting between FP16 and FP32 automatically.

learning_rate = 0.1
net_amp = amp.auto_mixed_precision(network=net, amp_level='O2')
loss_fn = nn.SoftmaxCrossEntropyWithLogits(reduction='mean')
optimizer = nn.SGD(params=net_amp.trainable_params(), learning_rate=learning_rate)

Defining our training loop with the Model API

MindSpore’s Model API gives us a clean, high-level abstraction to define our training loop. We’ve seen it before so let’s focus on the differences compared to previous runs.

1. Each worker subprocess receives its own set of training data and computes its gradients in parallel. This means each callback defined in callbacks is invoked once per worker across each step and epoch
2. As a result, an MLflow run is created and updated independently per worker. We customize the run_name based on the worker’s rank_id to distinguish both runs in the MLflow UI and log the num_shards, shard_id as additional hyperparameters

run_name = '00-mindspore-hccl-lenet'
network_type = 'lenet'
loss_fn_str = 'softmax_cross_entropy'
epochs = 10
model = Model(
    network=net_amp,
    loss_fn=loss_fn,
    optimizer=optimizer,
    metrics={'accuracy', 'loss'}
)
callbacks = [
    LossMonitor(per_print_times=10),
    MLflowLogging(
        run_name=f'{run_name}-rank-{rank_id}',
        network_type=network_type,
        loss_fn=loss_fn_str,
        learning_rate=learning_rate,
        batch_size=batch_size,
        epochs=epochs,
        num_shards=rank_size,
        shard_id=rank_id
    )
]
model.fit(
    epoch=epochs,
    train_dataset=train_ds,
    valid_dataset=test_ds,
    callbacks=callbacks,
    dataset_sink_mode=False
)

Running the distributed training job with msrun

Unlike single-NPU training jobs which are executed directly with Python, we need to leverage the msrun command provided by MindSpore to run our distributed training job in parallel across multiple subprocesses.

1. The scheduler process orchestrates the entire workflow and distributes the workload across multiple worker subprocesses, 1 worker per available NPU
2. Each worker subprocess runs the entire workflow with its allocated share of training data and model parameters. Workers coordinate via the network layout provided by the scheduler to synchronize the training progress, model weights and computed gradients among each other

With our distributed training job spread across 2 Ascend 310P NPUs, the following options were specified.

1. --worker_num=2: create a total of 2 worker subprocesses across all NPUs and nodes
2. --local_worker_num=2: since both NPUs are on the same node, create 2 workers on this node
3. --join=True: inform the scheduler to wait for all workers to complete. Otherwise, the scheduler process ends early which terminates our Kubernetes job prematurely

msrun \
    --worker_num=2 \
    --local_worker_num=2 \
    --join=True \
    python \
    /app/main.py

Demo and screenshots

Here’s a GIF animation of running the distributed training job on Kubernetes. View it on Asciinema or download the Asciicast recording.

The distributed training script logged separate MLflow runs and metrics per subprocess as well, shown below.

I’d like to take this opportunity to thank my teacher DeepSeek-V4-Pro for helping me adapt the original training script for multi-NPU scenario 😉

Concluding remarks and going further

We saw in this article how MindSpore and CANN provide native capabilities for distributing model training workloads across multiple Ascend NPUs on a single node with msrun and HCCL. In production environments with frontier models involving hundreds of billions of parameters, we can extend this distributed training workflow across multiple server nodes as well.

I hope you enjoyed reading this article as much as I did authoring it and stay tuned for updates! 😉