Skip to content
Spike Documentation

Computation Graph Synchronization between CPU and GPU

Instantiating a process involving a computation graph starts with creating an arbitrarily nested nn.Module object. In the example below, a simple convolutional neural network object consisting of chained modules is created. A dataset/dataloader is responsible for copying MNIST dataset batches to RAM, the Module object and batches are synchronized onto a device by means of the Module.to(device) and Tensor.to(device) methods, and module parameters are trained.

# Source: https://github.com/pytorch/examples/blob/main/mnist/main.py


import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR




class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)


    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output




def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break




def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()


    test_loss /= len(test_loader.dataset)


    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))




def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=14, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--no-mps', action='store_true', default=False,
                        help='disables macOS GPU training')
    parser.add_argument('--dry-run', action='store_true', default=False,
                        help='quickly check a single pass')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    use_mps = not args.no_mps and torch.backends.mps.is_available()


    torch.manual_seed(args.seed)


    if use_cuda:
        device = torch.device("cuda")
    elif use_mps:
        device = torch.device("mps")
    else:
        device = torch.device("cpu")


    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {'num_workers': 1,
                       'pin_memory': True,
                       'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)


    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)


    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)


    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()


    if args.save_model:
        torch.save(model.state_dict(), "mnist_cnn.pt")




if __name__ == '__main__':
    main()

Explicit device placement

Section titled “Explicit device placement”

In principle, all created tensors in PyTorch without explicit device placement get placed onto the CPU. This means you are responsible, by means of software enforcement, to make sure tensors reside at the correct device for optimal performance. In standard PyTorch object cases where the object is derived from nn.Module, this is effectively streamlined by the device parameter:

class Linear(nn.Module):


  def __init__(
         self,
         in_features: int,
         out_features: int,
         bias: bool = True,
         device=None,
         dtype=None,
     ) -> None:
    factory_kwargs = {"device": device, "dtype": dtype}


    ...


    self.weight = Parameter(
              torch.empty((out_features, in_features), **factory_kwargs)
    )
    if bias:
        self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
    else:
        self.register_parameter("bias", None)


    ...

Setting Module.to(device) and Tensor.to(device) for the data and label batches therefore ensures that all tensors get located onto the explicitly declared device. However, when we create our own custom nn.Module class and forget to set a device to a tensor instantiated in it, the tensor will be placed on the cpu. After profiling and identifying a placement bottleneck, it is recommended to explicitly print the string representation of tensors of interest. The string representation includes no device information when the placement device is the cpu and cuda identifier device information when placement is at the gpu:

x = torch.randn(3, 3, device='cpu')
print(x)
-->
tensor([[ 0.1443,  2.0319, -2.6817],
        [-1.6905, -0.1966,  1.4051],
        [ 0.4965,  0.1256, -0.6250]])


x = x.cuda(0)
print(x)
-->
tensor([[ 0.1443,  2.0319, -2.6817],
        [-1.6905, -0.1966,  1.4051],
        [ 0.4965,  0.1256, -0.6250]], device='cuda:0')

Control flow synchronization

Section titled “Control flow synchronization”

Sometimes, a computation graph involves more complicated control flow alongside data flow. For example, in physics-restricted neural networks, it is not uncommon to combine a neural network flow with a system solving flow. As an example, we assume calculation of a gain parameter using iterative solving alongside phase deviation determination with a neural network. This approach was suggested by Mateos-Ramos et al. in a paper called Model-Based End-to-End Learning for Multi-Target Integrated Sensing and Communication Under Hardware Impairments. Solving uses a termination condition for which solving should take place until this condition is met. Solving in this context is executed with a while loop control flow pattern:

import torch
from torch._higher_order_ops import while_loop


class While(torch.nn.Module):


    def cond_fn(self, iter_count, x):
        return iter_count.sum() > 0


    def body_fn(self, iter_count, x):
        return iter_count - 1, x * 2


    def forward(self, init_iter, init_x):
        final_iter, final_x = while_loop(self.cond_fn, self.body_fn, (init_iter, init_x))
        return final_iter, final_x


m = While()

The forum post From Tensorflow to Pytorch (while_loop) explains that the while_loop op is placed on the CPU and subsequently launches kernels to the GPU. Therefore, the chance for pipeline bottlenecking in this situation is very high due to the device synchronization required in every iteration of the while_loop op. In this situation, we recommend preventing termination based on a torch.Tensor value itself and to define a termination condition upfront as a configuration (or hyper-) parameter. Subsequently, we recommend unrolling the while loop iterations into separate chained pipeline layers. Even though this introduces excess memory usage to keep track of separate tensors and parameters per layer, it prevents synchronization because execution in this case moves kernel-to-kernel with data residing in global GPU vRAM:

import torch
import torch.nn as nn




class IterationLayer(nn.Module):


    def __init__(self):
        super().__init__()


    self.register_buffer('zero_int', torch.tensor(0, dtype=torch.long))
        self.register_buffer('one_int', torch.tensor(1, dtype=torch.long))
        self.register_buffer('two_float', torch.tensor(2.0))


    def forward(self, iter_count, x):


        cond_sum = torch.sum(iter_count)
        is_active = torch.gt(cond_sum, self.zero_int)


        decremented_iter = torch.sub(iter_count, self.one_int)
        multiplied_x = torch.mul(x, self.two_float)


        next_iter = torch.where(is_active, decremented_iter, iter_count)
        next_x = torch.where(is_active, multiplied_x, x)


        return next_iter, next_x




class UnrolledWhile(nn.Module):


  # NOTE: you have to tune num_layers to make sure the correct termination condition is reached!
    def __init__(self, num_layers: int):
        super().__init__()


    self.layers = nn.ModuleList([IterationLayer() for _ in range(num_layers)])


    def forward(self, init_iter, init_x):
        iter_count, x = init_iter, init_x


        # Sequentially route tensors through each PyTorch Module layer
        for layer in self.layers:
            iter_count, x = layer(iter_count, x)


        return iter_count, x




num_layers = 5  # Static hyperparameter
m = UnrolledWhile(num_layers=num_layers)

Context switching

Section titled “Context switching”

One of the reasons why Python is a very popular programming language has to do with the ability to combine packages contributed by third-party developers rather easily. Interpretation, rather than strict definition enforcement which is expected during compiling or tracing, in many cases, when developers define interpretation behavior upfront, makes combining packages feel almost seamless from the perspective of the user. However, interpretation comes at the cost of efficiency loss caused by both synchronization and type re-interpretation when switching context from one package type to another package type. As an example why re-interpretation is necessary, we hereby show both the Python List and Numpy ndarray datatypes (derived from Using Scientific Python Packages (NumPy, Pandas and more)):

Data structure difference between the Python List and Numpy ndarray datatypes.
Data structure difference between the Python List and Numpy ndarray datatypes.

Both datatypes each have their own niche strengths and weaknesses. Inserts are much cheaper in the Python List structure since copying is not necessary. Access efficiency is much higher in the Numpy ndarray structure because elements are contiguous layed out in RAM. In any case, we want to prevent re-interpretation (referred to as a context switch in this situation) causing the need to convert a contiguous memory block into a pointer memory table and spread values. Additionally, many optimized CUDA kernels do not support spread value access because the likelihood for encountering cache misses in shared threadblock memory drastically increases. Therefore, conversion to a datatype involving spread values also causes the data to be synchronized back to the CPU again.

We recommend sticking to a single package which involves datatypes with elements contiguous layed out in memory as much as possible. PyTorch tensors already lay out elements contiguously and sticking to pure tensor operations (not involving packages which are based on other datatypes) does not force re-interpretation. Additionally, pure tensor operations allows deterministic graph tracing to chained CUDA kernel launches without the need for additional re-interpretation and/or synchronization ops.