Thoughts on an integration with 🤗 Accelerate?

[muellerzr]

What is the problem this feature will solve?

Integration with 🤗 Accelerate opens up a wide variety of doors right at the get-go:

• Open-MMLab does not need to put in extra effort to get code going on GPU vs TPU vs M1, as Accelerate fully supports this out of the box
• Included in this is support for multi-gpu in all fashions, as Accelerate will handle all the distributed code for you
• Accelerate is trusted in the industry, being utilized already by the fastai library as well as multiple libraries by lucidrains including DALLE2-PyTorch, imagen-pytorch, as well as denoising-diffusion-pytorch

What is the feature you are proposing to solve the problem?

Modifying code to utilize 🤗 Accelerate is extremely straightforward, and leaves the code looking as close to plain PyTorch as possible. See below where the only changes are taking the code from a single-gpu and modifying it to be used across GPUs, TPUs, and M1:

+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
      inputs = inputs.to(device)
      targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()

To read more, check out these important documentation tutorials that describe various aspects of the library:

• Migrating code to 🤗 Accelerate
• Launching distributed code
• Launching distributed code from Jupyter Notebooks

Let me know if this is of interest to the team and we can assist as much as we can towards getting an integration with mmdetection going! 😄

What alternatives have you considered?

Lots and lots of custom code to get it working across all devices

[ZwwWayne]

Hi @muellerzr , Thanks for your kind suggestion. We are considering it. If we want do that, we could simply add it in mmengine, which could be used to accelerate all openmmlab projects.

[zhouzaida]

We have support DeepSpeed or FSDP in MMEngine. https://github.com/open-mmlab/mmengine/pull/1183

[hiyyg]

Wish to have Accelerate and Fabric in mmengine, they are very simple and effective!