Struggling to get PyTorch fast enough to use in public competition

fsschneider: <p>Hello everyone,</p>
<p>We, the Algorithms Working Group of MLCommons, are developing a <a href="https://arxiv.org/abs/2306.07179" rel="noopener nofollow ugc">large-scale competitive benchmark</a> for training algorithms (e.g. Adam, SGD, Lion, VeLO, etc.) with a $50,000 prize pool. We hope to support submissions in both the PyTorch and JAX frameworks. However, we’re finding it extremely challenging to get our PyTorch code to be as fast as our JAX code.</p>
<p>Specifically, our benchmark involves running submitted training algorithms on eight different deep learning workloads on a fixed 8xV100 system. A workload is essentially “a deep learning training problem”, consisting of a dataset, a model, a loss function, and a validation/test error goal. For 3 out of our 8 workloads, we have managed to achieve comparable training wall-clock runtimes between JAX and PyTorch (namely, ImageNet ResNet, ImageNet ViT, and LibriSpeech Conformer, the runtime difference is within 5%). Another three workloads (OGBG GNN, fastMRI U-Net, and LibriSpeech DeepSpeech) have somewhat comparable training wall-clock runtimes (with a runtime difference of &lt;12%) but we still see PyTorch consistently slower. However the most problematic is that, for the remaining two workloads (Criteo DLRMsmall and WMT Transformer), our PyTorch implementations seem significantly slower than our JAX implementations (&gt;25%):</p>
<ul>
<li>Criteo DLRMsmall: We are currently struggling with an OOM issue in our PyTorch 2 implementation that is not present in our JAX version. Using PyTorch 1.13, however, PyTorch was &gt;60% slower than JAX.</li>
<li>fastMRI U-Net: 10% slower than JAX (10% slower with PyTorch 1.13)</li>
<li>LibriSpeech DeepSpeech: 12% slower than JAX (17% slower with PyTorch 1.13)
<ul>
<li>Timing was done without torch.compile due to the torchdynamo error.</li>
</ul>
</li>
<li>OGBG GNN: 10% slower than JAX (30% slower with PyTorch 1.13)
<ul>
<li>Timing was done without torch.compile due to the torchdynamo error (for using GraphsTuple as inputs).</li>
</ul>
</li>
<li>WMT Transformer: 25% slower than JAX (21% slower with PyTorch 1.13)
<ul>
<li>Timing was done without torch.compile due to the torchdynamo error (on boolean masking).</li>
</ul>
</li>
</ul>
<p>All timings use PyTorch 2 (timing results for PyTorch 1.13 are shown in parentheses) and were run on the same 8xV100 machine using the AdamW optimizer with identical hyperparameters and CUDA version 11.8. The code for the individual workloads can be found <a href="https://github.com/mlcommons/algorithmic-efficiency/tree/main/algorithmic_efficiency/workloads" rel="noopener nofollow ugc">here in our codebase</a></p>
<p>Another requirement of our workload implementations is that they produce exactly the same results between the two frameworks. Therefore, we have tried to use the same data loading and preprocessing code as much as possible. Still, there are some challenges because it is more idiomatic for our JAX implementations to run only a single Python process. In contrast, our PyTorch implementations use DistributedDataParallel training with a Python process for each of the 8 GPUs that our benchmark requires.</p>
<p>We’re planning to issue a Call for Submission for our competition around September 1st. However, with the current speed differences between PyTorch and JAX, submissions in PyTorch will not be competitive. People will most likely not submit in PyTorch because the slowdown will significantly reduce their chance at the $50,000 prize pool.</p>
<p>We could really use any help possible from the PyTorch community since we are a small team with only a handful of PyTorch users performing engineering work. We are looking for suggestions on how to speed up our PyTorch implementations and make them competitive with our JAX versions (even better if these suggestions are feasible to implement before we issue the call for submissions). We can also use more extensive engineering contributions or new additions to our working group if anyone wants to join.</p>
<p>Thank you.</p>

smth: <p>Hi Frank, thanks for the detailed outline. We will try to profile and dig into whether you are using PyTorch optimally. Might send some PRs your way.</p>

ezyang: <p>This is not a good answer for the competition itself, but the very simplest thing you should try first is using a PyTorch nightly instead of the PT 2.0 public release. There have been a <em>lot</em> of improvements that haven’t made it out to a public release yet.</p>

FindHao: <p>Hi Frank, have you tried to run PyTorch models with <code>torch.compile()</code>? Or the comparison is between JAX and pure PyTorch eager mode?</p>
<p>Oh, I saw the call for <code>torch.compile()</code> in <code>submission_runner.py</code>.</p>

fsschneider: <p>That is amazing! Thanks a lot for digging into it.<br>
Please let us know if we can help you in any way, providing you with more details or logging results.</p>

fsschneider: <p>Yes, we are using <code>torch.compile()</code> whenever we can. As mentioned above, some of the models are currently incompatible with <code>torch.compile()</code>, e.g. due to our use of boolean masking or GraphsTuple as inputs.</p>

mikekgfb: <p>We have been running into a number of issues downloading and preparing the data sets.  <a class="mention" href="/u/lessw2020">@lessw2020</a>  created a PR for fastMRI =&gt; <a href="https://github.com/mlcommons/algorithmic-efficiency/pull/473" class="inline-onebox" rel="noopener nofollow ugc">Fastmri dataset setup fixes by lessw2020 · Pull Request #473 · mlcommons/algorithmic-efficiency · GitHub</a>.  For speech and wmt, we see sentencepiece segfaulting during the data prep setup step, as per this issue =&gt; <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/470" class="inline-onebox" rel="noopener nofollow ugc">Data Downloading segfaulting at sentencepiece · Issue #470 · mlcommons/algorithmic-efficiency · GitHub</a> - consequently, we have not been able to repro the torch.compile issue you flagged for wmt.</p>
<p>Beyond the problems at hand, I’m a bit concerned about the inefficiencies that TF data pipelines may introduce to PyTorch as per your description of the hack you used to enable them for PyTorch.  Is this putting PyTorch at an unfair disadvantage?  Is there a more neutral reading ecosystem, such as Ray Data or Databricks that might be used for both frameworks rather than one that developed for a particular framework flavor?</p>

ezyang: <p>It would be helpful to know what the absolute wall times were for the benchmarks you ran, so we can cross-reference our measurements with yours. (I didn’t have a V100 ready to go, but on an A100x8 system, I was able to torch.compile and get 7703.403180360794 score in 70 evals; no idea if this is good or not (still building TF haha). EDIT: Hmm, this is in the paper, I guess I need to read it more carefully. EDIT 2: I can’t compare with the paper because I’m on A100x8, but I also ran the JAX baseline and got 7703.041719198227.</p>

ezyang: <p>I have a question: if I’m comparing the PyTorch baseline vs the JAX baseline, do I expect the number of evals to be the same? (In my case, criteo1tb, they’re not: PyTorch does 70 evals, JAX does 92, but the score ends up being really close in the end anyway?)</p>

runame: <p>Not sure if it is expected for criteo1tb, but it is definitely possible. The score only takes the timing of the training into account and the eval frequency is based on time and not steps, so if the evals are slower in one framework than the other, it is possible that the number of evals differs while the final score is similar.</p>

fsschneider: <p>We could consider using a nightly PyTorch version for the competition, as long as we are able to pin to a specific version so that all results are reproducible.<br>
Thanks for the suggestion, we will investigate how the timings change with the nightly version.</p>

marksaroufim: <p>Hi everyone, just wanted to provide a macro update on the state of things so far and (not so) briefly give a mental model for <code>torch.compile()</code> that’ll help in squeezing out as much performance as possible</p>
<h1><a name="how-to-fix-torchcompile-bugs-1" class="anchor" href="#how-to-fix-torchcompile-bugs-1"></a>How to fix torch.compile bugs</h1>
<p>So you tried to <code>torch.compile()</code> something but it broke what now?</p>
<ol>
<li>Try the nightlies instead, the project is moving fast. I was using <code>torch                        2.1.0.dev20230814+cu118</code></li>
<li>Look at the error message and open up an issue on <a href="https://github.com/pytorch/pytorch" class="inline-onebox">GitHub - pytorch/pytorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration</a></li>
<li>While you’re waiting if you look at the offending line, can you rewrite it to sidestep the issue?</li>
<li>If your error only shows up with DDP then follow instructions here <a href="https://github.com/pytorch/pytorch/issues/107362#issuecomment-1682658771" class="inline-onebox">[Inductor] Run compiled model failed on 2023_08_17 nightly · Issue #107362 · pytorch/pytorch · GitHub</a></li>
<li>Try out <code>backend=aot_eager</code> it’ll help isolate inductor specific bugs but to get perf you essentially get inductor working which brings us to</li>
</ol>
<p>Ok so now your model compiles but it’s slow how come?</p>
<ol>
<li>Are you using the default <code>backend="inductor"</code></li>
<li>Does your model have graph breaks. i.e does it run with <code>fullgraph=True</code> - <code>torch.compile()</code> is an API that progressively unloads complexity onto users but as far as a competition goes getting to 0 graph breaks will be necessary</li>
<li>Do compilation times matter? If not then try <code>mode="max-autotune"</code></li>
</ol>
<p>I wish I had more performance numbers to actually share but unless we at the very least get inductor working those numbers will not be particularly meaningful.</p>
<p>Sometimes graph breaks will make sense if for example you notice print statement and it’s your responsibility to remove them to get good perf and other times there’s real bugs. Two noticeable ones I’d like to highlight are</p>
<ol>
<li>Optimizer support which I’ve tagged Jane on <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/487#issuecomment-1679732542" class="inline-onebox">torch.compile default settings · Issue #487 · mlcommons/algorithmic-efficiency · GitHub</a> which popped up in the WMT model</li>
<li>Incorrect tracing which popped up in the OGBG model <a href="https://github.com/pytorch/pytorch/issues/107040" class="inline-onebox">Dynamo not handling a NamedTuple · Issue #107040 · pytorch/pytorch · GitHub</a></li>
</ol>
<p>In both of those cases the fastest way to solve the problem is to open up a github issue while you try out simple 1 liners. I will also attend your meetings regularly to help unblock or triage open issues if there’s a lot of them.</p>
<h1><a name="minimal-repros-2" class="anchor" href="#minimal-repros-2"></a>Minimal repros</h1>
<p>So when you open up an issue, ideally we want a minimal reproduction and one challenge many of us faced was dealing with getting the preprocessing pipelines to work. I hope you’ll set some simple smoke tests <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/486" class="inline-onebox">Add dataset setup tests · Issue #486 · mlcommons/algorithmic-efficiency · GitHub</a> since that was the main reason I couldn’t get even more people involved in this effort.</p>
<ol>
<li>The segfault in dataset_setup.py affected many models but was fixed <a href="https://github.com/mlcommons/algorithmic-efficiency/pull/477" class="inline-onebox">Fix segfault in dataset_setup by msaroufim · Pull Request #477 · mlcommons/algorithmic-efficiency · GitHub</a></li>
<li>Less fixed the fastmri dataset pipeline <a href="https://github.com/mlcommons/algorithmic-efficiency/pull/473" class="inline-onebox">Fastmri dataset setup fixes by lessw2020 · Pull Request #473 · mlcommons/algorithmic-efficiency · GitHub</a></li>
<li>Ed found that the criteo pipeline also failed <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/479" class="inline-onebox">criteo preproc doesn't match input_pipeline · Issue #479 · mlcommons/algorithmic-efficiency · GitHub</a> and could benefit from having more options to run things locally for quick spot checks <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/475" class="inline-onebox">RuntimeError: shape '[213568]' is invalid for input of size 262144 on criteo1tb · Issue #475 · mlcommons/algorithmic-efficiency · GitHub</a></li>
<li>I could never get the librispeech pipeline to work and even then it would only fail after 1-2h, I’d like more ways to debug things fast locally and felt like I couldn’t get clever, I had to run the benchmark as is or it would probably fail <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/475" class="inline-onebox">RuntimeError: shape '[213568]' is invalid for input of size 262144 on criteo1tb · Issue #475 · mlcommons/algorithmic-efficiency · GitHub</a> and <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/471" class="inline-onebox">OGBG OOM'in in eager mode pytorch · Issue #471 · mlcommons/algorithmic-efficiency · GitHub</a></li>
<li>Getting an E2E perf number seems like it takes about 24h (at least for criteo), I wish we could have some way of spot checking expected performance as we iterate on the models</li>
</ol>
<p>One thing we did learn from you is that we should have more end to end data loading examples in our CI as well to catch more errors, In the coming weeks I’ll be looking at what we can upstream from you both in terms of specific models so they don’t regress but also your methodology which others internally have found appealing.</p>
<p>Ok that was not so brief but if I were to spend more time on this here’s how I would track the work per model, make sure that <code>torch.compile()</code> works and then make sure that I’ve removed graph breaks and added <code>mode="max-autotune"</code> this is the strategy we’re taking for models that people care about and it’s a bit time consuming because PyTorch will allow you to write non performant code instead of erroring out.</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>Task</th>
<th>Owner</th>
<th>Preprocessing pipeline</th>
<th>Made torch.compile work</th>
<th>Removed all graph breaks + cudagraph</th>
<th>Perf relative to JAX</th>
</tr>
</thead>
<tbody>
<tr>
<td>WMT</td>
<td>Mark/Michael</td>
<td>Fixed</td>
<td>Yes</td>
<td></td>
<td></td>
</tr>
<tr>
<td>Librispeech</td>
<td>Mark</td>
<td>Not fixed</td>
<td>No</td>
<td></td>
<td></td>
</tr>
<tr>
<td>OGDB</td>
<td>Mark</td>
<td>Fixed</td>
<td>NamedTuple dynamo bug</td>
<td></td>
<td></td>
</tr>
<tr>
<td>Criteo</td>
<td>Colin/Ed</td>
<td>Fixed</td>
<td>Yes</td>
<td></td>
<td>Matches</td>
</tr>
<tr>
<td>Fast-MRI</td>
<td>Less</td>
<td>Fixed</td>
<td>Yes</td>
<td></td>
<td></td>
</tr>
</tbody>
</table>
</div><h1><a name="per-model-updates-3" class="anchor" href="#per-model-updates-3"></a>Per Model updates</h1>
<p>None of the models were tested with cuda graphs or fullgraph yet so this will make us look worst vs JAX where things have to be fullgraph</p>
<ul>
<li>
<p>OGDB (Mark): NamedTuple not being traced correctly <a href="https://github.com/pytorch/pytorch/issues/107040" class="inline-onebox">Dynamo not handling a NamedTuple · Issue #107040 · pytorch/pytorch · GitHub</a></p>
</li>
<li>
<p>Criteo DLRM Small (Ed/Colin): Issues with data loading <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/475" class="inline-onebox">RuntimeError: shape '[213568]' is invalid for input of size 262144 on criteo1tb · Issue #475 · mlcommons/algorithmic-efficiency · GitHub</a></p>
<ul>
<li>Notes from Colin
<ul>
<li>How much of the latency is preproc/data loading vs forward backward?</li>
<li>It looks like their DLRM_small is a single embedding table. Why are they doing this? Even old dlrm (from 2019) used multiple and scaled down
<ul>
<li>It will be much simpler for them to use <a href="https://github.com/mlcommons/training/tree/master/recommendation_v2/torchrec_dlrm">https://github.com/mlcommons/training/tree/master/recommendation_v2/torchrec_dlrm</a> as a base (most current dlrm). They can simply scale down the DCN layers/embedding sizes as they see fit (although the full model should run on 8 V100s I believe)</li>
</ul>
</li>
<li>They are using vanilla nn.Embedding
<ul>
<li>I’m not sure why this isn’t a pooled embedding? Most dlrm tables are pooled</li>
<li>They will see <a href="https://github.com/pytorch/torchrec/tree/main/benchmarks">speedups</a> with switching to TorchRec (TBE backed)</li>
</ul>
</li>
<li>It looks like they are using DDP over embedding. They will likely see speedups by using a different sharding mechanism</li>
<li>They can also see speedup through optimizer fusion + pipelining, although this may be out of scope for them?</li>
<li>Tangential thought → how that torchRec/FBGEMM is more established, should we consider changing nn.EmbeddingBag / nn.Embedding implementation to use TBE? cc <a href="mailto:ezyang@meta.com">Edward Yang</a></li>
</ul>
</li>
</ul>
</li>
<li>
<p>Fast-mri (Less):</p>
<ul>
<li>broken data loading pipeline fixed <a href="https://github.com/mlcommons/algorithmic-efficiency/pull/473" class="inline-onebox">Fastmri dataset setup fixes by lessw2020 · Pull Request #473 · mlcommons/algorithmic-efficiency · GitHub</a></li>
<li>Model was already torch.compilable</li>
</ul>
</li>
<li>
<p>Librispeech (Mark):</p>
<ul>
<li>Segfault on dataloading <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/470" class="inline-onebox">Data Downloading segfaulting at sentencepiece · Issue #470 · mlcommons/algorithmic-efficiency · GitHub</a> - fixed</li>
<li>Data preprocessing script has some bugs and is crashing after 1h</li>
<li>Mlperf team got this to work with eager <a href="https://github.com/mlcommons/algorithmic-efficiency/pull/483" class="inline-onebox">librispeech_deepspeech eager compile by pomonam · Pull Request #483 · mlcommons/algorithmic-efficiency · GitHub</a> - they ran into a complex64 bug with inductor which should be fixed by falling back to eager in nightlies</li>
</ul>
</li>
<li>
<p>WMT transformer (Michael/Mark):</p>
<ul>
<li>Segfault on dataloading <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/470" class="inline-onebox">Data Downloading segfaulting at sentencepiece · Issue #470 · mlcommons/algorithmic-efficiency · GitHub</a> some interaction between boolean masks and torch.compile - fixed</li>
<li>Cannot repro boolean mask problem w/ hf model but only their custom implementation <a href="https://gist.github.com/msaroufim/946d7d26e89bab0bfe83d9929b533701" class="inline-onebox">gist:946d7d26e89bab0bfe83d9929b533701 · GitHub</a></li>
<li>Nightlies solved boolean mask problem</li>
<li>Now getting an AOT autograd issue when re-assigning shared embeddings P806943829</li>
<li>This seems fixed in the nightlies <a href="https://github.com/mlcommons/algorithmic-efficiency/issues/467#issuecomment-1675608043" class="inline-onebox">WMT slower in Pytorch than Jax · Issue #467 · mlcommons/algorithmic-efficiency · GitHub</a> - will check if mlperf team can repro too and indeed they can <a href="https://github.com/mlcommons/algorithmic-efficiency/pull/484" class="inline-onebox">PyTorch Workload Fixes by pomonam · Pull Request #484 · mlcommons/algorithmic-efficiency · GitHub</a></li>
</ul>
</li>
</ul>

Priya_Kasimbeg: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> thank you for the comprehensive update!</p>
<p>On a high level if I understand correctly:</p>
<ul>
<li>OGBG: not fully torch compiled yet so we may see further speed improvements.</li>
<li>Criteo: further speed ups possible by refactoring the embedding and not using DDP over embedding.</li>
<li>FastMRI: already torch compilable but not checked with fullgraph so we may or may not see further speed improvements.</li>
<li>Librispeech deepspeech: works with eager and nightly.</li>
<li>WMT transformer: full compilation blocked on AOT autograd issue but fixed nightlies.</li>
</ul>
<p>The workloads at highest risk from our perspective are Criteo which OOMs and WMT which is 22% slower than jax with current checked in fixes. Regarding further action items:</p>
<ul>
<li>Criteo1tb: It sounds like your team was blocked on running the code due to the data setup issues, so didn’t get to run the workload, correct? I’ll work with <a class="mention" href="/u/janeyx99">@janeyx99</a> to set them up with a VM on our GCP project to unblock debugging of the OOM issue. Once that is fixed we can work on the embedding refactoring.</li>
<li>WMT: Possibly fully torch compilable with nightly, is that correct? In that case I will re-benchmark the timing to get the latest update.</li>
<li>OGBG: full torch compile blocked on NamedTuple tracing issue.</li>
<li>other workloads : (AI for us) Rerun with torch.nightly and fullgraph to get up to date timing comparisons.</li>
</ul>
<p>Let me know if that sounds right or if you have any other suggestions.</p>

marksaroufim: <ul>
<li>Criteo: Yes if you can unblock Jane with a machine that will go a long way</li>
<li>WMT: Yes in fact I did try compiling a similar HF model and that worked just fine but most importantly here try to to make sure your implementation compiles, compiles with inductor and then finally has no graph breaks</li>
<li>OGBG: Yes, if you like comment here <a href="https://github.com/pytorch/pytorch/issues/107040" class="inline-onebox">Dynamo not handling a NamedTuple · Issue #107040 · pytorch/pytorch · GitHub</a> to +1 it. If you can also change the code to not use a NamedTuple here instead it might unblock you</li>
<li>Other workloads: fullgraph and change the mode to either reduce-overhead or max-autotune and most importantly make sure inductor is being used</li>
</ul>

Priya_Kasimbeg: <aside class="quote no-group" data-username="marksaroufim" data-post="12" data-topic="186015">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/marksaroufim/48/36747_2.png" class="avatar"> marksaroufim:</div>
<blockquote>
<p>torch 2.1.0.dev20230814+cu118</p>
</blockquote>
</aside>
<p>Is there an index url we can use to pip install the dev package versions?</p>
<p>Having some trouble with setuptools installing nightly.</p>

marksaroufim: <p>You can do something like this <code>pip3 install torch==2.1.0.dev20230814+cu118 --index-url https://download.pytorch.org/whl/nightly/cu118</code></p>
<p>This is the index url <a href="https://download.pytorch.org/whl/nightly/torch/">https://download.pytorch.org/whl/nightly/torch/</a></p>
<p>And the selector on the main <a href="http://pytorch.org">pytorch.org</a> webpage can show you the right way to install nightlies for different architectures or cuda versions here <a href="https://pytorch.org/">https://pytorch.org/</a></p>

marksaroufim: <p>For people following <a class="mention" href="/u/priya_kasimbeg">@Priya_Kasimbeg</a> invited me to their regular engineering syncs - will attend those and make sure we close on any loose ends</p>

Priya_Kasimbeg: <p>Hi Mark,<br>
I upgraded to torch nightly version torch.dev20230820 and reran all of the workloads. There are a couple of new issues with torch compile in the nightly versions on some of the other workloads.</p>
<ul>
<li><a href="https://github.com/mlcommons/algorithmic-efficiency/issues/496" rel="noopener nofollow ugc">ViT Torch compile: BackendCompilerFailed Cannot call sizes() on tensor with symbolic sizes strides on torch nightly</a></li>
<li><a href="https://github.com/mlcommons/algorithmic-efficiency/issues/498" rel="noopener nofollow ugc">Pytorch Deepspeech Torch compile TorchRuntimeError on torch nightly</a></li>
</ul>
<p>We’d like to pin the pytorch version in approximately 2 weeks and I am not sure which pytorch version to use, because I am not sure how long it will take to resolve the new issues on nightly.</p>
<h3><a name="current-state-on-stable-1" class="anchor" href="#current-state-on-stable-1"></a>Current state on stable:</h3>
<div class="md-table">
<table>
<thead>
<tr>
<th>Workload</th>
<th>Preprocessing pipeline</th>
<th>Made torch.compile work</th>
<th>Backend</th>
<th>Perf Pytorch relative to Jax (t_pytorch - t_jax)/t_jax</th>
</tr>
</thead>
<tbody>
<tr>
<td>criteo1tb</td>
<td>Not Fixed</td>
<td>No: OOM</td>
<td>aot_eager</td>
<td></td>
</tr>
<tr>
<td>fastmri</td>
<td>Fixed</td>
<td>Yes</td>
<td>inductor</td>
<td>9</td>
</tr>
<tr>
<td>imagenet_resnet</td>
<td></td>
<td>Yes</td>
<td>inductor</td>
<td>4</td>
</tr>
<tr>
<td>imagenet_vit</td>
<td></td>
<td>Yes</td>
<td>inductor</td>
<td>-10</td>
</tr>
<tr>
<td>librispeech_conformer</td>
<td>Fixed</td>
<td>Yes</td>
<td>eager</td>
<td>-1</td>
</tr>
<tr>
<td>librispeech_deepspeech</td>
<td>Fixed</td>
<td>Yes</td>
<td>eager</td>
<td>13</td>
</tr>
<tr>
<td>ogbg</td>
<td>Fixed</td>
<td>No: NamedTuple dynamo bug</td>
<td>inductor</td>
<td>11</td>
</tr>
<tr>
<td>wmt</td>
<td>Fixed</td>
<td>Yes</td>
<td>inductor</td>
<td>22</td>
</tr>
</tbody>
</table>
</div><h3><a name="current-state-on-nightly-torchdev08202023-2" class="anchor" href="#current-state-on-nightly-torchdev08202023-2"></a>Current state on nightly torch.dev08202023:</h3>
<div class="md-table">
<table>
<thead>
<tr>
<th>Workload</th>
<th>Preprocessing pipeline</th>
<th>Made torch.compile work</th>
<th>Backend</th>
<th>Perf Pytorch relative to Jax (t_pytorch - t_jax)/t_jax</th>
</tr>
</thead>
<tbody>
<tr>
<td>criteo1tb</td>
<td>Not Fixed</td>
<td>No: OOM</td>
<td>aot_eager</td>
<td>blocked</td>
</tr>
<tr>
<td>fastmri</td>
<td>Fixed</td>
<td>Yes</td>
<td>inductor</td>
<td>todo</td>
</tr>
<tr>
<td>imagenet_resnet</td>
<td></td>
<td>Yes</td>
<td>inductor</td>
<td>todo</td>
</tr>
<tr>
<td>imagenet_vit</td>
<td></td>
<td>No</td>
<td>inductor</td>
<td>blocked</td>
</tr>
<tr>
<td>librispeech_conformer</td>
<td>Fixed</td>
<td>No: OOM</td>
<td>eager</td>
<td>blocked</td>
</tr>
<tr>
<td>librispeech_deepspeech</td>
<td>Fixed</td>
<td>No</td>
<td>eager</td>
<td>blocked</td>
</tr>
<tr>
<td>ogbg</td>
<td>Fixed</td>
<td>No: NamedTuple dynamo bug</td>
<td></td>
<td>blocked</td>
</tr>
<tr>
<td>wmt</td>
<td>Fixed</td>
<td>Yes</td>
<td>inductor</td>
<td>todo</td>
</tr>
</tbody>
</table>
</div><p>Our highest priorities in terms of the workloads at the moment are to fix the Criteo pytorch workload. On our end we’re working on replacing the Embedding layer but we definitely could use your help still there because it’s strange that it was working on Pytorch 1.0 and we are not sure if changing the Embedding is going to fix it. We’re also working on fixing the data pipeline to unblock you all.<br>
Second highest priority workload is the WMT workload which is 22% slower still on pytorch. It compiles with inductor but we have to use torch stable 2.0.1. to check whether it compiles with fullgraph (pytorch/issues/107362).</p>