Pith. sign in

REVIEW 2 major objections 4 minor 31 references

Optimus rewrites deep-learning model graphs at the operator level with a small set of atomic patterns and greedy search, delivering up to 63% higher throughput plus lower memory and compile time on production recommendation models while kee

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

T0 review · grok-4.5

2026-07-12 05:53 UTC pith:OL2F6CB6

load-bearing objection Solid production systems paper: open PT2 plugin with real QPS/memory/compile gains on rec models; novelty is engineering, not theory. the 2 major comments →

arxiv 2607.02945 v1 pith:OL2F6CB6 submitted 2026-07-03 cs.PF

Optimus: A Generic Operator-Level PyTorch Model Transformation Framework

classification cs.PF
keywords PyTorch model transformationoperator-level fusiongraph rewriterecommendation systemskernel optimizationgreedy pattern matchingcompile-time reduction
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Industrial recommendation models are large, diverse and constantly changing, so hand-written module-level rewrite rules cannot keep pace. Optimus lives inside the PyTorch 2 compiler and instead matches a concise library of atomic operator patterns—batching independent linears or layernorms, eliminating redundant split-and-concat steps, padding GEMMs for better kernels, quantizing activations—then replaces each match with a faster equivalent. A greedy search that starts at anchor nodes and walks only local neighborhoods finds the matches quickly enough for real training runs. On five industry-scale models the approach yields as much as 63 percent higher queries per second, 6 percent lower peak memory, and more than 400 seconds shorter compile times. Because the patterns are small and the interface is open, engineers can keep adding new rules as architectures evolve without rewriting whole modules.

Core claim

Complex hand-engineered module-level graph rewrites can be replaced by a compact library of atomic operator-level patterns together with an efficient greedy neighborhood search; when these patterns are applied inside the PyTorch 2 stack they produce large, semantics-preserving gains in throughput, peak memory and compile time on real recommendation models.

What carries the argument

The greedy subgraph search that begins at designated anchor nodes, walks reverse-topological order for horizontal batch/group fusion and recursive child matching for vertical patterns, then rewrites the matched atomic motif while preserving inputs and outputs.

Load-bearing premise

The small, hand-written set of atomic patterns plus fixed search depths and fuse-set sizes will keep matching the real bottlenecks of future model architectures without missing profitable rewrites or breaking numerical results.

What would settle it

Run the same end-to-end measurements of QPS, peak memory and compile time on a new production recommendation model whose dominant kernels do not match any current Optimus patterns; if kernel counts stay high and no meaningful gains appear, the claim that the concise pattern set is sufficient fails.

Watch this falsifier — get emailed when new claim-graph text bears on it.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 4 minor

Summary. The paper presents Optimus, an operator-level graph transformation framework integrated into the PyTorch 2.x compiler stack. It decomposes optimizations into a library of hand-written atomic patterns (horizontal batch/group fusion and vertical rewrites such as split-cat elimination, pad_aten_mm, and activation quantization) and applies them via a greedy, anchor-driven search (Algorithms 1–2) over Torch-IR and ATen-IR. Semantic preservation is enforced by numerical checks and bisection controls. Evaluated on five industry-scale DLRM-style recommendation models and a public TorchBench illustrative model, Optimus reports up to 63% QPS gain, 6% peak-memory reduction, and >400 s compile-time reduction relative to TorchInductor, with open-source release inside torch/_inductor/fx_passes.

Significance. If the reported gains hold under production constraints, Optimus supplies a practical, extensible middle layer between module-level FX rewrites and full equality-saturation superoptimizers. The concrete strengths are (i) open-source integration into PT2 with a modular pattern-registration interface, (ii) GPU-trace evidence of kernel-count collapse (e.g., 12 772 → 24 training kernels on the illustrative model), (iii) explicit numerical-check / bisection machinery for semantic safety, and (iv) measured trade-offs between compile budget and QPS on real hardware (H100/B200). These make the work immediately usable by industrial ML engineers and a useful reference for compiler researchers targeting recommendation workloads.

major comments (2)
  1. [§5.1, Table 1] §5.1 and Table 1: All quantitative claims (63 % QPS, 6 % memory, 400 s compile-time) rest on five proprietary production models whose architectures are only loosely characterized as “DHEN-like / Wukong-like / …”. Without public weights, graphs, or a more detailed public proxy suite beyond the single TorchBench illustrative model, independent verification of the headline numbers is impossible. At minimum the authors should release the exact FX/ATen graphs (or a synthetic generator matching the reported operator histograms) used for the Table 1 measurements.
  2. [§4, Algorithms 1–2, Table 2] §4 / Algorithms 1–2 and Table 2: The free parameters min_fuse_set_size, max_depth and the activation-quantization thresholds are acknowledged but only partially ablated (compile-budget rows for Models A and E). Because every reported gain depends on these defaults remaining near-optimal for future model variants, a systematic sensitivity study (or an automated parameter-selection heuristic) is load-bearing for the claim of “general-purpose” applicability. The current evidence leaves open the possibility that the 63 % figure is tightly tuned to the five evaluated architectures.
minor comments (4)
  1. [Abstract] Abstract and §1: “over400second” and similar missing spaces appear repeatedly; a global proof-read is needed.
  2. [Fig. 9] Fig. 9 caption and surrounding text: absolute kernel counts are given, yet the y-axis units of the absolute-time plot are never stated; adding them would improve readability.
  3. [§6] §6 Related Work: TASO and TENSAT are correctly cited, but the discussion of why their search times are prohibitive for production recsys models would benefit from a single quantitative comparison (even on the public TorchBench model) rather than the qualitative statement that they “yield zero performance improvement”.
  4. [§3.2, Appendix 8.1] Listing 1 / Fig. 10a: the optimized split-concat example is trivial (identity); a slightly richer public example that still demonstrates the multi-level split elimination would better illustrate the vertical-pass power.

Circularity Check

0 steps flagged

No circularity: Optimus reports measured end-to-end gains from hand-written atomic rewrite rules and a greedy search; nothing is defined in terms of the reported speedups.

full rationale

The paper is a systems/compiler contribution. Its central claims (up to 63% QPS, 6% peak-memory reduction, >400 s compile-time decrease) are empirical measurements of TorchInductor graphs after applying a fixed library of manually designed, semantics-preserving operator-level rewrite rules (Table 1) via Algorithms 1–2. The patterns (split-cat elimination, batch linear/layernorm, pad_aten_mm, activation quantization, etc.) are obtained by GPU-trace analysis and model inspection, not by fitting free parameters to the reported deltas; the search depths and fuse-set sizes are user-tunable defaults, not quantities derived from the target metrics. Semantic preservation is enforced by construction (identical inputs/outputs, numerical checks, bisection). Self-citations are to prior Meta/PyTorch engineering blogs and papers that supply context or related tooling; none is invoked as a uniqueness theorem that forces the present results. The open-sourced code and TorchBench illustrative model further allow independent verification. Consequently no step reduces a claimed prediction to its own inputs by definition or by self-citation chain.

Axiom & Free-Parameter Ledger

3 free parameters · 4 axioms · 1 invented entities

The work is an engineering systems paper. It rests on standard compiler and graph-theory facts, the domain assumption that production recommendation graphs contain many small independent or split-cat motifs, and a handful of hand-chosen search hyperparameters. No new physical entities or free parameters fitted to scientific data are introduced; the free parameters are purely algorithmic knobs.

free parameters (3)
  • min_fuse_set_size = 5
    Default threshold (5) below which candidate sets are skipped for horizontal fusion; chosen by engineering judgment, not derived.
  • max_depth / min_depth (BFS) = max_depth=5 (default)
    Search-depth bounds that trade compile time against fusion opportunity; defaults given, user-tunable.
  • activation-quantization thresholds (size_in_mb, allowed_dtypes, use_scaling, exclude_primals) = defaults used in evaluation
    Configuration knobs that control which activations are quantized; affect memory vs numerical-equivalence trade-off.
axioms (4)
  • standard math Graph matching on ML DAGs is effectively tree inclusion and admits efficient recursive matching.
    Invoked in Section 4 to justify linear-space / sub-quadratic complexity of VerticalOpt.
  • domain assumption Semantic equivalence of each rewrite rule can be maintained by preserving inputs/outputs and by runtime numerical checks.
    Stated in Sections 2.2 and 3; underpins the claim that optimizations do not change model behavior.
  • domain assumption Production recommendation models contain many small, independent or split-cat operator motifs that dominate kernel-launch and compile cost.
    Motivates the entire pattern library and is the reason the greedy local search is claimed to be sufficient.
  • ad hoc to paper Hand-written atomic patterns plus sequential conflict-free application are sufficient; no global optimality search is required.
    Core design choice of Optimus (Sections 3–4) that distinguishes it from TASO/TENSAT-style exhaustive or equality-saturation approaches.
invented entities (1)
  • Optimus pattern library and greedy HorizontalOpt/VerticalOpt search independent evidence
    purpose: Provide a modular, extensible rewrite engine inside PT2 that avoids expensive subgraph substitution search.
    The concrete system and algorithm are the paper’s main artifacts; independent evidence is the open-source code and public TorchBench example.

pith-pipeline@v1.1.0-grok45 · 21387 in / 2912 out tokens · 21684 ms · 2026-07-12T05:53:53.049271+00:00 · methodology

0 comments
read the original abstract

In large-scale industrial applications, deep learning models that power recommendation and ranking have complex and diverse model architectures. These models are continuously developed and refined by large teams of machine learning engineers, rendering manual optimization infeasible. Consequently, graph-based optimization techniques have become an industry standard for boosting performance, with PyTorch FX transformations leading the charge. These transformations typically rely on a set of human-engineered module-level rewrite rules which are not scalable to diverse model architectures. To address this limitation, we introduce Optimus, a general-purpose model transformation framework built in the PyTorch 2.x (PT2) machine learning compiler. With a concise set of predefined patterns, Optimus applies an efficient greedy search algorithm for pattern matching and replacement, while preserving model semantic. It is designed and implemented as a highly customizable and extensible framework integrated into the PT2 stack. Our evaluation shows that the framework can achieve up to 63% speedup, 6% peak memory reduction, and over 400 second compile time decrease for our industry-scale recommendation models compared to baselines. Optimus is open-sourced together with PyTorch 2.x as a customizable model transformation layer.

Figures

Figures reproduced from arXiv: 2607.02945 by Brian Hirsh, Chi-Keung Luk, Daohang Shi, Edward Yang, Elias Ellison, Huaqing Xiong, Jason Ansel, Jia Chen Ren, Jia Liu, Jiaqi Xu, Junqing Zhou, Menglu Yu, Mingming Ding, Oguz Ulgen, Quanyu Zhu, Ruilin Chen, Shuai Yang, Will Feng, Xu Zhao, Yanbo Liang, Yuhang Yang, Yuzhen Huang.

Figure 2
Figure 2. Figure 2: Optimus architecture overview. 2.2 Semantic-Preserving Program Rewrite The core challenge in semantic-preserving program rewrite lies in defining effective pattern replacement rules, particularly when dealing with complex programs, while ensuring that these rules remain applicable across a wide range of program structures. An additional challenge is the design of an efficient search algorithm, without whic… view at source ↗
Figure 3
Figure 3. Figure 3: High-level diagram of model transformation. [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Horizontal transformation: batch vs group fusion. [PITH_FULL_IMAGE:figures/full_fig_p004_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Chained model transformations Listing 1: Optimized split-concat example. def fn(x): return x input = torch.randn(2, 6) output = torch.compile(fn)(input) We present an example in [PITH_FULL_IMAGE:figures/full_fig_p005_5.png] view at source ↗
Figure 8
Figure 8. Figure 8: Compile time overhead in sec￾ond on foundation model [PITH_FULL_IMAGE:figures/full_fig_p008_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: TorchBench evaluation results: (a) Execution profile highlighting absolute times, variance, and kernel counts; (b) [PITH_FULL_IMAGE:figures/full_fig_p009_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: (a) Graph transformation representation of a split [PITH_FULL_IMAGE:figures/full_fig_p010_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: Long computation kernel cutlass_75_xxx [PITH_FULL_IMAGE:figures/full_fig_p011_11.png] view at source ↗
Figure 14
Figure 14. Figure 14: GPU trace screenshot without split cat patterns. [PITH_FULL_IMAGE:figures/full_fig_p011_14.png] view at source ↗
Figure 16
Figure 16. Figure 16: Overview of activation quantization pattern. [PITH_FULL_IMAGE:figures/full_fig_p011_16.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

31 extracted references · 4 linked inside Pith

  1. [1]

    Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Re Re Gong, Michael Gschwind, and Soumith Chintala. 2024. PyTorch 2: Faster Machine Learning Through...

  2. [2]

    Newsha Ardalani, Sanyam Zheng, Ziyi Liu, Hao Hu, Jiarui Yu, Bowen Chen, Jiyan Yang, and Liqun Zhao. 2024. Wukong: Towards a Scaling Law for Large-Scale Recommendation. InProceedings of the 41st International Conference on Machine Learning (ICML). ACM / Curran Associates, 13570–13591. doi:10.5555/3692070. 3694525

  3. [3]

    Sorav Bansal and Alex Aiken. 2006. Automatic Generation of Peephole Superop- timizers.ACM SIGARCH Computer Architecture News34, 5 (2006), 394–403

  4. [4]

    Philip Bille and Inge Li Gortz. 2011. The tree inclusion problem: In linear space and faster.ACM Transactions on Algorithms (TALG)7, 3 (2011), 1–47

  5. [5]

    Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. InProceedings of the 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI). 578–594

  6. [6]

    Chillee. 2021. Where do the 2000+ PyTorch operators come from?: More than you wanted to know. https://dev-discuss.pytorch.org/t/where-do-the-2000-pytorch- operators-come-from-more-than-you-wanted-to-know/373

  7. [7]

    Will Constable, Xu Zhao, Victor Bittorf, Eric Christoffersen, Taylor – Robie, Eric Han, Peng Wu, Nick Korovaiko, Jason Ansel, Orion Reblitz-Richardson, and Soumith Chintala. 2020. TorchBench: A Collection of Open Source Benchmarks for PyTorch Performance and Usability Evaluation. https://github.com/pytorch/ benchmark. Accessed: 2020-09

  8. [8]

    Patrik Czakó, Gábor Kertész, and Sándor Szénási. 2025. Addressing Activation Outliers in LLMs: A Systematic Review of Post-Training Quantization Techniques. IEEE Access(2025)

  9. [9]

    Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’Aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. 2012. Large Scale Distributed Deep Networks. InAdvances in Neural Information Processing Systems. 1223–1231

  10. [10]

    Shi, Shuai Yang, Menglu Yu, Chunzhi Yang, and Jade Nie

    Mingming Ding, James Wu, Sam Larsen Ulgen, Bob Ren, Laith Sakka, Pian Pawakapan, Animesh Jain, Edward Yang, Yuzhen Huang, Ruilin Chen, Dao- hang ... Shi, Shuai Yang, Menglu Yu, Chunzhi Yang, and Jade Nie. 2025. Experience in Reducing PT2 Compilation Time for Meta Internal Work- loads. https://pytorch.org/blog/experience-in-reducing-pt2-compilation-time- f...

  11. [11]

    Jingzhi Fang, Yanyan Shen, Yue Wang, and Lei Chen. 2020. Optimizing DNN computation graph using graph substitutions.Proceedings of the VLDB Endowment 13, 12 (2020), 2734–2746

  12. [12]

    Garey, David S

    Michael R. Garey, David S. Johnson, and Larry Stockmeyer. 1976. The complexity of near-optimal graph coloring.SIAM J. Comput.5, 4 (1976), 743–765

  13. [13]

    Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia, and Alex Aiken. 2019. TASO: Optimizing Deep Learning Computation with Automatic Generation of Graph Substitutions. InProceedings of the 27th ACM Symposium on Operating Systems Principles (SOSP). 47–62

  14. [14]

    Rajeev Joshi, Greg Nelson, and Keith H. Randall. 2002. Denali: A Goal-Directed Superoptimizer. InProceedings of the ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI). ACM, Berlin, Germany, 304–314

  15. [15]

    Henry Massalin. 1987. Superoptimizer: A Look at the Smallest Program.ACM SIGARCH Computer Architecture News15, 5 (1987), 122–126

  16. [16]

    Mudigere, Jianyu Huang, Narayanan Sundara- man, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G

    Maxim Naumov, Dheevatsa ... Mudigere, Jianyu Huang, Narayanan Sundara- man, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G. Azzolini, et al. 2019. Deep learning recommendation model for personalization and recommendation systems.arXiv preprint arXiv:1906.00091(2019)

  17. [17]

    Luk, Xiaodong Wang, and Jackie (Jiaqi) Xu

    Jade Nie, C.K. Luk, Xiaodong Wang, and Jackie (Jiaqi) Xu. 2022. Opti- mizing Production PyTorch Models’ Performance with Graph Transforma- tions. https://pytorch.org/blog/optimizing-production-pytorch-performance- with-graph-transformations/

  18. [18]

    OpenAI. 2021. Introducing Triton: Open-source GPU programming for neural networks. https://openai.com/blog/triton/

  19. [19]

    Phitchaya Mangpo Phothilimthana, Aditya Thakur, Rastislav Bodik, and Dinakar Dhurjati. 2016. Greenthumb: Superoptimizer Construction Framework. InPro- ceedings of the 25th International Conference on Compiler Construction (CC). ACM, 261–262

  20. [20]

    Phitchaya Mangpo Phothilimthana, Aditya Thakur, Rastislav Bodik, and Dinakar Dhurjati. 2016. Scaling up Superoptimization. InProceedings of the Twenty-First International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS). ACM, 297–310

  21. [21]

    PyTorch Foundation. 2025. Current and New Activation Checkpointing Techniques in PyTorch. https://pytorch.org/blog/activation-checkpointing- techniques/

  22. [22]

    James Reed, Zachary DeVito, Horace He, Ansley Ussery, and Jason Ansel. 2022. torch.fx: Practical program capture and transformation for deep learning in python.Proceedings of Machine Learning and Systems (MLSys)4 (2022), 638–651

  23. [23]

    Sohini Roychowdhury, Doris Wang, Qian Ge, Joy Mu, and Srihari Reddy. 2026. COFFEE: COdesign Framework for Feature Enriched Embeddings in Ads-Ranking Systems. InProceedings of the 19th ACM International Conference on Web Search and Data Mining (WSDM ’26), Web and Graph Workshop. Association for Com- puting Machinery. doi:10.48550/arXiv.2601.02807

  24. [24]

    Raimondas Sasnauskas, Yang Chen, Paul Collingbourne, Jubi Ketema, Jay Taneja, and John Regehr. 2017. Souper: A Synthesizing Superoptimizer.arXiv preprint arXiv:1711.04422(2017). arXiv:1711.04422

  25. [25]

    Eric Schkufza, Rahul Sharma, and Alex Aiken. 2013. Stochastic Superoptimization. ACM SIGARCH Computer Architecture News41, 1 (2013), 305–316

  26. [26]

    Eric Schkufza, Rahul Sharma, and Alex Aiken. 2014. Stochastic Optimization of Floating-Point Programs with Tunable Precision.ACM SIGPLAN Notices49, 6 (2014), 53–64

  27. [27]

    Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al

    Hugo Touvron, Thibaut ... Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023. Llama: Open and efficient foundation language models.arXiv preprint arXiv:2302.13971(2023)

  28. [28]

    Liang, Jason Ansel, Chunzhi Yang, Jade Nie, Yuzhen Huang, C.K

    Jiaqi Xu, Yanbo ... Liang, Jason Ansel, Chunzhi Yang, Jade Nie, Yuzhen Huang, C.K. Luk, Xiaodong Wang, Lu Fang, Menglu Yu, Jinwon ... Lee, Daohang Shi, and Flavio Sales Truzzi. 2023. Empowering Models with Performance: The Art of Generalized Model Transformation Approach. https://pytorch.org/blog/ empowering-models-performance/

  29. [29]

    Yisu Remy Yang, Phitchaya Mangpo Phothilimthana, Yihan Wang, Max Willsey, Saurabh Roy, and Jacques Pienaar. 2021. Equality Saturation for Tensor Graph Superoptimization. InProceedings of Machine Learning and Systems (MLSys), Vol. 3. 255–268

  30. [30]

    Zhichen Zeng, Wayne Xin Zhao, Yaliang Zhang, Hao Zhou, Jiayi Guo, Weidong Lin, Chih-Hsin He, Jiyan Yang, and Houfeng Ji. 2025. InterFormer: Effective Heterogeneous Interaction Learning for Click-Through Rate Prediction. InPro- ceedings of the 34th ACM International Conference on Information and Knowl- edge Management (CIKM ’25). Association for Computing ...

  31. [31]

    Li, Zeliang Chen, Weilin Zhang, Xiaohan Wei, and Ellie Wen

    Buyun Zhang, Liangchen Luo, Xi Liu, Jay ... Li, Zeliang Chen, Weilin Zhang, Xiaohan Wei, and Ellie Wen. 2022. DHEN: A Deep and Hierarchical Ensemble Network for Large-Scale Click-Through Rate Prediction. InProceedings of the 4th Workshop on Deep Learning Practice and Theory for High-Dimensional Sparse and Imbalanced Data (DLP-KDD). Association for Computi...