Distribution and aggregation of graph operations over multiple (fractions of) GPUs for advanced users
Graph operation distribution
Section titled “Graph operation distribution”Operation distribution is not a new machine learning concept. Back in the 1980’s, performance as a result of hardware development (CPUs and complex instruction set machines) at a particular point in time was coming to a halt. In response, Inmos designed the transputing architecture, including the associated transputer hardware modules. In the transputer architecture, one transputer acts as the central authority for booting a system containing a number of connected transputers. The transputers in principle communicate over separate serial links, one transputer message being communicated to another transputer at a time. The loader would boot the central transputer, which would then begin sending boot code to the other transputers in the network, and could customize the code sent to each one.
Both communication and computation can be abstracted in software by means of a message passing interface. As an example, we discuss the interface called openmpi:
#include <stdio.h>#include "mpi.h"
int main(argc, argv)int argc;char **argv;{ int rank, value;
MPI_Init(&argc, &argv); MPI_Comm_rank(MPI_COMM_WORLD, &rank); //sets rank of the current processor, root = 0, to rank variable printf("hello I am process %d\n", rank); //output rank or process ID of the this process
if (rank == 0) { printf("Enter a value to send to processor %d:\n", destination); scanf("%d", &buffer); MPI_Isend(&buffer, count, MPI_INT, destination, tag, MPI_COMM_WORLD, &request); //non blocking send to destination process }
if (rank == destination) { MPI_Irecv(&buffer, count, MPI_INT, 0, tag, MPI_COMM_WORLD, &request); //destination process receives }
MPI_Wait(&request, &status); //bloks and waits for destination process to receive data
if (rank == 0) { printf("processor %d sent %d\n", rank, buffer); } if (rank == destination) { printf("processor %d got %d\n", rank, buffer); // Perform some computations ... }
MPI_Finalize();
return 0;}In principle, all ‘transputers’ receive the same source code when the openmpi software interface would be used. The rank and the world size determine which transputer is going to wait for information, compute, etc. When a particular transputer reaches the state of finalization, it is assumed to be finished and switches off. The same communication and computation principles in the situation of computation distribution are used in machine learning frameworks today.
Why do specific machine learning computation graphs implemented with these frameworks require access to operation distribution? Attention-based neural network architectures do not contain a severely restrictive inductive bias. This means that as long as the amount of data used for training and the number of model parameters can be scaled up, this neural network architecture generalizes better to the task it needs to perform. The time required to process the entire dataset on a single GPU in a single epoch might be too long, or the amount of model parameters causes the computation graph not to fit anymore in GPU vRAM. In this situation, the same model parameters have to be processed in parallel on different GPUs for different batches, or computation graph sections have to be placed on different GPUs and intermediate batch representations have to be synchronized between the GPUs accordingly. Data, in contrast to the model parameters, might not be allowed to leave the administrative domain of a particular system. Therefore, a communication mechanism has to be set up in which model parameters are synchronized to other systems so that a model can still be trained. This form of model parameter training is typically referred to as federated learning. In federated learning, local data processing and sharing model parameters is required due to privacy concerns.
For the explanations that follow, we assume you are familiar with the simple model training example in Beginner -> CPU and GPU synchronization. In the example, 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. If you would like to read more in-depth material about all potential different distribution configurations, including quantitative results depicting the effect on computation and communication efficiency, we recommend the The Ultra-Scale Playbook.
Configure data loading for distribution
Section titled “Configure data loading for distribution”To make sure that particular GPUs do not receive the same input batches, it is important to limit the batch indices that particular GPUs have access to throughout an epoch. The torch.utils.data.DataLoader class can be given a sampler as constructor parameter. The sampler is responsible for controlling the sequence of indices/keys used in data loading. Samplers represent iterable objects over the indices to datasets. The torch.utils.data.distributed.DistributedSampler can be used to limit the indices to be sampled for the source code responsible for graph computation for a particular rank. The rank subsequently determines on which GPU the source code is placed and executed:
...
dataset1 = datasets.MNIST('../data', train=True, download=True, transform=transform)dataset2 = datasets.MNIST('../data', train=False, transform=transform)
sampler1 = DistributedSampler(dataset1)sampler2 = DistributedSampler(dataset2)
train_loader = torch.utils.data.DataLoader(dataset1, sampler=sampler1, **train_kwargs)test_loader = torch.utils.data.DataLoader(dataset2, sampler=sampler2, **test_kwargs)Please note that a particular GPU receiving an input batch is subject to the operation distribution type deployed. For example, because in model pipeline parallelism GPUs receive different pipeline stage representations rather than input batches, the limit is different (no limit when 1 pipeline is involved).
Data parallel
Section titled “Data parallel”When the amount of involved model parameters is limited, however, the amount of data tensors the parameters should be trained on is large (in the order of several TBs), data parallel distribution can be applied. The idea behind data parallel distribution is that every GPU with attached process involved in the training process receives the same neural network architecture, including initial parameters, and stores this in vRAM. Subsequently, every training iteration involves a fork-join and synchronization process. Inside the fork step, every process samples their own batch tensor, performs a local forward-pass through the local neural network architecture replica, and subsequently performs a local backward-pass for computation of local replica gradients. Inside the join step, the local replica gradient calculation operation is awaited for by means of synchronization and the local replica gradients are subsequently all-reduced. Subsequently, the all-reduced global gradient present at all processes now is used to perform an optimizer step on the local neural network parameters on the GPU. An example on how this can be applied onto the example in Beginner -> CPU and GPU synchronization in a single-node multi-gpu situation is provided below:
import osimport argparseimport torchimport torch.nn as nnimport torch.distributed as distimport torch.multiprocessing as mpimport torch.nn.functional as Fimport torch.optim as optimfrom torchvision import datasets, transformsfrom torch.optim.lr_scheduler import StepLRfrom torch.nn.parallel import DistributedDataParallel
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 worker(rank, world_size, gpu_count, device_type, bsize, test_bsize, lr, gamma, save_model, epochs, fraction):
# Please note that the Nvidia NCCL library is currently limited to a single process per GPU device. # Therefore, fraction should always equate to 1.0 when using NCCL (this might work with other message passing interfaces) # We keep applying a fraction part of the example because this might change in the future in case fractioning is adopted by NCCL if device_type == "cuda": dist.init_process_group("nccl", rank=rank, world_size=world_size) device_id = rank % gpu_count device = torch.device(device_type, device_id) torch.cuda.set_per_process_memory_fraction(fraction, device)
else: dist.init_process_group("gloo", rank=rank, world_size=world_size) device = torch.device(device_type)
train_kwargs = {'batch_size': bsize, 'transform': transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])} test_kwargs = {'batch_size': test_bsize, 'transform': transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])}
if device_type == "cuda": cuda_kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True, } train_kwargs.update(cuda_kwargs) test_kwargs.update(cuda_kwargs)
// Configure train_loader, test_loader according to the dataset/dataloader example provided above // Note: provide train_kwargs, test_kwargs to dataloaders as well ...
model = Net().to(device) ddp_model = DistributedDataParallel(model, device_ids=device)
optimizer = optim.Adadelta(ddp_model.parameters(), lr=lr) scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
for epoch in range(1, epochs + 1): train(..., ddp_model, device, train_loader, optimizer, epoch) test(ddp_model, device, test_loader) scheduler.step()
if save_model:
# Ensure that all processes have reached this point such that no updates are still being performed dist.barrier()
if rank == 0: torch.save(ddp_model.module.state_dict(), "mnist_cnn.pt")
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('--fraction', type=float, default=1.0, metavar='M', help='offset GPU count by applied software-enforced fraction (default: 1.0)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA 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()
if use_cuda: device_str = "cuda" else: device_str = "cpu"
gpu_count = torch.cuda.device_count() process_count = int(gpu_count / args.fraction) if use_cuda else 8
mp.spawn(worker, args=(process_count, gpu_count, device_str, args.batch_size, args.test_batch_size, args.lr, args.gamma, args.save_model, args.epochs, args.fraction), nprocs=process_count, join=True)
if __name__ == '__main__': # Rendezvous location where process access information is communicated to os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29500" main()Model pipeline parallel
Section titled “Model pipeline parallel”There are situations where the amount of involved model parameters is very large. This is typically the case with foundational attention-based representation models. The required batch size in combination with the parameter count may cause vRAM occupation to reach outside the capacity of a single GPU. In the situations sketched above, distributed pipelining can be employed. The idea behind distributed pipelining is that graph execution is staged into different pipeline steps. Pipeline step execution can subsequently be overlapped across different GPUs. This means that while execution of pipeline step 1 on GPU 1 for a particular input batch is happening, execution of pipeline step 2 on GPU 2 for another batch can happen simultaneously. In the case of model pipeline parallelism, a step typically refers to a layer (or layers) of a particular neural network architecture used in the model. An example on how this can be applied onto the example in Beginner -> CPU and GPU synchronization in a single-node multi-gpu situation is provided below:
import osimport argparseimport torchimport torch.nn as nnimport torch.distributed as distimport torch.multiprocessing as mpimport torch.nn.functional as Fimport torch.optim as optimfrom torchvision import datasets, transformsfrom torch.optim.lr_scheduler import StepLRfrom torch.distributed.pipelining import PipelineStage, ScheduleGPipe
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):
# The goal of the if checks is to enable correct layer functionality # Correct functionality entails that the layer or op should not be skipped when a layer is active in a particular stage
x = self.conv1(x) if self.conv1 else x x = F.relu(x) if self.conv1 else x
x = self.conv2(x) if self.conv2 else x x = F.relu(x) if self.conv2 else x x = F.max_pool2d(x, 2) if self.conv2 else x
x = self.dropout1(x) if self.dropout1 else x
x = torch.flatten(x, 1) if self.fc1 else x x = self.fc1(x) if self.fc1 else x x = F.relu(x) if self.fc1 else x
x = self.dropout2(x) if self.dropout2 else x
x = self.fc2(x) if self.fc2 else x output = F.log_softmax(x, dim=1) if self.fc2 else x return output
def train(args, model, device, train_loader, optimizer, epoch, rank):
model.stage.submodule.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
# Note: in case of more granular pipelining, intermediate layers call .step() without involvement of data, target if rank == 0: data = data.to(device) output = model.step(data)
elif rank == 1: loss_list = [] target = target.to(device) output = model.step(target=target, losses=loss_list)
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_list[0].item()))
else: raise ValueError("This distributed pipeline program requires 2 GPUs to function.")
optimizer.step()
# If we need to break for a test run, break all processes after the barrier has been reached if batch_idx % args.log_interval == 0: if args.dry_run: dist.barrier() break
def test(model, device, test_loader, rank):
test_loss = 0 correct = 0
model.stage.submodule.eval() with torch.no_grad():
for data, target in test_loader:
if rank == 0: data = data.to(device) output = model.step(data)
elif rank == 1: loss_list = [] target = target.to(device) output = model.step(target=target, losses=loss_list)
test_loss += torch.sum(loss_list[0]).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()
else: raise ValueError("This distributed pipeline program requires 2 GPUs to function.")
if rank == 1: 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 worker(rank, world_size, gpu_count, device_type, bsize, test_bsize, lr, gamma, save_model, epochs, fraction):
if device_type == "cuda": dist.init_process_group("nccl", rank=rank, world_size=world_size) device = torch.device(device_type, rank)
else: dist.init_process_group("gloo", rank=rank, world_size=world_size) device = torch.device(device_type)
train_kwargs = {'batch_size': bsize, 'transform': transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])} test_kwargs = {'batch_size': test_bsize, 'transform': transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])}
if device_type == "cuda": cuda_kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True, } train_kwargs.update(cuda_kwargs) test_kwargs.update(cuda_kwargs)
// Configure train_loader, test_loader according to the dataset/dataloader example provided above // Note: provide train_kwargs, test_kwargs to dataloaders as well // Note: since the GPUs ingest intermediate pipeline stage representations rather than different batches, limit is not necessary ...
model = Net() if rank == 0:
model.dropout2 = None model.fc1 = None model.fc2 = None
elif rank == 1:
model.conv1 = None model.conv2 = None model.dropout1 = None
else: raise ValueError("This distributed pipeline program does not supported processes with a rank >= 2")
stage = PipelineStage(model, rank, world_size, device) model.to(device) p_schedule = ScheduleGPipe(stage, n_microbatches=4, loss_fn=F.nll_loss)
optimizer = optim.Adadelta(model.parameters(), lr=lr) o_scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
for epoch in range(1, epochs + 1): train(..., p_schedule, device, train_loader, optimizer, epoch, rank) test(p_schedule, device, test_loader, rank) o_scheduler.step()
if save_model:
# Ensure that all processes have reached this point such that no updates are still being performed dist.barrier()
# Note: checkpointing happens per stage model_state, optimizer_state = dist.checkpoint.state_dict.get_state_dict( model, optimizer, options=dist.checkpoint.state_dict.StateDictOptions(full_state_dict=True) ) dist.checkpoint.save( {"model": model_state, "optimizer": optimizer_state}, checkpoint_id=f"checkpoint-dist-mnist-cnn", # each rank will save its own file )
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('--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()
if use_cuda: device_str = "cuda" else: device_str = "cpu"
gpu_count = torch.cuda.device_count() if use_cuda and gpu_count < 2: raise ValueError("This distributed pipeline program requires 2 GPUs to function.")
process_count = 2
mp.spawn(worker, args=(process_count, gpu_count, device_str, args.batch_size, args.test_batch_size, args.lr, args.gamma, args.save_model, args.epochs, args.fraction), nprocs=process_count, join=True)
if __name__ == '__main__': # Rendezvous location where process access information is communicated to os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29500" main()Tensor parallel
Section titled “Tensor parallel”In most typical machine learning pipelines, tensor parallelism is not necessary. The idea behind tensor parallelism is that parts of a matrix multiply-accumulate operation are sharded across different GPUs. This is helpful in situations where the world size of a distributed machine learning training procedure becomes so large to the extent that communication latency becomes a bottleneck. In this case, the decision can be made to limit communication by introducing intra-node tensor parallelism while applying another distributed training procedure inter-node. Because the amount of GPUs inside the Spike-1 cluster is limited, this level of world size is never going to be reached. Please recall from a linear algebra class that matrix multiplication entails element-wise multiplication and subsequent accumulation reduction for all row-column combinations. In column-wise parallelism, you shard one matrix across column groups, multiply with the entire matrix you wish to multiply with, and subsequently perform an all-gather and concatenation operation. In row-wise parallelism, you shard one matrix across row groups, multiply with the entire matrix you wish to multiply with, and subsequently perform an all-reduce operation. Because support for tensor parallelism is limited to nn.Linear, nn.Embedding, nn.LayerNorm, nn.Dropout, we will not provide a simple example. In what follows, we briefly show how parallelism can be configured for portions of a nn.Module class depicting a transformer architecture for a GPU group while replicating over multiple GPU groups in data parallel form.
...
import torchimport torch.distributed as distfrom torch.nn.parallel import DistributedDataParallelfrom torch.distributed.device_mesh import init_device_meshfrom torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel
...
def worker(rank, world_size, gpu_count, device_type, bsize, test_bsize, lr, gamma, save_model, epochs, fraction):
if device_type == "cuda": dist.init_process_group("nccl", rank=rank, world_size=world_size) device = torch.device(device_type, rank)
else: dist.init_process_group("gloo", rank=rank, world_size=world_size) device = torch.device(device_type)
# Create process groups to manage 2-D like parallel pattern where 4 groups of 2 GPUs are used for data parallelism. # Inside every group, 2 GPUs are used for tensor parallel sharding. mesh = init_device_mesh( device_type=device_type, mesh_shape=(4, 2), mesh_dim_names=("dp", "tp") ) tp_group = mesh["tp"].get_group() dp_group = mesh["dp"].get_group()
...
model = TransformerModel() for layer_id, transformer_block in enumerate(model.layers):
layer_tp_plan = { # by default ColwiseParallel input layouts is replicated # and RowwiseParallel output layouts is replicated "attention.wq": ColwiseParallel(use_local_output=False), "attention.wk": ColwiseParallel(use_local_output=False), "attention.wv": ColwiseParallel(use_local_output=False), "attention.wo": RowwiseParallel(), "feed_forward.w1": ColwiseParallel(), "feed_forward.w2": RowwiseParallel(), "feed_forward.w3": ColwiseParallel() }
parallelize_module( module=transformer_block, device_mesh=tp_group, parallelize_plan=layer_tp_plan )
model = parallelize_module( model, tp_group, { "tok_embeddings": RowwiseParallel( input_layouts=Replicate() ), "output": ColwiseParallel( output_layouts=Replicate() ), } )
model.to(device) ddp_model = DistributedDataParallel(model, device_ids=device, process_group=dp_group)
optimizer = optim.Adadelta(ddp_model.parameters(), lr=lr) scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
for epoch in range(1, epochs + 1): train(..., ddp_model, device, train_loader, optimizer, epoch) test(ddp_model, device, test_loader) scheduler.step()
...
...Full sharding
Section titled “Full sharding”For a dataset and neural network architecture that keeps scaling, it is typically advised to increase the number of involved model replicas across GPUs in a data parallel training procedure. Additionally, further configuration of sharding alongside a GPU mesh group for both (a set of) layer(s) and tensor is typically advised. As you might have already seen in the earlier examples, additional configuration of a mesh group and static layer splitting feels clunky and does not follow the code modularization standard. This limits widespread adoption of very-large model and dataset training to a handful of engineers that have been instructed to apply configuration because the stakeholder value to be gained is high enough. In response to the clunkiness and the desire to further allow scaling according to the modularization standard, Fully Sharded Data Parallelism (FSDP) was invented as a new distribution procedure.
FSDP works with help of a sharding factor F.
When F = 1, FSDP acts in a normal distributed data parallel way (see the subsection on Data parallel distribution) by limiting distribution to single GPU neural network replica groups.
When 1 < F <= n_gpus, a device mesh according to the factor is enabled whereby arbitrary sharding is applied in a particular inner mesh group.
Arbitrary sharding involves proportional distribution of for example rows and or columns for all layers along GPUs in the inner mesh group.
For the sake of simplicity, we consider loss and gradient distribution outside the scope of this discussion.
Prior to applying a forward or backward pass, gather-scatter communication operations are performed for the row-wise and or column-wise parallel matrix multiplication calculations.
The passes are awaited for by means of a join operation after which gradients are all-reduced.
Finding the correct shards required for the computation requires a sophisticated GPU routing mechanism.
For example, matrix row and or column indexes with layer name can be given to a hashing algorithm.
The hash is subsequently used to calculate a GPU index alongside a ring (for more information see an example of the distributed hash table algorithm).
In what follows, we provide a brief small example on how FSDP can be configured:
import osimport argparseimport functoolsimport torchimport torch.nn as nnimport torch.distributed as distimport torch.multiprocessing as mpimport torch.nn.functional as Fimport torch.optim as optimfrom torchvision import datasets, transformsfrom torch.optim.lr_scheduler import StepLRfrom torch.distributed.fsdp import FullyShardedDataParallel, ShardingStrategy, FullStateDictConfig, StateDictType
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 custom_auto_wrap_policy( module: nn.Module, recurse: bool, nonwrapped_numel: int, # Additional custom arguments min_num_params: int = int(1e8),) -> bool: return nonwrapped_numel >= min_num_params
def worker(rank, world_size, gpu_count, device_type, bsize, test_bsize, lr, gamma, save_model, epochs, fraction):
# Please note that the Nvidia NCCL library is currently limited to a single process per GPU device. # Therefore, fraction should always equate to 1.0 when using NCCL (this might work with other message passing interfaces) # We keep applying a fraction part of the example because this might change in the future in case fractioning is adopted by NCCL if device_type == "cuda": dist.init_process_group("nccl", rank=rank, world_size=world_size) device_id = rank % gpu_count device = torch.device(device_type, device_id) torch.cuda.set_per_process_memory_fraction(fraction, device)
else: dist.init_process_group("gloo", rank=rank, world_size=world_size) device = torch.device(device_type)
train_kwargs = {'batch_size': bsize, 'transform': transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])} test_kwargs = {'batch_size': test_bsize, 'transform': transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])}
if device_type == "cuda": cuda_kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True, } train_kwargs.update(cuda_kwargs) test_kwargs.update(cuda_kwargs)
// Configure train_loader, test_loader according to the dataset/dataloader example provided above // Note: provide train_kwargs, test_kwargs to dataloaders as well ...
model = Net().to(device) my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5)) fsdp_model = FullyShardedDataParallel(model, device_id=device, auto_wrap_policy=my_auto_wrap_policy, sharding_strategy=ShardingStrategy.HYBRID_SHARD)
optimizer = optim.Adadelta(fsdp_model.parameters(), lr=lr) scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
for epoch in range(1, epochs + 1): train(..., fsdp_model, device, train_loader, optimizer, epoch) test(fsdp_model, device, test_loader) scheduler.step()
if save_model:
# Ensure that all processes have reached this point such that no updates are still being performed dist.barrier()
# Ensure that the entire model, in RAM, once unsharded, fits save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_policy): state_dict = model.state_dict()
if rank == 0: torch.save(state_dict, "trained_fsdp_model.pt")
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('--fraction', type=float, default=1.0, metavar='M', help='offset GPU count by applied software-enforced fraction (default: 1.0)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA 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()
if use_cuda: device_str = "cuda" else: device_str = "cpu"
gpu_count = torch.cuda.device_count() process_count = int(gpu_count / args.fraction) if use_cuda else 8
mp.spawn(worker, args=(process_count, gpu_count, device_str, args.batch_size, args.test_batch_size, args.lr, args.gamma, args.save_model, args.epochs, args.fraction), nprocs=process_count, join=True)
if __name__ == '__main__': # Rendezvous location where process access information is communicated to os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29500" main()Multi-GPU run:ai launching
Section titled “Multi-GPU run:ai launching”Run:ai user interface standard training configuration
Section titled “Run:ai user interface standard training configuration”
Run:ai user interface configuration of multiple GPUs
Section titled “Run:ai user interface configuration of multiple GPUs”
Run:ai command-line interface
Section titled “Run:ai command-line interface”# See command-line option --gpu-devices-requestrunai training submit a-test-workload -i harbor.spike.tue.nl/aiteam/pretrainingcontainer:latest -p aiteam --existing-pvc claimname=aiteam,path=/data --user-group-source fromTheImage --cpu-core-limit 7 --cpu-core-request 7 --cpu-memory-limit 64G --cpu-memory-request 64G --gpu-devices-request 4 --command -- bash -lc 'cd /data; python3 feature_extraction_eval.py --start_index 18999 --end_index 20000 --model_name exported_step_100000.pt'Elastic run:ai launching
Section titled “Elastic run:ai launching”So far in all aforementioned examples we have been using the torch.multiprocessing.spawn method for spinning up multiple processes locally. Each process is asked to execute the implemented worker method with provided arguments. Unfortunately, torch.multiprocessing.spawn does not scale beyond the confines of a single DGX B200 node. Therefore, torch.multiprocessing.spawn is only suitable for distributed training and evaluation processes which require access to at most eight GPUs. To account for message passing alongside both local and inter-node processes, PyTorch includes the torch.distributed.run console script. This console script can in turn be given a worker training and evaluation script (it should therefore ONLY include the functionality of the worker function in previous examples) which you can configure yourself.
When we take a deeper look into the torch.distributed.run console script, we can see that the argparse package is used to read node, process, and process rendezvous settings. Additionally, when settings are not provided by the user as script arguments, the custom torch.distributed.argparse_util.env argparse action method reads these arguments from environment variables. These environment variables are injected automatically into the pod container by the GPU Operator stack when configuring distributed training in the Run:ai interface. Another interesting observation is that during worker launch, a health check server is setup. This health server accounts for a potentially long rendezvous procedure before a network information key-value store is ready by making sure network ports are healthy. Therefore, if you are considering distributed training with at most eight GPUs, to circumvent initial long experiment initialization time, we recommend distributed launching with torch.multiprocessing.spawn.
Even though rendezvous introduces additional experiment time for both initialization and information synchronization, it opens up the door to some interesting new distribution procedures. Recall from the sections Data parallel and Fully sharded data parallel that each GPU (or GPU mesh group) uses a graph replica for processing of an additional batch size tensors. In case the DGX nodes are underutilized at some point, it is nice that the distributed training process can account for this by increasing the replica group, thereby being able to process more batches. This is typically referred to as elastic distributed launching. Additionally, when there is no need for inter-process synchronization (typically the case for bag-of-task processes), a distributed training process can also be performed master-less.
Run:ai user interface distributed training configuration
Section titled “Run:ai user interface distributed training configuration”
Run:ai user interface distributed training elastic configuration of multiple GPUs
Section titled “Run:ai user interface distributed training elastic configuration of multiple GPUs”
Run:ai command-line interface distributed training configuration
Section titled “Run:ai command-line interface distributed training configuration”# See the change from training submit to training pytorch submit, the command-line option --workers, and options --min-replicas, --max-replicasrunai training pytorch submit a-test-workload -i harbor.spike.tue.nl/aiteam/pretrainingcontainer:latest -p aiteam --existing-pvc claimname=aiteam,path=/data --user-group-source fromTheImage --cpu-core-limit 7 --cpu-core-request 7 --cpu-memory-limit 64G --cpu-memory-request 64G --gpu-devices-request 1 --workers 16 --min-replicas 1 --max-replicas 16 --command -- bash -lc 'cd /data; python3 feature_extraction_eval.py --start_index 18999 --end_index 20000 --model_name exported_step_100000.pt'Effect of NVLink and NCCL on communication efficiency
Section titled “Effect of NVLink and NCCL on communication efficiency”In standard servers, GPUs typically interface over a PCIe interconnect. The total PCIe interconnect bandwidth for Gen 5.0 16 slots is limited bidirectionally to 128 GB/s. Nvidia also offers their own GPU-to-GPU interconnect system called NVLink. The total NVLink interconnect bandwidth for Gen 5.0 (compatible with the Blackwell architecture) is limited bidirectionally to 1800 GB/s. This means that GPU-to-GPU inter-process communication is an order of magnitude faster over NVLink, provided that the volume of information to be communicated scales accordingly. The Nvidia communication library NCCL has the ability to recognize NVLink and to fallback to PCIe communication when access to a NVLink interconnect network is not found. All Spike-1 DGX nodes contain B200 GPUs which have been interconnected with NVLink technology using NVswitch hardware in a so-called mesh configuration. This is especially helpful in a situation where arbitrary sharding can occur, especially in the situation of FSDP.