Getting Started with Training on Intel Gaudi

This guide provides simple steps for preparing a PyTorch model to run training on Intel® Gaudi® AI accelerator.

Make sure to install the PyTorch packages provided by Intel Gaudi. To set up the PyTorch environment, refer to the Installation Guide.The supported PyTorch versions are listed in the Support Matrix.

Once you are ready to migrate PyTorch models that run on GPU-based architecture to run on Gaudi, you can use the GPU Migration Toolkit. The GPU Migration toolkit automates the process of migration by replacing all Python API calls that have dependencies on GPU libraries with Gaudi-specific API calls, so you can run your model with fewer modifications.

Note

Creating a Simple Training Example

The following sections provide two training examples using Eager mode with torch.compile and Lazy mode. For more details, refer to PyTorch Gaudi Theory of Operations.

Note

For more detailed training examples, refer to the MNIST model or to the PyTorch Torchvision.

Example with Eager Mode and torch.compile

Create a file named torch_compile.py with the code below:

 1import torch
 2import torch.nn as nn
 3import torch.nn.functional as F
 4import torch.optim as optim
 5from torchvision import datasets, transforms
 6from torch.optim.lr_scheduler import StepLR
 7import os
 8import sys
 9import habana_frameworks.torch.core as htcore
10
11class Net(nn.Module):
12    def __init__(self):
13        super(Net, self).__init__()
14        self.conv1 = nn.Conv2d(1, 32, 3, 1)
15        self.conv2 = nn.Conv2d(32, 64, 3, 1)
16        self.dropout1 = nn.Dropout(0.25)
17        self.dropout2 = nn.Dropout(0.5)
18        self.fc1 = nn.Linear(7744, 128)
19        self.fc2 = nn.Linear(128, 10)
20
21    def forward(self, x):
22        x = self.conv1(x)
23        x = F.relu(x)
24        x = self.conv2(x)
25        x = F.relu(x)
26        x = F.max_pool2d(x, 3, 2)
27        x = self.dropout1(x)
28        x = torch.flatten(x, 1)
29        x = self.fc1(x)
30        x = F.relu(x)
31        x = self.dropout2(x)
32        x = self.fc2(x)
33        output = F.log_softmax(x, dim=1)
34        return output
35
36def train(model, device, train_loader, optimizer, epoch):
37    model.train()
38    model = torch.compile(model,backend="hpu_backend")
39
40    def train_function(data, target):
41        optimizer.zero_grad()
42        output = model(data)
43        loss = F.nll_loss(output, target)
44        loss.backward()
45        optimizer.step()
46        return loss
47
48    training_step = 0
49    for batch_idx, (data, target) in enumerate(train_loader):
50        data, target = data.to(device), target.to(device)
51        loss = train_function(data, target)
52        if batch_idx % 10 == 0:
53            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
54                epoch, batch_idx *
55                len(data), len(train_loader.dataset),
56                100. * batch_idx / len(train_loader), loss.item()))
57
58def main():
59    device = torch.device("hpu")
60
61    model = Net().to(device)
62
63    optimizer = optim.Adadelta(model.parameters(), lr=1.0)
64
65    transform = transforms.Compose([
66        transforms.ToTensor(),
67        transforms.Normalize((0.1307,), (0.3081,))
68    ])
69
70    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
71    train_loader = torch.utils.data.DataLoader(dataset, batch_size=500)
72
73    scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
74    for epoch in range(0,1):
75        train(model, device, train_loader, optimizer, epoch)
76        scheduler.step()
77
78    print("torch.compile training completed.")
79
80if __name__ == '__main__':
81    main()

This is a simple PyTorch CNN model with torch.compile enabled. The Gaudi-specific lines are explained below.

  • Line 9 - Import the Intel Gaudi PyTorch framework:

    import habana_frameworks.torch.core as htcore

  • Line 38 - Wrap the model in torch.compile function and set the backend to hpu_backend:

    model = torch.compile(model,backend="hpu_backend")

    If you want to run the model in Eager mode without torch.compile, comment out this line.

  • Line 59 - Target the Gaudi device:

    device = torch.device("hpu")

Executing the Example

After creating the torch_compile.py, perform the following:

  1. Set PYTHON to Python executable:

    export PYTHON=/usr/bin/python3.10

    Note

    The Python version depends on the operating system. Refer to the Support Matrix for a full list of supported operating systems and Python versions.

  2. Execute the torch_compile.py:

    $PYTHON torch_compile.py

Example with Lazy Mode

Create a file named example.py with the code below:

  1import torch
  2import torch.nn as nn
  3import torch.optim as optim
  4import torch.nn.functional as F
  5import torchvision
  6import torchvision.transforms as transforms
  7import os
  8
  9# Import Habana Torch Library
 10import habana_frameworks.torch.core as htcore
 11
 12class SimpleModel(nn.Module):
 13    def __init__(self):
 14        super(SimpleModel, self).__init__()
 15
 16        self.fc1   = nn.Linear(784, 256)
 17        self.fc2   = nn.Linear(256, 64)
 18        self.fc3   = nn.Linear(64, 10)
 19
 20    def forward(self, x):
 21
 22        out = x.view(-1,28*28)
 23        out = F.relu(self.fc1(out))
 24        out = F.relu(self.fc2(out))
 25        out = self.fc3(out)
 26
 27        return out
 28
 29def train(net,criterion,optimizer,trainloader,device):
 30
 31    net.train()
 32    train_loss = 0.0
 33    correct = 0
 34    total = 0
 35
 36    for batch_idx, (data, targets) in enumerate(trainloader):
 37
 38        data, targets = data.to(device), targets.to(device)
 39
 40        optimizer.zero_grad()
 41        outputs = net(data)
 42        loss = criterion(outputs, targets)
 43
 44        loss.backward()
 45        
 46        # API call to trigger execution
 47        htcore.mark_step()
 48        
 49        optimizer.step()
 50
 51        # API call to trigger execution
 52        htcore.mark_step()
 53
 54        train_loss += loss.item()
 55        _, predicted = outputs.max(1)
 56        total += targets.size(0)
 57        correct += predicted.eq(targets).sum().item()
 58
 59    train_loss = train_loss/(batch_idx+1)
 60    train_acc = 100.0*(correct/total)
 61    print("Training loss is {} and training accuracy is {}".format(train_loss,train_acc))
 62
 63def test(net,criterion,testloader,device):
 64
 65    net.eval()
 66    test_loss = 0
 67    correct = 0
 68    total = 0
 69
 70    with torch.no_grad():
 71
 72        for batch_idx, (data, targets) in enumerate(testloader):
 73
 74            data, targets = data.to(device), targets.to(device)
 75
 76            outputs = net(data)
 77            loss = criterion(outputs, targets)
 78
 79            # API call to trigger execution
 80            htcore.mark_step()
 81
 82            test_loss += loss.item()
 83            _, predicted = outputs.max(1)
 84            total += targets.size(0)
 85            correct += predicted.eq(targets).sum().item()
 86
 87    test_loss = test_loss/(batch_idx+1)
 88    test_acc = 100.0*(correct/total)
 89    print("Testing loss is {} and testing accuracy is {}".format(test_loss,test_acc))
 90
 91def main():
 92
 93    epochs = 20
 94    batch_size = 128
 95    lr = 0.01
 96    milestones = [10,15]
 97    load_path = './data'
 98    save_path = './checkpoints'
 99
100    if(not os.path.exists(save_path)):
101        os.makedirs(save_path)
102    
103    # Target the Gaudi HPU device
104    device = torch.device("hpu")
105    
106    # Data
107    transform = transforms.Compose([
108        transforms.ToTensor(),
109    ])
110
111    trainset = torchvision.datasets.MNIST(root=load_path, train=True,
112                                            download=True, transform=transform)
113    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
114                                            shuffle=True, num_workers=2)
115    testset = torchvision.datasets.MNIST(root=load_path, train=False,
116                                        download=True, transform=transform)
117    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
118                                            shuffle=False, num_workers=2)
119
120    net = SimpleModel()
121    net.to(device)
122
123    criterion = nn.CrossEntropyLoss()
124    optimizer = optim.SGD(net.parameters(), lr=lr,
125                        momentum=0.9, weight_decay=5e-4)
126    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.1)
127
128    for epoch in range(1, epochs+1):
129        print("=====================================================================")
130        print("Epoch : {}".format(epoch))
131        train(net,criterion,optimizer,trainloader,device)
132        test(net,criterion,testloader,device)
133
134        torch.save(net.state_dict(), os.path.join(save_path,'epoch_{}.pth'.format(epoch)))
135
136        scheduler.step()
137
138if __name__ == '__main__':
139    main()

The example.py presents a basic PyTorch code example. The Gaudi-specific lines are explained below.

  • Line 10 - Import habana_frameworks.torch.core:

    import habana_frameworks.torch.core as htcore

  • Line 104 - Target the Gaudi device:

    device = torch.device("hpu")

  • Lines 47, 52, 80 - In Lazy mode, mark_step() must be added in all training scripts right after loss.backward() and optimizer.step(). For further details on mark_step, refer to mark_step section.

    htcore.mark_step()

Executing the Example

After creating the example.py, perform the following:

  1. Set PYTHON to Python executable:

    export PYTHON=/usr/bin/python3.10

    Note

    The Python version depends on the operating system. Refer to the Support Matrix for a full list of supported operating systems and Python versions.

  2. Execute the example.py:

    PT_HPU_LAZY_MODE=1 $PYTHON example.py

    To use Lazy mode, the PT_HPU_LAZY_MODE=1 environment variable must be set since Eager mode with torch.compile is the default mode. Refer to Runtime Environment Variables for more information.

The following should appear as part of the output:

Epoch 1/5
469/469 [==============================] - 1s 3ms/step - loss: 1.2647 - accuracy: 0.7208
Epoch 2/5
469/469 [==============================] - 1s 2ms/step - loss: 0.7113 - accuracy: 0.8433
Epoch 3/5
469/469 [==============================] - 1s 2ms/step - loss: 0.5845 - accuracy: 0.8606
Epoch 4/5
469/469 [==============================] - 1s 2ms/step - loss: 0.5237 - accuracy: 0.8688
Epoch 5/5
469/469 [==============================] - 1s 2ms/step - loss: 0.4865 - accuracy: 0.8749
313/313 [==============================] - 1s 2ms/step - loss: 0.4482 - accuracy: 0.8869

Since the first iteration includes graph compilation time, you can see the first iteration takes longer to run than later iterations. The software stack compiles the graph and saves the recipe to cache. Unless the graph changes or a new graph comes in, no recompilation is needed during the training. Typically, the graph compilation happens at the beginning of the training and at the beginning of the evaluation.

Saving Model Checkpoints with torch.save

When working with convolutional neural networks, using torch.save to save an HPU model’s state_dict or tensors may result in errors. This issue can occur for tensors with NCHW layout which are internally permuted on the device.

As a workaround, move the model or output tensors to CPU before calling torch.save:

# Move model to CPU before saving
torch.save(model.cpu().state_dict(), 'model_checkpoint.pth')

# Move tensor to CPU before saving
torch.save(output_tensor.cpu(), 'output_tensor.pth')

Torch Multiprocessing for DataLoaders

If training scripts use multiprocessing with multiple workers for PyTorch dataloader, change the start method to spawn or forkserver using the PyTorch API multiprocessing.set_start_method(...). For example:

torch.multiprocessing.set_start_method('spawn')

Default start method is fork which may result in undefined behavior.