shrinath-suresh/blogs-docs-splitted huggingface dataset:
2.12
   
10.51
   
14.2
   



A10

-2.34
   
8.99
   
10.57
   


V100

18.63
   
6.39
   
10.43
   


A100

38.5
   
20.33
   
12.17
   


One can notice the following:

The improvements are significant for powerful GPUs like A100 and V100. For those GPUs the improvement is most pronounced for batch size 1
For less powerful GPUs we observe smaller speedups (or in two cases slight regressions). The batch size trend is reversed here: improvement is larger for larger batches

In the following sections we describe the applied optimizations and provide detailed benchmarking data, comparing the generation time with various optimization features on/off.
Optimizations
Here we’ll go into more detail about the optimizations introduced into the model code. These optimizations rely on features of PyTorch 2.0 which has been released recently. 
Optimized Attention
One part of the code which we optimized is the scaled dot-product attention. Attention is known to be a heavy operation: naive implementation materializes the attention matrix, leading to time and memory complexity quadratic in sequence length. It is common for diffusion models to use attention (CrossAttention) as part of Transformer blocks in multiple parts of the U-Net. Since the U-Net runs at every sampling step, this becomes a critical point to optimize. Instead of custom attention implementation one can use torch.nn.MultiheadAttention, which in PyTorch 2 has optimized attention implementation is integrated into it. This optimization schematically boils down to the following pseudocode:
```
class CrossAttention(nn.Module):
    def init(self, ...):
def init(self, ...):
        # Create matrices: Q, K, V, out_proj
        ...
    def forward(self, x, context=None, mask=None):
       # Compute out = SoftMax(Q*K/sqrt(d))V
       # Return out_proj(out)
       …

gets replaced with


class CrossAttention(nn.Module):
    def init(self, ...):
        self.mha = nn.MultiheadAttention(...)
    def forward(self, x, context):
    return self.mha(x, context, context)
```
return self.mha(x, context, context)
```
The optimized implementation of attention was available already in PyTorch 1.13 (see here) and widely adopted (see e.g. HuggingFace transformers library example). In particular, it integrates memory-efficient attention from the xFormers library and flash attention from https://arxiv.org/abs/2205.14135. PyTorch 2.0 expands this to additional attention functions such as cross attention and custom kernels for further acceleration, making it applicable to diffusion models.
Flash attention is available on GPUs with compute capability SM 7.5 or SM 8.x - for example, on T4, A10, and A100, which are included in our benchmark (you can check compute capability of each NVIDIA GPU here). However, in our tests on A100 the memory efficient attention performed better than flash attention for the particular case of diffusion models, due to the small number of attention heads and small batch size.  PyTorch understands this and in this case chooses memory efficient attention over flash attention when both are available (see the logic here). For full control over the attention backends (memory-efficient attention, flash attention, “vanilla math”, or any future ones), power users can enable and disable them manually with the help of the context manager torch.backends.cuda.sdp_kernel.
Compilation
Compilation is a new feature of PyTorch 2.0, enabling significant speedups with a very simple user experience. To invoke the default behavior, simply wrap a PyTorch module or a function into torch.compile:
model = torch.compile(model)

PyTorch compiler then turns Python code into a set of instructions which can be executed efficiently without Python overhead. The compilation happens dynamically the first time the code is executed. With the default behavior, under the hood PyTorch utilized TorchDynamo to compile the code and TorchInductor to further optimize it. See this tutorial for more details.
Although the one-liner above is enough for compilation, certain modifications in the code can squeeze a larger speedup. In particular, one should avoid so-called graph breaks - places in the code which PyTorch can’t compile. As opposed to previous PyTorch compilation approaches (like TorchScript), PyTorch 2 compiler doesn’t break in this case. Instead it falls back on eager execution - so the code runs, but with reduced performance. We introduced a few minor changes to the model code to get rid of graph breaks. This included eliminating functions from libraries not supported by the compiler, such as inspect.isfunction and einops.rearrange. See this doc to learn more about graph breaks and how to eliminate them.
Theoretically, one can apply torch.compileon the whole diffusion sampling loop. However, in practice it is enough to just compile the U-Net. The reason is that torch.compile doesn’t yet have a loop analyzer and would recompile the code for each iteration of the sampling loop. Moreover, compiled sampler code is likely to generate graph breaks - so one would need to adjust it if one wants to get a good performance from the compiled version.
Note that compilation requires GPU compute capability >= SM 7.0 to run in non-eager mode. This covers all GPUs in our benchmarks -  T4, V100, A10, A100 - except for P100 (see the full list). 
Other optimizations
Other optimizations
In addition, we have improved efficiency of GPU memory operations by eliminating some common pitfalls, e.g. creating a tensor on GPU directly rather than creating it on CPU and later moving to GPU. The places where such optimizations were necessary were determined by line-profiling and looking at CPU/GPU traces and Flame Graphs.
Benchmarking setup and results summary
We have two versions of code to compare: original and optimized. On top of this, several optimization features (xFormers, PyTorch memory efficient attention, compilation) can be turned on/off. Overall, as mentioned in the introduction, we will be benchmarking 5 configurations:

Original code without xFormers
Original code with xFormers
Optimized code with vanilla math attention backend and no compilation
Optimized code with memory-efficient attention backend and no compilation
Optimized code with memory-efficient attention backend and compilation

As the original version we took the version of the code which uses PyTorch 1.12 and a custom implementation of attention. The optimized version uses nn.MultiheadAttention in CrossAttention and PyTorch 2.0.0.dev20230111+cu117. It also has a few other minor optimizations in PyTorch-related code. 
The table below shows runtime of each version of the code in seconds, and the percentage improvement compared to the _original with xFormers. _The compilation time is excluded.
Runtimes for batch size 1. In parenthesis - relative improvement with respect to the “Original with xFormers” row



Configuration

P100

T4

A10

V100

A100
Original without xFormers

30.4s (-19.3%)
   
29.8s (-77.3%)
   
13.0s (-83.9%)
   
10.9s (-33.1%)
   
8.0s (-19.3%)
   


Original with xFormers

25.5s (0.0%)
   
16.8s (0.0%)
   
7.1s (0.0%)
   
8.2s (0.0%)
   
6.7s (0.0%)
   


Optimized with vanilla math attention, no compilation

27.3s (-7.0%)
   
19.9s (-18.7%)
   
13.2s (-87.2%)
   
7.5s (8.7%)
   
5.7s (15.1%)
   


Optimized with mem. efficient attention, no compilation

26.5s (-3.8%)
   
16.8s (0.2%)
   
7.1s (-0.8%)
   
6.9s (16.0%)
   
5.3s (20.6%)
5.3s (20.6%)
   



Optimized with mem. efficient attention and compilation

-
   
16.4s (2.1%)
   
7.2s (-2.3%)
   
6.6s (18.6%)
   
4.1s (38.5%)
   


Runtimes for batch size 2



Configuration

P100

T4

A10

V100

A100




Original without xFormers

58.0s (-21.6%)
   
57.6s (-84.0%)
   
24.4s (-95.2%)
   
18.6s (-63.0%)
   
12.0s (-50.6%)
   


Original with xFormers

47.7s (0.0%)
   
31.3s (0.0%)
47.7s (0.0%)
   
31.3s (0.0%)
   
12.5s (0.0%)
   
11.4s (0.0%)
   
8.0s (0.0%)
   



Optimized with vanilla math attention, no compilation

49.3s (-3.5%)
   
37.9s (-21.0%)
   
17.8s (-42.2%)
   
12.7s (-10.7%)
   
7.8s (1.8%)
   


Optimized with mem. efficient attention, no compilation

47.5s (0.4%)
   
31.2s (0.5%)
   
12.2s (2.6%)
   
11.5s (-0.7%)
   
7.0s (12.6%)
   


Optimized with mem. efficient attention and compilation

-
   
28.0s (10.5%)
   
11.4s (9.0%)
   
10.7s (6.4%)
   
6.4s (20.3%)
   


Runtimes for batch size 4
Runtimes for batch size 4



Configuration

P100

T4

A10

V100

A100




Original without xFormers

117.9s (-20.0%)
   
112.4s (-81.8%)
   
47.2s (-101.7%)
   
35.8s (-71.9%)
   
22.8s (-78.9%)
   


Original with xFormers

98.3s (0.0%)
   
61.8s (0.0%)
   
23.4s (0.0%)
   
20.8s (0.0%)
   
12.7s (0.0%)
   


Optimized with vanilla math attention, no compilation

101.1s (-2.9%)
   
73.0s (-18.0%)
   
28.3s (-21.0%)
   
23.3s (-11.9%)
23.3s (-11.9%)
   
14.5s (-13.9%)
   



Optimized with mem. efficient attention, no compilation

92.9s (5.5%)
   
61.1s (1.2%)
   
23.9s (-1.9%)
   
20.8s (-0.1%)
   
12.8s (-0.9%)
   


Optimized with mem. efficient attention and compilation

-
   
53.1s (14.2%)
   
20.9s (10.6%)
   
18.6s (10.4%)
   
11.2s (12.2%)
To minimize fluctuations and external influence on the performance of the benchmarked code, we ran each version of the code one after another, and then repeated this sequence 10 times: A, B, C, D, E,  A, B, … So the results of a typical run would look like the one in the picture below.. Note that one shouldn’t rely on comparison of absolute run times between different graphs, but comparison of run times_ inside_ one graph is pretty reliable, thanks to our benchmarking setup.
Each run of text-to-image generation script produces several batches, the number of which is regulated by the CLI parameter --n_iter. In the benchmarks we used n_iter = 2, but introduced an additional “warm-up” iteration, which doesn’t contribute to the run time. This was necessary for the runs with compilation, because compilation happens the first time the code runs, and so the first iteration is much longer than all subsequent. To make comparison fair, we also introduced this additional “warm-up” iteration to all other runs. 
The numbers in the table above are for number of iterations 2 (plus a “warm-up one”), prompt ”A photo”, seed 1, PLMS sampler, and autocast turned on.
Benchmarks were done using P100, V100, A100, A10 and T4 GPUs. The T4 benchmarks were done in Google Colab Pro. The A10 benchmarks were done on g5.4xlarge AWS instances with 1 GPU.
Conclusions and next steps
Conclusions and next steps
We have shown that new features of PyTorch 2 - compiler and optimized attention implementation - give performance improvements exceeding or comparable with what previously required installation of an external dependency (xFormers). PyTorch achieved this, in particular, by integrating memory efficient attention from xFormers into its codebase. This is a significant improvement for user experience, given that xFormers, being a state-of-the-art library, in many scenarios requires custom installation process and long builds.
There are a few natural directions in which this work can be continued: 

The optimizations we implemented and described here are only benchmarked for text-to-image inference so far. It would be interesting to see how they affect training performance. PyTorch compilation can be directly applied to training; enabling training with PyTorch optimized attention is on the roadmap
We intentionally minimized changes to the original model code. Further profiling and optimization can probably bring more improvements
At the moment compilation is applied only to the U-Net model inside the sampler. Since there is a lot happening outside of U-Net (e.g. operations directly in the sampling loop), it would be beneficial to compile the whole sampler. However, this would require analysis of the compilation process to avoid recompilation at every sampling step
Current code only applies compilation within the PLMS sampler, but it should be trivial to extend it to other samplers
Besides text-to-image generation, diffusion models are also applied to other tasks - image-to-image and inpainting. It would be interesting to measure how their performance improves from PyTorch 2 optimizations 

See if you can increase performance of open source diffusion models using the methods we described, and share the results! 
Resources
Resources

PyTorch 2.0 overview, which has a lot of information on torch.compile: https://pytorch.org/get-started/pytorch-2.0/ 
Tutorial on torch.compile: https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
General compilation troubleshooting: https://pytorch.org/docs/master/dynamo/troubleshooting.html
Details on graph breaks: https://pytorch.org/docs/master/dynamo/faq.html#identifying-the-cause-of-a-graph-break
Details on guards: https://pytorch.org/docs/master/dynamo/guards-overview.html
Video deep dive on TorchDynamo https://www.youtube.com/watch?v=egZB5Uxki0I
Tutorial on optimized attention in PyTorch 1.12: https://pytorch.org/tutorials/beginner/bettertransformer_tutorial.html 

Acknowledgements
We would like to thank Geeta Chauhan, Natalia Gimelshein, Patrick Labatut, Bert Maher, Mark Saroufim, Michael Voznesensky and Francisco Massa for their valuable advice and early feedback on the text.
Special thanks to Yudong Tao initiating the work on using PyTorch native attention in diffusion models.
layout: blog_detail
title: "New Library Updates in PyTorch 2.0"

Summary
We are bringing a number of improvements to the current PyTorch libraries, alongside the PyTorch 2.0 release. These updates demonstrate our focus on developing common and extensible APIs across all domains to make it easier for our community to build ecosystem projects on PyTorch. 
Along with 2.0, we are also releasing a series of beta updates to the PyTorch domain libraries, including those that are in-tree, and separate libraries including TorchAudio, TorchVision, and TorchText. An update for TorchX is also being released as it moves to community supported mode. Please find the list of the latest stable versions and updates below.
Latest Stable Library Versions (Full List)


TorchArrow 0.1.0
   
TorchRec 0.4.0
   
TorchVision 0.15
TorchVision 0.15
   



TorchAudio 2.0
   
TorchServe 0.7.1
   
TorchX 0.4.0
   


TorchData 0.6.0
   
TorchText 0.15.0
   
PyTorch on XLA Devices 1.14
   


*To see prior versions or (unstable) nightlies, click on versions in the top left menu above ‘Search Docs’.
TorchAudio
[Beta] Data augmentation operators
The release adds several data augmentation operators under torchaudio.functional and torchaudio.transforms:
* torchaudio.functional.add_noise
* torchaudio.functional.convolve
* torchaudio.functional.deemphasis
* torchaudio.functional.fftconvolve
* torchaudio.functional.preemphasis
* torchaudio.functional.speed
* torchaudio.transforms.AddNoise
* torchaudio.transforms.Convolve
* torchaudio.transforms.Deemphasis
* torchaudio.transforms.FFTConvolve
* torchaudio.transforms.Preemphasis
* torchaudio.transforms.Speed
torchaudio.transforms.Speed
torchaudio.transforms.SpeedPerturbation

The operators can be used to synthetically diversify training data to improve the generalizability of downstream models.
For usage details, please refer to the functional and transform documentation and Audio Data Augmentation tutorial.
[Beta] WavLM and XLS-R models
The release adds two self-supervised learning models for speech and audio.

WavLM that is robust to noise and reverberation.
XLS-R that is trained on cross-lingual datasets.

Besides the model architectures, torchaudio also supports corresponding pre-trained pipelines:

torchaudio.pipelines.WAVLM_BASE
torchaudio.pipelines.WAVLM_BASE_PLUS
torchaudio.pipelines.WAVLM_LARGE
torchaudio.pipelines.WAVLM_LARGE
torchaudio.pipelines.WAV2VEC_XLSR_300M
torchaudio.pipelines.WAV2VEC_XLSR_1B
torchaudio.pipelines.WAV2VEC_XLSR_2B

For usage details, please refer to the factory function and pre-trained pipelines documentation.
TorchRL
The initial release of torchrl includes several features that span across the entire RL domain. TorchRL can already be used in online, offline, multi-agent, multi-task and distributed RL settings, among others. See below:
[Beta] Environment wrappers and transforms
torchrl.envs includes several wrappers around common environment libraries. This allows users to swap one library with another without effort. These wrappers build an interface between these simulators and torchrl:

dm_control: 
Gym
Brax
EnvPool
Jumanji
Habitat
Gym
Brax
EnvPool
Jumanji
Habitat

It also comes with many commonly used transforms and vectorized environment utilities that allow for a fast execution across simulation libraries. Please refer to the documentation for more detail.
[Beta] Datacollectors
Data collection in RL is made easy via the usage of single process or multiprocessed/distributed data collectors that execute the policy in the environment over a desired duration and deliver samples according to the user’s needs. These can be found in torchrl.collectors and are documented here.
[Beta] Objective modules
Several objective functions are included in torchrl.objectives, among which: 

A generic PPOLoss class and derived ClipPPOLoss and KLPPOLoss
SACLoss and DiscreteSACLoss
DDPGLoss
DQNLoss
REDQLoss
A2CLoss
TD3Loss
ReinforceLoss
Dreamer
A2CLoss
TD3Loss
ReinforceLoss
Dreamer

Vectorized value function operators also appear in the library. Check the documentation here.
[Beta] Models and exploration strategies
We provide multiple models, modules and exploration strategies. Get a detailed description in the doc.
[Beta] Composable replay buffer
A composable replay buffer class is provided that can be used to store data in multiple contexts including single and multi-agent, on and off-policy and many more.. Components include:

Storages (list, physical or memory-based contiguous storages)
Samplers (Prioritized, sampler without repetition)
Writers
Possibility to add transforms

Replay buffers and other data utilities are documented here.
[Beta] Logging tools and trainer
We support multiple logging tools including tensorboard, wandb and mlflow.
We provide a generic Trainer class that allows for easy code recycling and checkpointing.
These features are documented here.
TensorDict
TensorDict is a new data carrier for PyTorch.
[Beta] TensorDict: specialized dictionary for PyTorch
TensorDict allows you to execute many common operations across batches of tensors carried by a single container. TensorDict supports many shape and device or storage operations, and  can readily be used in distributed settings. Check the documentation to know more.
[Beta] @tensorclass: a dataclass for PyTorch
Like TensorDict, tensorclass provides the opportunity to write dataclasses with built-in torch features such as shape or device operations. 
[Beta] tensordict.nn: specialized modules for TensorDict
The tensordict.nn module provides specialized nn.Module subclasses that make it easy to build arbitrarily complex graphs that can be executed with TensorDict inputs. It is compatible with the latest PyTorch features such as functorch, torch.fx and torch.compile.
TorchRec
[Beta] KeyedJaggedTensor All-to-All Redesign and Input Dist Fusion
We observed performance regression due to a bottleneck in sparse data distribution for models that have multiple, large KJTs to redistribute. 
To combat this we altered the comms pattern to transport the minimum data required in the initial collective to support the collective calls for the actual KJT tensor data. This data sent in the initial collective, ‘splits’ means more data is transmitted over the comms stream overall, but the CPU is blocked for significantly shorter amounts of time leading to better overall QPS.
Furthermore, we altered the TorchRec train pipeline to group the initial collective calls for the splits together before launching the more expensive KJT tensor collective calls. This fusion minimizes the CPU blocked time as launching each subsequent input distribution is no longer dependent on the previous input distribution.
With this feature, variable batch sizes are now natively supported across ranks. These features are documented here.
TorchVision
[Beta] Extending TorchVision’s Transforms to Object Detection, Segmentation & Video tasks
TorchVision is extending its Transforms API! Here is what’s new:

You can use them not only for Image Classification but also for Object Detection, Instance & Semantic Segmentation and Video Classification.
You can use new functional transforms for transforming Videos, Bounding Boxes and Segmentation Masks.
Learn more about these new transforms from our docs, and submit any feedback in our dedicated issue.
TorchText
[Beta] Adding scriptable T5 and Flan-T5 to the TorchText library with incremental decoding support!
TorchText has added the T5 model architecture with pre-trained weights for both the original T5 paper and Flan-T5. The model is fully torchscriptable and features an optimized multiheaded attention implementation. We include several examples of how to utilize the model including summarization, classification, and translation.
For more details, please refer to our docs.
TorchX
layout: blog_detail
title: "Deprecation of CUDA 11.6 and Python 3.7 Support"

For the upcoming PyTorch 2.0 feature release (target March 2023), we will target CUDA 11.7 as the stable version and CUDA 11.8 as the experimental version of CUDA and Python >=3.8, <=3.11. 
If you are still using or depending on CUDA 11.6 or Python 3.7 builds, we strongly recommend moving to at least CUDA 11.7 and Python 3.8, as it would be the minimum versions required for PyTorch 2.0.
Please note that as of Feb 1, CUDA 11.6 and Python 3.7  are no longer included in the nightlies
Please refer to the Release Compatibility Matrix for PyTorch releases:


PyTorch Version

Python

Stable CUDA

Experimental CUDA



2.0
   
>=3.8, <=3.11
   
CUDA 11.7, CUDNN 8.5.0.96
   
CUDA 11.8, CUDNN 8.7.0.84
1.13
   
>=3.7, <=3.10
   
CUDA 11.6, CUDNN 8.3.2.44
   
CUDA 11.7, CUDNN 8.5.0.96
   


1.12
   
>=3.7, <=3.10
   
CUDA 11.3, CUDNN 8.3.2.44
   
CUDA 11.6, CUDNN 8.3.2.44
   


As of 2/1/2023
For more information on PyTorch releases, updated compatibility matrix and release policies, please see (and bookmark) Readme.
layout: blog_detail
title: 'torchvision 0.3: segmentation, detection models, new datasets and more..'
author: Francisco Massa
redirect_from: /2019/05/23/torchvision03.html

PyTorch domain libraries like torchvision provide convenient access to common datasets and models that can be used to quickly create a state-of-the-art baseline. Moreover, they also provide common abstractions to reduce boilerplate code that users might have to otherwise repeatedly write. The torchvision 0.3 release brings several new features including models for semantic segmentation, object detection, instance segmentation, and person keypoint detection, as well as custom C++ / CUDA ops specific to computer vision.



New features include:
New features include:
Reference training / evaluation scripts: torchvision now provides, under the references/ folder, scripts for training and evaluation of the following tasks: classification, semantic segmentation, object detection, instance segmentation and person keypoint detection. These serve as a log of how to train a specific model and provide baseline training and evaluation scripts to quickly bootstrap research.
torchvision ops: torchvision now contains custom C++ / CUDA operators. Those operators are specific to computer vision, and make it easier to build object detection models. These operators currently do not support PyTorch script mode, but support for it is planned for in the next release. Some of the ops supported include:

roi_pool (and the module version RoIPool)
roi_align (and the module version RoIAlign)
nms, for non-maximum suppression of bounding boxes
box_iou, for computing the intersection over union metric between two sets of bounding boxes