openeuler/deepspeed repository overview

# Update DEVICE according to your device (/dev/davinci[0-7])
export DEVICE=/dev/davinci0
# Update the deepspeed image
export IMAGE=openeuler/deepspeed:0.19.9-torch_npu2.9.0-cann9.1.0-python3.12-oe2403sp4
docker run --rm \
--name deepspeed \
--device $DEVICE \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /root/.cache:/root/.cache \
-it $IMAGE bash

python3 -c "import torch, torch_npu; print('torch=', torch.__version__); print('torch_npu=', torch_npu.__version__); print('is_available:', torch.npu.is_available()); print('count:', torch.npu.device_count())"
python3 -c "import deepspeed; print('deepspeed', deepspeed.__version__)"
python3 -c "from deepspeed.accelerator import get_accelerator; print('accelerator:', get_accelerator()._name)"

import os
import urllib.request
import zipfile
from pathlib import Path

import torch
import torch.nn as nn
import torchvision
from torchvision import transforms
import deepspeed

MICRO_BATCH = 8

# CIFAR10 is downloaded from the ModelScope mirror on first use
data_dir = Path('./data')
if not (data_dir / 'cifar-10-batches-py').exists():
    data_dir.mkdir(parents=True, exist_ok=True)
    archive = data_dir / 'cifar-10-batches-py.zip'
    urllib.request.urlretrieve(
        'https://modelscope.cn/api/v1/datasets/studyhard1/cifar10-dataset/repo?Revision=master&FilePath=cifar-10-batches-py.zip',
        archive)
    with zipfile.ZipFile(archive) as zf:
        zf.extractall(data_dir)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    transforms.ConvertImageDtype(torch.bfloat16),
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=MICRO_BATCH, shuffle=True)

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

# deepspeed.initialize() injects the optimizer, ZeRO memory optimization and
# BF16 mixed precision; WORLD_SIZE is provided by the deepspeed launcher
world_size = int(os.environ.get('WORLD_SIZE', 1))
ds_config = {
    'train_batch_size': MICRO_BATCH * world_size,
    'optimizer': {
        'type': 'Adam',
        'params': {'lr': 0.001},
    },
    'zero_optimization': {'stage': 1},
    'bf16': {'enabled': True},
}
model = Net()
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config=ds_config)

criterion = nn.CrossEntropyLoss()
for epoch in range(1):
    running_loss = 0.0
    for i, data in enumerate(trainloader):
        inputs, labels = data[0].to(model_engine.device), data[1].to(model_engine.device)
        outputs = model_engine(inputs)
        loss = criterion(outputs, labels)
        model_engine.backward(loss)
        model_engine.step()
        running_loss += loss.item()
        if i % 100 == 99:
            if model_engine.global_rank == 0:
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')
            running_loss = 0.0
if model_engine.global_rank == 0:
    print('Finished Training')

deepspeed --num_gpus 1 train_cifar10.py