Compiling Deepchem Torch Models ¶

Deep Learning models typically involve millions or even billions of parameters (as in the case of LLMs) that need to be fine-tuned through an iterative training processes. During training, these models process vast amounts of data to learn patterns and features effectively. This data-intensive nature, combined with the computational complexity of operations like matrix multiplications and backpropagation, leads to lengthy training times that can span days, weeks, or even months on standard hardware. Additionally, the need for multiple training runs to experiment with the different model and hyperparameter configurations further extends the overall development time.

Effective optimization techniques can significantly reduce training times, lower computational costs, and improve model performance. This makes optimization particularly crucial in research and industrial settings where faster iterations can accelerate scientific discoveries, product development, and the deployment of AI solutions. Moreover, as models grow larger and more sophisticated, optimization plays a vital role in making advanced AI accessible and practical for a wider range of applications and environments.

To address the need for optimization of Deep Larning models and as an improvement over existing methods, PyTorch introduced the torch.compile() function in PyTorch 2.0 to allow faster training and inference of the models. torch.compile() works by compiling PyTorch code into optimised kernels using a JIT(Just in Time) compiler. Different models show varying levels of improvement in run times depending on their architecture and batch size when compiled. Compared to existing methods like TorchScript or FX tracing , compile() also offers advantages such as the ability to handle arbitrary Python code and conditional graph-breaking flow of the inputs to the models. This allows compile() to work with minimal or no code modification to the model.

DeepChem has builtin support for compiling PyTorch models using torch.compile() and using this feature, users can efficiently run PyTorch models and achieve significant performance gains. This tutorial contains the steps for compiling DeepChem PyTorch models, benchmarking and evaluating their performance with the uncompiled models.

NOTE: DeepChem contains many models with varying architecture and complexity. Not all models will show significant improvements in run times when compiled. It is recommended to test the models with and without compilation to determine the performance improvements.

Colab ¶

This tutorial and the rest in this sequence can be done in Google colab. If you'd like to open this notebook in colab, you can use the following link.

Compilation Process ¶

This section gives an introductory explanation about the compilation process of PyTorch models and assumes prior knowledge about forward pass, backward pass and computational graphs in neural networks. If you're unfamiliar with these concepts, you can refer to these slides for a basic understanding. Alternatively, you can proceed to the next section to learn how to compile and benchmark DeepChem models without delving into the internal details of the compilation process.

Image taken from PyTorch2.0 Introductory Blog

The compilation process is split into multiple steps which uses many new technologies that were introduced in PyTorch 2.0. The process is as follows:

1. Graph Acquisition: During the compilation process, TorchDynamo and AOTAutograd are used for capturing the forward and backward pass graphs respectively. AOTAutograd allows the backward graph to be captured ahead of time without needing a backward pass to be performed.

2. Graph Lowering: The captured graph that could be composed of the 2000+ PyTorch operators is lowered into a collection of ~250 Prim and ~750 ATen operators.

3. Graph Compilation: In this step optimised low-level kernels are generated for the target accelerator using a suitable backend compiler. TorchInductor is the default backend compiler used for this purpose.

Deepchem uses the torch.compile() function that implements all the above steps internally to compile the models. The compiled model can be used for training, evaluation and inference.

For more information on the compilation process, refer to PyTorch2.0 Introductory Blog that does a deep dive into the compilation process, technical decisions and future features for the compile function. You can also refer to the Huggingface blog, Optimize inference using torch.compile() that benchmarks many common PyTorch models and shows the performance improvements when compiled.

Compiling Models ¶

The compile function is only available in DeepChem for models that use PyTorch as the backend (i.e inherits TorchModel class). You can see the complete list of models that are available in DeepChem and their backends in the DeepChem Documentation here .

This tutorial contains the steps to load a DeepChem model, compile it and evaluate the performance improvements when compiled for both training and inference. Refer to the documentation of DeepChem's compile function to read more about the different parameters you can pass to the function and their usage.

If you just want to compile the model, you can add the line model.compile() after initialising the model. You DO NOT have to make any changes to the rest of your code.

Some of the things to keep in mind when compiling models are:

1. Selecting the right mode: The modes can be default , reduce-overhead , max-autotune or max-autotune-no-cudagraphs . Out of this reduce-overhead and max-autotune modes requires triton to be installed. Refer to the PyTorch docs on torch.compile for more information on the modes.

2. Setting fullgraph parameter: If True (default False ), torch.compile will require that the entire function be capturable into a single graph. If this is not possible (that is, if there are graph breaks), then the function will raise an error.

3. Experimenting with different parameter configuration: Different parameter configurations can give different speedups based on the model, batch size and the device used for training/inference. Experiment with a few parameter combinations to check which one gives better results.

In this tutorial, we will be using DMPNN model and Freesolv Dataset for training and inference of the models.

The below line is the only addition you have to make to the code for compiling the model. You can pass in the other arguments too to the compile() function if they are required.

Benchmarking model Speedups ¶

This section contains the steps for benchmarking the performance of models after compilation process for both training and inference. We are using the same model(DMPNN) and dataset(FreSolv) in this section too. The steps for compilation and benchmarking is same for other models as well.

To account for the initial performance overhead of kernel compilation in compiled models, median values are employed as the performance metric throughout the tutorial for calculating speedup.

The below two functions, time_torch_function and get_time_track_callback can be used for tracking the time taken for inference and training respectively.

The implementation of time_torch_function is taken from the PyTorch official torch.compile tutorial here .

We use get_time_track_callback to make a callback that can track the time taken for each batch during training as DeepChem does not provide a direct way to track the time taken per batch during training. We can use this callback by passing it as an argument to model.fit() function.

Tracking Training Time ¶

Looking at the graph, there is a significant difference in the time taken for the compiled and uncompiled versions of the model for the starting many steps. After that the time taken by the compiled model stabilises below the uncompiled model. This is because the compilation is done JIT when the model is first run and the optimized kernels are generated after a few passes.

Tracking Inference Time ¶

As with the results we got training, the first few runs for inference also takes significantly more time due to the same reason as mentioned before.

Congratulations! Time to join the Community! ¶

Congratulations on completing this tutorial notebook! If you enjoyed working through the tutorial, and want to continue working with DeepChem, we encourage you to finish the rest of the tutorials in this series. You can also help the DeepChem community in the following ways:

Star DeepChem on GitHub ¶

This helps build awareness of the DeepChem project and the tools for open source drug discovery that we're trying to build.

Join the DeepChem Discord ¶

The DeepChem Discord hosts a number of scientists, developers, and enthusiasts interested in deep learning for the life sciences. Join the conversation!