Using Huggingface Accelerate with Weights and Biases

Accelerate is this amazing little framework that simplifies your PyTorch training scripts enabling you to train with all the tricks out there!

• Quickly convert your code to support multiple hardward (GPUS, TPUs, Metal,...)
• One code to support mixed precision, bfloat16 and even 8 bit Adam.

Minimal code and no boilerplate. Weights and Biases integration out of the box!

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator(log_with="wandb")
+ accelerator.init_trackers("my_wandb_project", config=cfg)
- device = 'cpu'
+ device = accelerator.device

  model = torch.nn.Transformer().to(device)
  optimizer = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optimizer, data = accelerator.prepare(model, optimizer, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

Training and Image Classifier

Store your configuration parameters

setup transforms

Create a simple CNN

Wrap everything into a training functions (this is necessary to run on multiple GPUS, if it is only one, you can skip the wrapping)

Let's train on 2 GPUs! This is really nice, as accelerate will take care of only calling log on the main process, so only one run get's created, so no need to manually check the rank of the process when using multiple GPUs.