Nebulgym, a new open-source that accelerates AI training (~1.5-2x)

3 points

4 years ago

Training always takes too long. If it takes an hour, it would be better if it took 30 minutes, or maybe 15 minutes... or just 1 minute, why not? And if you want to speed up training, the techs available usually require to increase the complexity of the training process, whether it's making trade-off in terms of accuracy or time for the developer to learn a new framework. Often times it's trial and error, playing with parameters, training recipes, or switching framework/model. That's definitely not ideal.

“Fast & easy-to-use” These were keywords that motivated me to work on a new way of doing training, the library nebulgym, which now is open-source (https://github.com/nebuly-ai/nebulgym).

FAST

Training should be fast, period. Wouldn't it be great if in the near future you could train a GPT3 from scratch on your laptop? Or a large EfficientNet in a fraction of a minute? Nebulgym was built to try to bring developers closer to that future. This open-source optimizes the full training computation stack, from efficient data loading to faster forward and backward passes and earlier convergence. For example, by saving data samples in the cache on the first data read, it speeds up the full data loading process and eliminates what can become the bottleneck for the training process. Nebulgym also leverages techniques such as partial compilation of some calculations and smart sparse gradients to speed up forward and backward gradient propagations. And many more features will be implemented soon. And please let me know / open issues if you have ideas for making nebulgym even faster :)

EASY-TO-USE

"Not another framework, please, there're already 1000". That's a call for help from many developers, so nebulgym has been developed with this in mind. Nebulgym let you use the training setup you've always used, and works "on top". This is made possible with the use of class decorators (like Java's annotations). In short, you can just add these decorators before defining the model classes, and nebulgym will make sure that you use your computing resources to the fullest.

Here's a snippet of training with nebulgym decorators (@accelerate_dataset and @accelerate_model)

--------------------------

import torch

from nebulgym.decorators.torch_decorators import accelerate_model, accelerate_dataset

@accelerate_dataset()

class MyDataset(torch.utils.data.Dataset): #Your Dataset definition

@accelerate_model()

class MyModel(torch.nn.Module): # Your model definition

# Train your model as you usually do

--------------------------

And that's it. Give it a try, and leave a star , it's a little contribution to show some love for open-source projects :) Also feedback would be super appreciated!

https://github.com/nebuly-ai/nebulgym

1 comment