![Revisit consent button](../images/body-df29beacd7496220cc94.svg) We value your privacy We use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking "Accept All", you consent to our use of cookies. Customise Reject All Accept All Customise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg) We use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below. The cookies that are categorised as "Necessary" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more NecessaryAlways Active Necessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data. * Cookie BCTempID * Duration 10 minutes * Description No description available. * Cookie __cf_bm * Duration 1 hour * Description This cookie, set by Cloudflare, is used to support Cloudflare Bot Management. * Cookie AWSALBCORS * Duration 7 days * Description Amazon Web Services set this cookie for load balancing. * Cookie _cfuvid * Duration session * Description Cloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services * Cookie li_gc * Duration 6 months * Description Linkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes. * Cookie __hssrc * Duration session * Description This cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session. * Cookie __hssc * Duration 1 hour * Description HubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie. * Cookie wpEmojiSettingsSupports * Duration session * Description WordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly. * Cookie BCSessionID * Duration 1 year 1 month 4 days * Description Blueconic sets this cookie as a unique identifier for the BlueConic profile. * Cookie _octo * Duration 1 year * Description No description available. * Cookie logged_in * Duration 1 year * Description No description available. * Cookie __Secure-YEC * Duration past * Description YouTube sets this cookie to stores the user's video player preferences using embedded YouTube video * Cookie __eoi * Duration 6 months * Description Description is currently not available. * Cookie AWSALBTGCORS * Duration 7 days * Description No description available. * Cookie login-status-p * Duration past * Description Description is currently not available. * Cookie AWSALBTG * Duration 7 days * Description No description available. * Cookie csrf_token * Duration session * Description No description available. * Cookie token_v2 * Duration 1 day * Description Description is currently not available. * Cookie D * Duration 1 year * Description Description is currently not available. * Cookie PHPSESSID * Duration session * Description This cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed. * Cookie VISITOR_PRIVACY_METADATA * Duration 6 months * Description YouTube sets this cookie to store the user's cookie consent state for the current domain. * Cookie cookietest * Duration session * Description The cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience. * Cookie __Host-airtable-session * Duration 1 year * Description This cookie is used to enable us to integrate the services of Airtable. * Cookie __Host-airtable-session.sig * Duration 1 year * Description This cookie is used to enable us to integrate the services of Airtable. * Cookie m * Duration 1 year 1 month 4 days * Description Stripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly. * Cookie BIGipServer* * Duration session * Description Marketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. * Cookie __cfruid * Duration session * Description Cloudflare sets this cookie to identify trusted web traffic. * Cookie _GRECAPTCHA * Duration 6 months * Description Google Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks. * Cookie __Secure-YNID * Duration 6 months * Description Google cookie used to protect user security and prevent fraud, especially during the login process. * Cookie cookieyes-consent * Duration 1 year * Description CookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors. Functional Functional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features. * Cookie lidc * Duration 1 day * Description LinkedIn sets the lidc cookie to facilitate data center selection. * Cookie brw * Duration 1 year * Description No description available. * Cookie brwConsent * Duration 5 minutes * Description Description is currently not available. * Cookie WMF-Uniq * Duration 1 year * Description Description is currently not available. * Cookie loom_anon_comment * Duration 1 year * Description No description available. * Cookie loom_referral_video * Duration session * Description Description is currently not available. * Cookie VISITOR_INFO1_LIVE * Duration 6 months * Description A cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface. * Cookie yt-remote-connected-devices * Duration Never Expires * Description YouTube sets this cookie to store the user's video preferences using embedded YouTube videos. * Cookie ytidb::LAST_RESULT_ENTRY_KEY * Duration Never Expires * Description The cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future. * Cookie yt-remote-device-id * Duration Never Expires * Description YouTube sets this cookie to store the user's video preferences using embedded YouTube videos. * Cookie yt-remote-session-name * Duration session * Description The yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video. * Cookie yt-remote-fast-check-period * Duration session * Description The yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos. * Cookie yt-remote-session-app * Duration session * Description The yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player. * Cookie yt-remote-cast-available * Duration session * Description The yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player. * Cookie yt-remote-cast-installed * Duration session * Description The yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video. * Cookie cp_session * Duration 3 months * Description Codepen sets this cookie for Help systems found in the website. * Cookie loid * Duration 1 year 1 month 4 days * Description This cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform. Analytics Analytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc. * Cookie __hstc * Duration 6 months * Description Hubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session). * Cookie hubspotutk * Duration 6 months * Description HubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts. * Cookie _ga * Duration 1 year 1 month 4 days * Description Google Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors. * Cookie _ga_* * Duration 1 year 1 month 4 days * Description Google Analytics sets this cookie to store and count page views. * Cookie __Host-psifi.analyticsTrace * Duration 6 hours * Description Description is currently not available. * Cookie __Host-psifi.analyticsTraceV2 * Duration 6 hours * Description Description is currently not available. * Cookie _gh_sess * Duration session * Description GitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form. * Cookie YSC * Duration session * Description YSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages. * Cookie ajs_anonymous_id * Duration 1 year * Description This cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before. * Cookie vuid * Duration 1 year 1 month 4 days * Description Vimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. Performance Performance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors. * Cookie AWSALB * Duration 7 days * Description AWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target. * Cookie acq * Duration past * Description Description is currently not available. * Cookie acq.sig * Duration past * Description Description is currently not available. * Cookie ptc * Duration 2 years * Description No description available. Advertisement Advertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns. * Cookie muc_ads * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to collect user behaviour and interaction data to optimize the website. * Cookie guest_id_marketing * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to identify and track the website visitor. * Cookie guest_id_ads * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to identify and track the website visitor. * Cookie personalization_id * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting. * Cookie guest_id * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences. * Cookie bcookie * Duration 1 year * Description LinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs. * Cookie __Secure-ROLLOUT_TOKEN * Duration 6 months * Description YouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment. * Cookie yt.innertube::nextId * Duration Never Expires * Description YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen. * Cookie yt.innertube::requests * Duration Never Expires * Description YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen. * Cookie session_tracker * Duration session * Description This cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website. * Cookie edgebucket * Duration session * Description Reddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content. * Cookie did * Duration 1 year * Description Arbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year. Uncategorised Other uncategorised cookies are those that are being analysed and have not been classified into a category as yet. No cookies to display. Reject All Save My Preferences Accept All Toggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation [](https://towardsdatascience.com/) Publish AI, ML & data-science insights to a global community of data professionals. Sign in[Submit an Article](https://write.insightmediagroup.io) [Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/) * * * [Write For TDS](https://towardsdatascience.com/submissions/) [LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search [Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/) # Maximizing AI/ML Model Performance with PyTorch Compilation Practical Tips for Getting the Most Out of torch.compile [Chaim Rand](https://towardsdatascience.com/author/chaimrand/) August 18, 202526 min read ![](../images/000.webp) [Photo](https://unsplash.com/photos/red-boat-near-mosque-painting--1h_NN3nqzI) by Mg Cthu on Unsplash Since its inception in [PyTorch 2.0](https://towardsdatascience.com/tips-and-tricks-for-upgrading-to-pytorch-2-3127db1d1f3d/) in March 2023, the evolution of torch.compile has been one of the most exciting things to follow. Given that PyTorch’s popularity was due to its “Pythonic” nature, its ease of use, and its line-by-line (a.k.a., eager) execution, the success of a just-in-time (JIT) graph compilation mode should not have been taken for granted. And yet, just over two years later, the importance of this feature cannot be overstated: It is an essential tool in optimizing the runtime performance of AI/ML workloads. Unfortunately, the use of torch.compile still feels a bit like a dark art. When it works it is awesome and everyone is happy. However, when it doesn’t, figuring out the reason can be difficult. It has multiple API controls, but knowing which ones to apply and when — can seem like black magic. Moreover, its documentation is currently somewhat decentralized, with the details of many of its key features scattered across multiple posts and tutorials. Although covered in a [previous post](https://medium.com/data-science/tips-and-tricks-for-upgrading-to-pytorch-2-3127db1d1f3d), we felt that the rapid evolution of torch.compile warranted a renewed discussion. This post attempts to unveil some of the mystique surrounding torch.compile. We will review how it works, demonstrate its use, discuss a few strategies for how to apply it most effectively, and evaluate the impact of some of its features on the runtime performance of a toy model. We will cover the following topics: Learn this step by step with the interactive [Machine Learning](https://roadmap.sh/machine-learning) roadmap. * techniques for avoiding the two “compilation-killers”, graph-breaks and recompilations, * strategies for debugging compilation issues * squeezing maximum performance using some of torch.compile’s advanced features and configuration settings, * making the most of the torch.compile logs to debug compilation issues, * modular application of torch.compile, * methods for reducing compilation time, * and more. As in our previous posts, we will define a toy PyTorch model which we will use to demonstrate the application and impact of torch.compile. We will run our experiments on an [Amazon EC2 p4d.96xlarge](https://aws.amazon.com/ec2/instance-types/p4/) instance (containing 8 NVIDIA A100 GPUs) running a [PyTorch (2.7) Deep Learning AMI (DLAMI)](https://docs.aws.amazon.com/dlami/latest/devguide/aws-deep-learning-ami-gpu-pytorch-2.7-ubuntu-22-04.html). ### Disclaimers: PyTorch compilation is a complex topic with a continuously growing set of features. This post makes no attempt to encompass the full scope of torch.compile, but rather aims to offer some practical tips on how to approach it. For a complete reference, please see the [official PyTorch documentation](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html). But keep in mind that you may need to surf through multiple pages to collect all the information you need (e.g., [here](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) for the API documentation, [here](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) for an introductory tutorial, [here](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_deepdive.html) for a deep-dive on TorchDynamo, [here](https://docs.pytorch.org/tutorials/compilers_index.html) and [here](https://docs.pytorch.org/docs/stable/torch.compiler.html#read-more) for indices to many other pages covering a wide range of compilation features, etc.). If you prefer a single source with a comprehensive overview of torch.compile, its inner workings, and detailed examples of its use, we recommend chapter 14 of the book [AI Systems Performance Engineering](https://github.com/cfregly/ai-performance-engineering/tree/main), by Chris Fregly. The code we will share is intended for demonstrative purposes and should not be relied on for correctness or optimality — especially for other projects. Please do not interpret our choice of platform, framework, or any other tool or library as an endorsement for its use. The impact of torch.compile can vary greatly based on the details of the AI/ML model and runtime environment. The results we will share on our toy model may not be indicative of the results you will get on your own model. In fact, compilation of some models may result in worse performance. When applied correctly, torch.compile should not affect the quality of your model (in the case of inference) or its ability to converge (in the case of training). However, there are likely to be numerical differences due to the use of different compute kernels. It is essential that you verify that applying torch.compile does not degrade your quality-performance metrics before deploying it to a production environment. Importantly, torch.compile continues to evolve with each PyTorch release. The contents of this post are based on PyTorch 2.7. Staying up-to-date with latest PyTorch releases is essential for taking advantage of the latest and greatest available optimization opportunities. ## PyTorch Compilation: How it Works In PyTorch’s default _eager execution_ mode, each line of Python code is processed independently. While this mode of execution is extremely user-friendly — making it easy to follow and debug line-per-line what the model is doing — it misses a great deal of opportunity to optimize performance, e.g.: 1. GPU operations are performed independently. This misses the opportunity for _operator fusion_ where GPU operations are combined into a single, more efficient, GPU kernel. 2. Potential optimizations from ahead-of-time (AOT) compilation, such as out-of-order execution and memory layout optimizations, are missed. 3. The Python runtime is involved in all stages of the model execution. Every time an operation is launched on the GPU, control is passed from the Python interpreter to the CUDA backend and back. This can introduce significant overhead. ### How torch.compile Fixes This First introduced in PyTorch 2.0, torch.compile acts as a just-in-time (JIT) compiler: The first time you call a compiled function, the compiler traces the Python code and converts it into an intermediate graph representation (IR) using [TorchDynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html), sometimes referred to as an [FX Graph](https://docs.pytorch.org/docs/stable/fx.html#torch.fx.Graph). If the compiled function requires backpropagation, the FX Graph is passed to the [AOTAutograd](https://docs.pytorch.org/functorch/stable/notebooks/aot_autograd_optimizations.html) library which captures the backward pass ahead-of-time (AOT) and generates a combined forward and backward graph. The FX Graph is then passed to the compiler backend which performs kernel fusion, out-of-order execution, and other techniques to generate machine code that is highly optimized for the target hardware. The default PyTorch compiler backend is [TorchInductor](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747) which supports both GPU and CPU targets. When compiling for NVIDIA GPUs, TorchInductor uses: 1) the [Triton](https://openai.com/index/triton/) compiler (previously covered in [this post](https://towardsdatascience.com/unleashing-the-power-of-triton-mastering-gpu-kernel-optimization-in-python-160a3f52701e/)) to create optimal GPU kernels and 2) [CUDA Graphs](https://developer.nvidia.com/blog/cuda-graphs/) (whenever possible) to combine multiple GPU kernels into efficient, re-playable sequences. The final, machine-specific computation graph is [cached](https://docs.pytorch.org/tutorials/recipes/torch_compile_caching_configuration_tutorial.html) and used for each subsequent invocation of the compiled function/model. Note that although the bulk of the compilation is performed on the first invocation, several additional _warm-up_ passes are often required to reach peak performance. The combined JIT and AOT properties of torch.compile allow it to maximize opportunities for graph optimization, while the use of the compiled execution graph avoids the line-by-line involvement of the Python interpreter — thereby addressing the three aforementioned inefficiencies of eager execution. ## Avoiding Compilation Pitfalls Usually, applying torch.compile will boost your model throughput (e.g., see the [TorchInductor performance dashboard](https://hud.pytorch.org/benchmark/compilers)). However, sometimes you may find that torch compilation results in the same or even worse performance than in eager mode. There could be a number of reasons for this: 1. There may be a bottleneck in the training step that is overshadowing the torch.compile optimization, e.g., a data input pipeline bottleneck. This can be identified and solved through appropriate performance analysis and optimization. 2. Your function or model might already be so efficient that the application of torch.compile is negligible. 3. You may be suffering from one of two compilation killers, graph-breaks and recompilations, which we elaborate on in the next sections. ### PyTorch Compilation Killer #1: Graph-Breaks Graph-breaks are one of the most common events that interfere with efficient torch compilation. Graph-breaks occur when the TorchDynamo or AOTAutograd libraries encounter Python operations that they cannot convert into a graph operation. In such cases, the sections of code before and after the problematic operation, are compiled separately and the resultant graph is said to contain a _graph-break_. Graph-breaks interfere with the compiler’s capacity for optimization in two primary ways: First, optimizations such as kernel fusion cannot be performed across graph breaks and, second, a graph break implies a return of control to the Python interpreter. The presence of a large number of graph breaks can completely cancel out the potential benefit of torch.compile. Common examples of graph breaks include print() operations, conditional logic, and asserts. What is frustrating is that, more often than not, graph-breaks can be easily avoided. What is even more frustrating is that the default behavior is to handle graph breaks by silently falling back to eager execution for the problematic code segment. ### Avoiding Graph-Breaks The first step to handling graph-breaks is to configure the compiler to report them. Here are several ways of doing this: 1. Apply the torch._dynamo.explain operator to your (uncompiled) model and run it on a sample input (as demonstrated [here](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html#torchdynamo-and-fx-graphs)). This will result in a log containing a list of all of the graph-breaks. 2. Set the _TORCH_LOGS_ environment variable to include “ _graph_breaks_ ”. This will cause the compiler to print the graph-breaks it encounters during compilation. 3. Call with torch.compile with _fullgraph_ =_True_. This will cause the compilation to fail each time it encounters a graph-break — thereby forcing the developer to acknowledge its presence and potentially fix it. While our personal preference is option three, it is important to note that there are times where graph-breaks cannot be avoided, which means that we may need to disable _fullgraph_ in a production setting. The best example of this is distributed training (e.g., [DDP](https://docs.pytorch.org/tutorials/intermediate/ddp_tutorial.html) and [FSDP](https://docs.pytorch.org/docs/stable/fsdp.html)) where the computation group includes communication calls which (as of the time of this writing) are not supported by torch.compile and, thus, result in graph-breaks. With knowledge of the location of our graph breaks, we address each one individually. We remove redundant prints and assertions, replace conditional blocks with graph-friendly alternatives such as [torch.where](https://docs.pytorch.org/docs/stable/generated/torch.where.html) or [torch.cond](https://docs.pytorch.org/docs/stable/generated/torch.cond.html), and adjust our model implementation to minimize untraceable Python control flow and native operations. In some cases, we may desire to maintain some of the prints or assertions for running in eager mode; in this case, we can wrap them in a conditional check like `if not torch.compiler.is_compiling()`. There may be cases (e.g., DDP) where graph-breaks are unavoidable. See [here](https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#resolving-graph-breaks) for more on avoiding graph-breaks. ### PyTorch Compilation Killer #2: Recompilations The second potential compilation killer is the graph recompilation. During the initial graph compilation phase, several assumptions are made and relied upon for generating the resultant graph. In torch.compile lingo these assumptions are referred to as _guards_. Common guards include the data types and shapes of input tensors. On each iteration, these guards are verified on the current tensor inputs and training state. If one of the guards is violated, the current graph is deemed invalid for the current state and a new graph is generated, i.e., the graph is recompiled. Graph compilation takes an extremely long time relative to the time it takes to execute a compiled graph. Consequently, multiple recompilations is likely to erase any potential performance gains from torch.compile. Moreover, torch.compile has a recompilation limit (the default is 8) after which it will raise a torch._dynamo.exc.RecompileLimitExceeded exception and fall back to eager mode. ### Avoiding Recompiles Here too, the first step is identifying the causes of the recompilations. Once again, there are several options: 1. Use [torch_compiler.set_stance](https://docs.pytorch.org/tutorials/recipes/torch_compiler_set_stance_tutorial.html) operator to fail on recompile: `torch.compiler.set_stance(“fail_on_recompile”)`. In practice, this option can sometimes prove to be too limiting. 2. Set the _TORCH_LOGS_ environment variable to include “ _recompiles_ ”. This will cause the compiler to report each time it performs recompilation along with the guards that were violated. ### Compiling Graphs with Variable-Shaped Tensors One of the most common causes of recompilations is the presence of tensors with dynamic shapes. The first time a graph is compiled it creates guards according to the shapes of the tensors it traced. When a tensor changes shape in a subsequent step, the guard is violated and the graph is recompiled. There are multiple ways of handling tensors with dynamic shapes: 1. **Default Compilation Behavior** : If the _dynamic_ field of the torch.compile call is not set (or set to _None_), each time the compiler encounters a new dynamic tensor, it will perform recompilation to generate a new graph that supports the dynamism it identified. In this option, the graph modification is applied surgically, allowing for “static” optimizations to be applied to other portions of the graph. If new dynamism is discovered in multiple iterations, we may hit the recompilation limit and fall back to eager execution. Consequently, this option should only be used for models with limited dynamism. 2. **Mark Dynamic Tensors** : Another option is to explicitly mark the dynamic tensors and associated dynamic axis using the torch._dynamo.mark_dynamic API. This informs the compiler to build a graph that supports the reported dynamism and prevents recompilations altogether. This is a great option in situations in which you know upfront what your dynamic shapes are (which you absolutely should!!). 3. **Dynamic Compilation** : The third option is to apply torch.compile with _dynamic_ =_True_. This instructs the compiler to construct a graph that is as dynamic as possible in order to avoid recompilations. When enabled, dynamic shape tracing is applied to all of the tensors in the graph. This is often overkill. Keep in mind that many graph optimization techniques (e.g., CUDA graphs) assume static shapes. These are automatically disabled when this setting is applied. This option should be avoided whenever possible. 4. **Generate a Limited Number of Static Graphs** : When torch.compile is applied with _dynamic_ =_False,_ the compiler will never generate dynamic graphs. Each time a guard is violated a new static graph is created, supporting the newly encountered tensor shape, and added to the compilation cache. While limited (by the recompilation limit) in the number of shapes it can support, this option is compelling due to the fact that it allows for optimizations that assume a static graph. To benefit from this capability, a common approach is to remove dynamism from the graph by padding dynamic tensors to a fixed length. A more advanced approach that reduces the amount of padding is to set a number of fixed length values (e.g., powers of two) and pad the variable shaped tensors to the closest length. The number of length values should not exceed the recompilation limit. This will result in a fixed number of recompilations and a fixed number of highly optimized graphs. We can ensure that all graphs are created during the model warmup phase. As before, there are some situations where graph recompilations cannot be avoided, and we may have no choice but to run our model in eager mode. See [here](https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#dealing-with-recompilations) for more on avoiding recompilations and [here](https://docs.pytorch.org/docs/stable/torch.compiler_dynamic_shapes.html) for details on how torch.compile handles dynamic shapes. ## Debugging Compilation Issues Inevitably, you will encounter some situations where torch compilation fails. Often, you will get a long error message and callstack, but it may as well be in a foreign language. You will likely be encouraged to _Set TORCH_LOGS=”+dynamo” and TORCHDYNAMO_VERBOSE=1_ but you may find that this does little to help you solve the problem. The [torch.compile troubleshooting](https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html) guide offers multiple tips for diagnosing compilation errors (e.g., by compiling with “eager”, “aot_eager” and “inductor” backends), for fixing or avoiding them, and if all else fails, for [reporting](https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#reporting-issues) them to PyTorch. In this post we call out two different approaches for tackling tough compilation issues. ### Top-Down VS. Bottom-Up Approach In a top-down approach, we apply torch compilation at the highest-level function/model — come what may. We then begin to work through the compilation issues as they come up by either fixing them or removing them from the graph via the [torch.compiler.disable](https://docs.pytorch.org/docs/stable/generated/torch.compiler.disable.html) utility. This approach assumes that we are sufficiently able to decipher the compilation logs — at least well enough to navigate to the problematic line of code. In a bottom-up approach, we begin by applying compilation to a few low-level components and slowly increase the scope of compilation until we hit an error. This approach makes it easy to pinpoint the sources of the compilation issue. An additional advantage is that we can benefit from the results of a partially compiled graph while we continue to work on additional optimizations. This is contrary to the Top-Down approach where we will only have a workable graph once all issues are addressed. The best approach depends on the model at hand and your personal inclination.Often, a combination of the two delivers the best results: for example, identifying issues via a bottom-up approach, resolving them, and then testing if the full graph compilation works. ## Tuning for Maximal Performance Once you have succeeded in compiling your model, there are a wide range of controls for trying to squeeze out even better performance. In this section we will cover some of the available options. It should be noted that the additional performance gains from these options are usually a small fraction of the gains from the initial application of standard compilation. ### Advanced Compiler Modes and Options The [torch.compile API](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) allows for tuning the compiler-backend behavior via via the `mode` and `options` parameters. There are dozens of knobs that can be applied and assessed. Some of the most notable ones are “reduce-overhead” — that optimizes more aggressively to further reduce the overhead of the kernel loading and Python interpreter, and “max-autotune” — the most aggressive optimization option that performs benchmarking of multiple kernel options before choosing the most efficient one. Both of these, particularly “max-autotune”, increase the compilation time, but usually result in more efficient graphs. ### Varying the Compiler Backend The default compiler backend is TorchInductor which supports a variety of target devices. You can specify the compiler backend via the _backend_ parameter of the [torch.compile API](https://docs.pytorch.org/docs/stable/generated/torch.compile.html). While other backends are unlikely to beat TorchInductor when running on NVIDIA GPUs, you may find them to perform better on other hardware devices (e.g., the [ _ipex_](https://github.com/intel/intel-extension-for-pytorch) backend includes optimizations that leverage the unique capabilities of Intel® CPUs). ### Applying Modular Compilation While it is usually recommended to apply compilation to the entire model, there are times where the model can be broken into submodules that respond very differently to the compiler controls. For example, if your model contains one component that includes many tensors with dynamic shapes and another component that is static, you may find that compiling the first in “max-autotune-no-cudagraphs” mode and the second in “max-autotune” mode, results in maximum performance. ### Compiling the Optimizer In addition to compiling the model execution, as of PyTorch 2.2, you can further optimize your training workload by [compiling the optimizer](https://docs.pytorch.org/tutorials/recipes/compiling_optimizer.html). This will be demonstrated below. ### New Compiler Features Since the initial release of torch.compile in PyTorch 2.0, each PyTorch release has included enhancements to the torch.compile offering. Sometimes released as “prototypes”, new features offerings challenge developers to extract even greater performance out of graph compilation. For example, the [PyTorch 2.7 release](https://pytorch.org/blog/pytorch-2-7/) included the [foreach_map](https://docs.pytorch.org/tutorials/recipes/foreach_map.html) prototype feature, the use of which we will demonstrate below. ## Reducing Compilation Time While the initial compilation and warm-up time can be quite long compared to the subsequent training steps, it is usually negligible compared to the overall life of the model (i.e., the training or inference time). In some cases, however, the lengthy compilation time can become an issue. If the model is extremely large and we are tuning for optimal performance, compilation could take hours. If we are using our model in an inference server setup, the model start-up time could have a direct impact on the server response time and user experience. In this section we cover two techniques for reducing model compilation time: [compile-time caching](https://docs.pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html) and [regional compilation](https://docs.pytorch.org/tutorials/recipes/regional_compilation.html). ### Compile Time Caching In compile-time caching we upload the results of the local graph compilation to persistent storage. Every time we need to run the same model in the same runtime environment (e.g., same hardware and same library versions) we pull the cache state from persistent storage to the local disk, instead of compiling from scratch. ### Regional Compilation Regional compilation relies on the fact that large models typically consist of computation blocks that are repeated multiple times. In regional compilation, torch.compile is applied to the repeating block, instead of the entire model. The result is a single, relatively small graph that is created and reused for each of the blocks. ## How to Configure the TORCH_LOGS Environment Variable Torch compilation supports a wide variety of [logging controls](https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#summary-of-torch-logs-options). While the log reports can be extremely useful for debugging issues and maximizing performance, it’s important to find the right balance where the logs are helpful but not excessive. In this post we propose using the following initial configuration and adapting as needed: bash ``` export TORCH_LOGS="graph_breaks,recompiles,perf_hints" ``` * “graph_breaks” — reports each time a graph-break is encountered (see above) * “recompiles” — reports each time a recompilation is performed along with the guard-violation that triggered it. * “perf_hints” — outputs performance logs from the inductor backend along with hints for additional optimizations Note that sometimes “perf_hints” will flood the console with unactionable messages, in which case you may opt to disable it. ## A Toy PyTorch Model: Image Captioning To demonstrate torch.compile in action, we define a toy image captioning model using the popular [Hugging Face](https://huggingface.co/) [transformers](https://huggingface.co/docs/transformers/en/index) library (version 4.54.1). Specifically, we define an image-to-text model using a [VisionEncoderDecoderModel](https://huggingface.co/docs/transformers/v4.53.3/en/model_doc/vision-encoder-decoder#transformers.VisionEncoderDecoderModel), with a [Vision Transformer (ViT)](https://huggingface.co/papers/2010.11929) encoder and a [GPT-2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) decoder, and train it on a synthetic dataset of fixed-sized images and random sequences (“captions”) of variable length. We begin by defining our image-to-text model: python ``` import os, shutil, time, random, torch from transformers import ( VisionEncoderDecoderModel, VisionEncoderDecoderConfig, AutoConfig ) torch.manual_seed(42) random.seed(42) BATCH_SIZE = 64 NUM_WORKERS = 12 NUM_TOKENS = 1024 MAX_SEQ_LEN = 256 PAD_ID = 0 START_ID = 1 END_ID = 2 # set up image-to-text model def get_model(): config = VisionEncoderDecoderConfig.from_encoder_decoder_configs( encoder_config=AutoConfig.for_model("vit"), # vit encoder decoder_config=AutoConfig.for_model("gpt2") # gpt2 decoder ) config.decoder.vocab_size = NUM_TOKENS config.decoder.use_cache = False config.decoder_start_token_id = START_ID config.pad_token_id = PAD_ID config.eos_token_id = END_ID config.max_length = MAX_SEQ_LEN model = VisionEncoderDecoderModel(config=config) # remove unused pooler model.encoder.pooler = None # uncomment to specify the loss function # from transformers.loss.loss_utils import ForCausalLMLoss # model.loss_function = ForCausalLMLoss return model ``` Next, we define a synthetic dataset that generates pairs of random images of fixed size and random sequences of variable size. We use a weighted distribution for the sequence length to mimic a scenario where the vast majority of sequences are short. Given the varying length of the input captions, we require a strategy for dealing with dynamically shaped input. Here, we offer two alternatives, both of which use padding: padding to the maximum input length and padding to the length of the longest sequence in the batch, along with an option to align it to a given multiple. Please see our [previous post](https://towardsdatascience.com/optimizing-transformer-models-for-variable-length-input-sequences-19fb88fddf71/) for additional strategies for handling variable-length input sequences. python ``` from torch.utils.data import Dataset, DataLoader from functools import partial # A synthetic Dataset with random images and captions class FakeDataset(Dataset): def __init__(self): self.length_dist = { 'short': {'range': (5, 32), 'weight': 0.90}, 'medium': {'range': (33, 64), 'weight': 0.09}, 'long': {'range': (65, 256), 'weight': 0.01} } super().__init__() def __len__(self): return 1000000 def __getitem__(self, index): length_bin = random.choices( list(self.length_dist.keys()), weights=[d['weight'] for d in self.length_dist.values()], k=1 )[0] range_start, range_end = self.length_dist[length_bin]['range'] image = torch.randn(3, 224, 224) length = random.randint(range_start, range_end - 1) labels = torch.cat([torch.randint(1, NUM_TOKENS, (length,)), torch.tensor([END_ID])], dim=0) input_ids = torch.cat([torch.tensor([START_ID]), labels[:-1]], dim=0) return { 'image': image, 'input_ids': input_ids, 'labels': labels } def pad_sequence(sequence, length, pad_val): return torch.nn.functional.pad( sequence, (0, length - sequence.shape[0]), value=pad_val ) def collate_with_padding(batch, pad_to_longest=False, align=None): padded_inputs = [] padded_labels = [] if pad_to_longest: pad_len = max([b['input_ids'].shape[0] for b in batch]) if align: pad_len = ((pad_len + align - 1) // align) * align else: pad_len = MAX_SEQ_LEN for b in batch: input_ids = b['input_ids'] labels = b['labels'] padded_inputs.append(pad_sequence(input_ids, pad_len, PAD_ID)) padded_labels.append(pad_sequence(labels, pad_len, -100)) padded_inputs = torch.stack(padded_inputs, dim=0) padded_labels = torch.stack(padded_labels, dim=0) images = torch.stack([b['image'] for b in batch], dim=0) return { 'pixel_values': images, 'decoder_input_ids': padded_inputs, 'labels': padded_labels, 'decoder_attention_mask': (padded_inputs != PAD_ID) } def get_dataloader(pad_to_longest=False, align=None): return DataLoader( dataset=FakeDataset(), batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, collate_fn=partial( collate_with_padding, pad_to_longest=pad_to_longest, align=align ) ) ``` Last, we define our training step and main training function: python ``` def copy_to_device(batch, device): return { key: val.to(device=device, non_blocking=True) for key, val in batch.items() } def train_step(model, device, optimizer, batch): # copy data to device batch = copy_to_device(batch, device) optimizer.zero_grad() with torch.amp.autocast('cuda', dtype=torch.bfloat16): outputs = model(**batch) loss = outputs.loss loss.backward() optimizer.step() return loss def train(local_rank=0, world_size=1, compile=False): # specify log settings torch._logging.set_logs( graph_breaks=True, recompiles=True, perf_hints=True ) torch.cuda.set_device(local_rank) device = torch.cuda.current_device() if world_size > 1: # DDP setup import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = str(2222) dist.init_process_group('nccl', rank=local_rank, world_size=world_size) # configure pad_to_longest and optional alignment dataloader = get_dataloader(pad_to_longest=False, align=None) model = get_model() model = model.to(device) if world_size > 1: model = DDP(model, [local_rank]) optimizer = torch.optim.Adam(model.parameters()) if compile: # uncomment to run pre-compile warmup - required for some optimizations # batch = next(iter(dataloader)) # train_step(model, device, optimizer, batch) model, optimizer = apply_compilation(model, optimizer) warmup = 20 active = 100 total_steps = warmup + active t0 = time.perf_counter() for idx, batch in enumerate(dataloader, start=1): # apply train step train_step(model, device, optimizer, batch) if idx == warmup: torch.cuda.synchronize() print(f'warmup time: {time.perf_counter()-t0}') t0 = time.perf_counter() elif idx == total_steps: break if local_rank == 0: torch.cuda.synchronize() total_time = time.perf_counter() - t0 print(f'average throughput: {active / total_time}') if world_size > 1: dist.destroy_process_group() if __name__ == '__main__': # specify inductor cache dir inductor_cache_dir = '/tmp/inductor_cache' os.environ['TORCHINDUCTOR_CACHE_DIR'] = inductor_cache_dir # clean up compiler cache torch._dynamo.reset() shutil.rmtree(inductor_cache_dir, ignore_errors=True) world_size = 1 torch.multiprocessing.spawn( fn=train, args=(world_size,), nprocs=world_size, join=True ) ``` ### Baseline Performance Running the training script without compilation yields the following baseline performance results: ![](../images/001.webp) Baseline Model Performance (by Author) We can see clearly that the collation strategy that reduces padding results in much better performance. ## Applying Model Compilation In this section we will apply torch compilation with different configurations and measure its impact on the training throughput. We will begin by applying compilation without dynamism, i.e., when padding all inputs to max sequence length. In the following section we will evaluate its impact in the case of inputs with dynamic shapes. ### Model Compilation Step #1: Fixing Graph Breaks We introduce the following compilation utility function and apply it to our model: python ``` def apply_compilation(model, optimizer): model = torch.compile(model, fullgraph=True) return model, optimizer ``` The _fullgraph_ setting ensures that compilation will fail whenever it encounters a graph break. Sure enough, our first compilation attempt results in an error coming from the transformer library. Here is a small snippet: shell-session ``` from user code: File "/opt/pytorch/lib/python3.12/site-packages/transformers/models/vision_encoder_decoder/modeling_vision_encoder_decoder.py", line 574, in forward loss = self.loss_function( File "/opt/pytorch/lib/python3.12/site-packages/transformers/modeling_utils.py", line 5776, in loss_function ``` The reason for this error is that when the [VisionEncoderDecoderModel](https://huggingface.co/docs/transformers/v4.53.3/en/model_doc/vision-encoder-decoder#transformers.VisionEncoderDecoderModel) loss function is not specified, the transformers library uses native Python code to determine what loss function to apply. This is easy to fix by specifying the model loss function, as follows: python ``` from transformers.loss.loss_utils import ForCausalLMLoss model.loss_function = ForCausalLMLoss ``` Following this fix, model compilation succeeds. The resultant throughput is 5.17 steps per second — a 66% speed-up over the baseline (fixed-input) throughput. Note that in the current scenario of a static graph, the compiler did not report any recompilations, but it did report the following _perf_hint_ : shell-session ``` I0805 13:37:52.406000 51587 torch/_inductor/codegen/simd.py:1976] [0/0] [__perf_hints] Reduction over non-contiguous dims.I0805 13:37:52.406000 51587 torch/_inductor/codegen/simd.py:1976] [0/0] [__perf_hints] Consider setting config.triton.tile_reductions to True. ``` However, applying the suggested configuration results in a compilation error, so we ignore it going forward. ### Model Compilation Step #2: Tuning the Compiler Configuration Let’s try to increase the performance further by applying some of the advanced compilation controls. The code block below includes three alternative modifications: python ``` # reduce-overhead model = torch.compile(model, fullgraph=True, mode="reduce-overhead") # max-autotune model = torch.compile(model, fullgraph=True, mode="max-autotune") # shapes padding model = torch.compile(model, fullgraph=True, options={"shape_padding":True}) ``` The results are captured in the table below: ![](../images/body-71769591bb9c75f4216b.webp) torch.compile results (by Author) The subsequent experiments in this section will be run with the “max-autotune” optimization. ### Model Compilation Step #3: Compiling the Optimizer Next, we extend our solution to apply compilation to the optimizer. Since optimizer compilation currently requires graph-breaks, we apply it without the _fullgraph_ flag: python ``` def apply_compilation(model, optimizer): model = torch.compile(model, fullgraph=True, mode="max-autotune") optimizer.step = torch.compile(optimizer.step) return model, optimizer ``` Compiling the optimizer further increases the throughput to 5.54 steps per second!! When compiling the optimizer, the following performance hint is printed: shell-session ``` will be copied during cudagraphs execution.If using cudagraphs and the grad tensor addresses will be the same across runs, use torch._dynamo.decorators.mark_static_address to elide this copy. ``` The proposal is to fix the addresses of gradient tensors and mark them. To implement the suggestion, we introduce the following two utility functions: python ``` # this replaces default optimizer.zero_grad() and verifies reuse # of same gradient tensors def zero_grads(model): for p in model.parameters(): if p.grad is not None: p.grad.zero_() # uses dynamo utility to mark each of the gradient tensors as static def mark_static_address(optimizer): for group in optimizer.param_groups: for p in group['params']: if p.grad is not None: torch._dynamo.mark_static_address(p.grad) ``` The updated training step appears below: python ``` def train_step(model, device, optimizer, batch): # copy data to device batch = copy_to_device(batch, device) zero_grads(model) with torch.amp.autocast('cuda', dtype=torch.bfloat16): outputs = model(**batch) loss = outputs.loss loss.backward() mark_static_address(optimizer) optimizer.step() return loss ``` In our case, implementing the performance hint decreases the throughput result to 5.32 steps per second — so we disregard it. ### Model Compilation Step #4: Foreach Map Optimization Constantly be on the lookout for torch.compile enhancements and additions. Here we will apply [horizontal fusion with foreach_map](https://docs.pytorch.org/tutorials/recipes/foreach_map.html) — an optimization introduced in the latest PyTorch release — to the optimizer step. Using the utility functions from the [Foreach Map tutorial](https://docs.pytorch.org/tutorials/recipes/foreach_map.html), we create an optimized Adam optimizer step function, and apply it to our optimizer: python ``` def get_compiled_adam_step(optimizer): compiled_adam = torch.compile(foreach_map_adam) inputs = get_inputs(optimizer) def compiled_adam_step(): compiled_adam(*inputs) return compiled_adam_step def apply_compilation(model, optimizer): model = torch.compile(model, fullgraph=True, mode="max-autotune") optimizer.step = get_compiled_adam_step(optimizer) return model, optimizer ``` This optimization requires use of the _zero_grads_ utility from above. It also requires that we run a warmup training step before compilation to populate all of the gradient tensors. The modified optimizer step results in a reduced throughput of 5.28 steps per second. We presume that our toy model is too small to reap the benefit of the new compilation feature. Our best result, 5.54 steps per second, is 78% faster than our baseline result. Let’s see what happens when we extend our solution to multiple GPUs. ### Model Compilation Step #5: Extending to DDP The final step is to extend the training script to use all 8 GPUs. For this step we need to disable the _fullgraph_ setting since the cross-GPU gradient sharing requires graph-breaking communication calls. The resultant throughput is 4.59 steps per second, nearly two times faster than our baseline result. ### Results The table below summarizes the results of our static-graph experiments: ![](../images/body-f2b817fa04effe14232c.webp) Static Graph Compilation Results (by Author) Thus far, all of our experiments have assumed fixed-sized input tensors. Since the vast majority of input sequences are small, our graph is performing a huge amount of wasteful computation. In the next section we will evaluate torch.compile when padding to variable-length inputs. ## Dynamic Model Compilation In this section we introduce dynamism into our toy model definition by padding the inputs sequences in each batch to the length of the longest sequence. In a previous section we described several strategies for compiling dynamic graphs. We will apply these strategies and assess their impact on the training throughput. The experiments in this section were run on a single NVIDIA A100 GPU. ### Option #1: Auto-Detect Dynamism The default behavior (_dynamic_ =_None_) of torch.compile is to auto-detect dynamism and recompile the graph accordingly. When running in this setting, we indeed see the recompilation due to the variation in input size, but we also get the following print: shell-session ``` V0806 09:31:00.624000 175763 torch/_dynamo/guards.py:2997] [0/1] [__recompiles] - 0/1: ((decoder_input_ids.size()[1]*decoder_input_ids.size()[1]) % 8) != 0 # attn_output = torch.nn.functional.scaled_dot_product_attention( # transformers/integrations/sdpa_attention.py:89 in sdpa_attention_forward (_dynamo/utils.py:3284 in run_node) ``` The source of this recompilation is the [scaled_dot_product_attention](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) operator, which requires that input shapes be aligned to multiples of eight for optimal use. To address this issue and avoid the recompilation, we modify our padding operation to pad to a multiple of eight. To avoid the recompilation that is triggered by the variable-length inputs, we define the following utility and apply it to the input tensors: python ``` def mark_dynamic(batch): for key in ['decoder_input_ids', 'labels', 'decoder_attention_mask']: torch._dynamo.mark_dynamic(batch[key], 1) def train_step(model, device, optimizer, batch): # copy data to device batch = copy_to_device(batch, device) # mark inputs as dynamic to avoid recompilation mark_dynamic(batch) optimizer.zero_grad() with torch.amp.autocast('cuda', dtype=torch.bfloat16): outputs = model(**batch) loss = outputs.loss loss.backward() optimizer.step() return loss ``` This option results in a throughput of 7.78 steps per second, 64% higher than the baseline throughput (4.73). An additional speed-up is achieved when we apply the “max-autotune” mode — 8.13 steps per second. ### Option #2: Dynamic Compilation Another way to avoid recompilations is to call torch.compile with _dynamic_ =_True_ : python ``` def apply_compilation(model, optimizer): model = torch.compile(model, fullgraph=True, dynamic=True) optimizer.step = torch.compile(optimizer.step) return model, optimizer ``` This results in a throughput of 7.77 steps per second. Since setting _dynamic_ =_True_ precludes the use of CUDA graphs, we attempt to optimize further by setting _mode_ =”max-autotune-no-cudagraphs”. This results in a throughput of 7.89 steps per second. ### Option #3: Compile a Fixed Number of Static Graphs The last option we explore is to set a fixed number of supported input shapes and compile a corresponding fixed number of static graphs. Since the default number of recompilations supported is eight, we program our collator to emit eight different tensor shapes by aligning the padding to multiples of 32. To force the recompilations, we set _dynamic_ =_False_. The resultant throughputs are for 7.77 steps per second for the default mode and 8.04 for _mode_ =”max-autotune”. Note that this option may require a greater number of warmup steps to ensure that all shape variations are processed. (An alternative is to manually feed the model with all shape variations before starting the training loop.) ### Modular Compilation Since our model naturally splits into two submodules — a static encoder and a dynamic decoder — it is tempting to explore the option of applying separate compilation to each component. Note that in an inference setting, it is essential to compile the encoder and decoder separately, since the encoder is called only once, while the decoder is called repeatedly in an auto-regressive loop. python ``` def apply_compilation(model, optimizer): model.encoder = torch.compile(model.encoder, fullgraph=True) model.decoder = torch.compile(model.decoder, fullgraph=True) model.loss_function = torch.compile(model.loss_function, fullgraph=True) optimizer.step = torch.compile(optimizer.step) return model, optimizer ``` The result of this strategy is a throughput of 7.93, which is slightly higher than the result we got (in default mode) when compiling the entire model. One advantage to this approach is the ability to tune the compilation controls to each submodule independently. For example, setting _mode_ =”max-autotune” to just the encoder, further increased the throughput to 8.04 steps per second. ### Results We summarize the results of our dynamic-graph experiments in the table below: ![](../images/body-1c165d4497c53b3d25b2.webp) Dynamic Graph Compilation Results (by Author) The best result was 8.13 steps per second, 72% higher than the baseline result (4.73). It is likely that further tuning could result in additional gains. Keep in mind that the impact of torch.compile can vary greatly based on the details of the model and the runtime environment. ## Reducing Compilation Time We now turn our attention to the duration of the torch.compile warmup. We will assess the two optimizations discussed above, [compile-time caching](https://docs.pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html) and [regional compilation](https://docs.pytorch.org/tutorials/recipes/regional_compilation.html). We limit our experiments to a single GPU. We use the default application of torch.compile and measure the duration of the first 20 training iterations. ### Pre-Loading Compilation Cache In the following demonstration of compile-time caching, we use an [Amazon S3](https://aws.amazon.com/s3/) bucket as our persistent storage location: python ``` import boto3 S3_BUCKET = "" S3_KEY = "" def download_cache(): s3_client = boto3.client('s3') t0 = time.perf_counter() try: response = s3_client.get_object(Bucket=S3_BUCKET, Key=S3_KEY) artifact_bytes = response['Body'].read() torch.compiler.load_cache_artifacts(artifact_bytes) print(f"Cache restored. Time: {time.perf_counter()-t0} sec") except: return False return True def upload_cache(): s3_client = boto3.client('s3') artifact_bytes, cache_info = torch.compiler.save_cache_artifacts() s3_client.put_object( Bucket=S3_BUCKET, Key=S3_KEY, Body=artifact_bytes ) if __name__ == '__main__': # specify inductor cache dir inductor_cache_dir = '/tmp/inductor_cache' os.environ['TORCHINDUCTOR_CACHE_DIR'] = inductor_cache_dir # clean up compiler cache torch._dynamo.reset() shutil.rmtree(inductor_cache_dir, ignore_errors=True) # upload the compilation artifacts download_cache() # train the model train() # upload the compilation artifacts upload_cache() ``` This method reduces the compilation warmup from 196 seconds to 56 seconds — a 3.5X speed-up. ### Regional Compilation To implement regional compilation, we apply compilation to the internal blocks of both the encoder and the decoder: python ``` def apply_compilation(model, optimizer): model.encoder.encoder.layer = torch.nn.ModuleList( [torch.compile(layer, fullgraph=True) for layer in model.encoder.encoder.layer] ) model.decoder.transformer.h = torch.nn.ModuleList( [torch.compile(layer, fullgraph=True) for layer in model.decoder.transformer.h] ) model.loss_function = torch.compile(model.loss_function, fullgraph=True) optimizer.step = torch.compile(optimizer.step) return model, optimizer ``` This change reduces the throughput from 7.78 steps per second to 7.61 steps per second. On the other hand, the compilation warmup drops from 196 seconds to 80 seconds — a 2.45X speed-up. In the case of our toy model — which is extremely small by today’s standards — the gains we have demonstrated are modest. But for large models, these types of compilation-time optimization techniques could prove essential. ## Summary As AI/ML models grow in size to hundreds of billions or even trillions of parameters, optimizing their runtime performance becomes increasingly essential. For PyTorch models, torch.compile is one of the most powerful optimization tools at your disposal. This post has aimed to ease the adoption of torch.compile by addressing some of its intricacies and demonstrating its practical use. Some of the main techniques we covered were: * Reducing graph-breaks and recompilations * Tuning compilation settings to maximize performance gains * Effective use of the PyTorch logs * Top-down vs. bottom-up debugging strategies * Modular application of torch.compile * Reducing the duration of compilation warmup PyTorch compilation is a complex and nuanced topic. In this post we have covered just some of its many features. For more on the topic, be refer to the [official documentation](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html). ## Register Now to Continue Reading ### Unlock Towards Data Science with a Free Account All articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more. Register for freeSign in to your account * * * ![](../images/body-97f5483d57a9b8741146.jpeg) Written By ## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/) [See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/) [Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Deep Learning](https://towardsdatascience.com/tag/deep-learning/)[Python](https://towardsdatascience.com/tag/python/)[Pytorch](https://towardsdatascience.com/tag/pytorch/) Share This Article * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fmaximizing-ai-ml-model-performance-with-pytorch-compilation%2F&title=Maximizing%20AI%2FML%20Model%20Performance%20with%20PyTorch%20Compilation) * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fmaximizing-ai-ml-model-performance-with-pytorch-compilation%2F&title=Maximizing%20AI%2FML%20Model%20Performance%20with%20PyTorch%20Compilation) * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fmaximizing-ai-ml-model-performance-with-pytorch-compilation%2F&text=Maximizing%20AI%2FML%20Model%20Performance%20with%20PyTorch%20Compilation) Towards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program. [Write for TDS](https://towardsdatascience.com/submissions/) ## Related Articles * [![](../images/002.png)](https://towardsdatascience.com/time-series-transformations-selection-using-edaspy-d97e0dc1fcca/) ## [Time series transformations selection using EDAspy](https://towardsdatascience.com/time-series-transformations-selection-using-edaspy-d97e0dc1fcca/) [](https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/) Easy way to find the optimum Time series transformations for a forecasting model in a Machine Learning project. [Vicente P. Soloviev](https://towardsdatascience.com/author/vicente-perez-soloviev/) November 25, 20204 min read * [![](../images/003.jpeg)](https://towardsdatascience.com/want-to-break-into-data-science-start-building-db5ed24a7d77/) ## [Want to break into data science? Start building](https://towardsdatascience.com/want-to-break-into-data-science-start-building-db5ed24a7d77/) [Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/) Why you should build first and ask questions later [Russell Pollari](https://towardsdatascience.com/author/russellpollari/) October 19, 20204 min read * [![](../images/004.jpg)](https://towardsdatascience.com/how-to-select-loss-function-and-activation-function-for-classification-problems-b7f32c724507/) ## [How to Select Loss Function and Activation Function for Classification Problems](https://towardsdatascience.com/how-to-select-loss-function-and-activation-function-for-classification-problems-b7f32c724507/) [Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/) Understanding how to choose activation function in the final layer and loss function for neural network [Ibrahim Kovan](https://towardsdatascience.com/author/ibrahimkovan/) April 14, 20226 min read * [![](../images/005.jpg)](https://towardsdatascience.com/three-main-components-of-machine-learning-modeling-c9186658f7fe/) ## [Three Main Components of Machine Learning Modeling](https://towardsdatascience.com/three-main-components-of-machine-learning-modeling-c9186658f7fe/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) The art of learning through data [Abiodun Olaoye](https://towardsdatascience.com/author/aolaoye/) October 24, 20224 min read * [![](../images/006.jpg)](https://towardsdatascience.com/teaching-is-hard-how-to-train-small-models-and-outperforming-large-counterparts-f131f9d463e1/) ## [Teaching is Hard: How to Train Small Models and Outperforming Large Counterparts](https://towardsdatascience.com/teaching-is-hard-how-to-train-small-models-and-outperforming-large-counterparts-f131f9d463e1/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) Distilling the knowledge of a large model is complex but a new method shows incredible performances [Salvatore Raieli](https://towardsdatascience.com/author/salvatore-raieli/) November 11, 202311 min read * [![](../images/007.jpeg)](https://towardsdatascience.com/exploring-reddit-wallstreetbets-posts-data-bed613386029/) ## [Exploring Reddit WallStreetBets Posts Data](https://towardsdatascience.com/exploring-reddit-wallstreetbets-posts-data-bed613386029/) [Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/) EDA and Sentiment Analysis of Reddit Data [Sadrach Pierre, Ph.D.](https://towardsdatascience.com/author/spierre91/) February 8, 20215 min read * [![](../images/008.png)](https://towardsdatascience.com/linear-regression-in-time-series-sources-of-spurious-regression/) ## [Linear Regression in Time Series: Sources of Spurious Regression](https://towardsdatascience.com/linear-regression-in-time-series-sources-of-spurious-regression/) [Data Science](https://towardsdatascience.com/category/data-science/) Why does the autocorrelation of the errors term matter? [JUNIOR JUMBONG](https://towardsdatascience.com/author/jumbongjunior1999/) March 10, 202513 min read * [![](../images/009.png)](https://towardsdatascience.com/how-can-machine-learning-help-visually-impaired-people-4fcdc76816b2/) ## [How machine learning can help visually impaired people](https://towardsdatascience.com/how-can-machine-learning-help-visually-impaired-people-4fcdc76816b2/) [Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/) An outdoor obstacle identification project proposal for blind people with DeepLabv3+ [Shuyang Xiang](https://towardsdatascience.com/author/vanillaxiangshuyang/) January 4, 20226 min read * [![](../images/010.png)](https://towardsdatascience.com/the-end-of-anchors-improving-object-detection-models-and-annotations-73828c7b39f6/) ## [The End of Anchors - Improving Object Detection Models and Annotations](https://towardsdatascience.com/the-end-of-anchors-improving-object-detection-models-and-annotations-73828c7b39f6/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) If you have ever had to tinker with anchor boxes, you were probably frustrated, confused and saying to yourself, "There must be another... [Anders Christiansen](https://towardsdatascience.com/author/andersasac/) April 2, 20194 min read * [YouTube](https://www.youtube.com/c/TowardsDataScience) * [X](https://x.com/TDataScience) * [LinkedIn](https://www.linkedin.com/company/towards-data-science/) * [Threads](https://www.threads.net/@towardsdatascience) * [Bluesky](https://bsky.app/profile/towardsdatascience.com) [](https://towardsdatascience.com/) Your home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals. Subscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/) © Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network) #### Recommended Articles Close * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ "Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1") * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ "The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs") * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ "Does Your Company Have a Data Strategy?") * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ "Productivity Tips, Data Career Insights, and Other Recent Must-Reads")