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 →
Optimus: A Generic Operator-Level PyTorch Model Transformation Framework
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
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)
- [§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.
- [§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)
- [Abstract] Abstract and §1: “over400second” and similar missing spaces appear repeatedly; a global proof-read is needed.
- [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.
- [§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”.
- [§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
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
free parameters (3)
- min_fuse_set_size =
5
- max_depth / min_depth (BFS) =
max_depth=5 (default)
- activation-quantization thresholds (size_in_mb, allowed_dtypes, use_scaling, exclude_primals) =
defaults used in evaluation
axioms (4)
- standard math Graph matching on ML DAGs is effectively tree inclusion and admits efficient recursive matching.
- domain assumption Semantic equivalence of each rewrite rule can be maintained by preserving inputs/outputs and by runtime numerical checks.
- domain assumption Production recommendation models contain many small, independent or split-cat operator motifs that dominate kernel-launch and compile cost.
- ad hoc to paper Hand-written atomic patterns plus sequential conflict-free application are sufficient; no global optimality search is required.
invented entities (1)
-
Optimus pattern library and greedy HorizontalOpt/VerticalOpt search
independent evidence
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
Reference graph
Works this paper leans on
-
[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...
2024
-
[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
doi:10.5555/3692070 2024
-
[3]
Sorav Bansal and Alex Aiken. 2006. Automatic Generation of Peephole Superop- timizers.ACM SIGARCH Computer Architecture News34, 5 (2006), 394–403
2006
-
[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
2011
-
[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
2018
-
[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
2021
-
[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
2020
-
[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)
2025
-
[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
2012
-
[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...
2025
-
[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
2020
-
[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
1976
-
[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
2019
-
[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
2002
-
[15]
Henry Massalin. 1987. Superoptimizer: A Look at the Smallest Program.ACM SIGARCH Computer Architecture News15, 5 (1987), 122–126
1987
-
[16]
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)
Pith/arXiv arXiv 2019
-
[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/
2022
-
[18]
OpenAI. 2021. Introducing Triton: Open-source GPU programming for neural networks. https://openai.com/blog/triton/
2021
-
[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
2016
-
[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
2016
-
[21]
PyTorch Foundation. 2025. Current and New Activation Checkpointing Techniques in PyTorch. https://pytorch.org/blog/activation-checkpointing- techniques/
2025
-
[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
2022
-
[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]
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
Pith/arXiv arXiv 2017
-
[25]
Eric Schkufza, Rahul Sharma, and Alex Aiken. 2013. Stochastic Superoptimization. ACM SIGARCH Computer Architecture News41, 1 (2013), 305–316
2013
-
[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
2014
-
[27]
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)
Pith/arXiv arXiv 2023
-
[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/
2023
-
[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
2021
-
[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]
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...
Pith/arXiv arXiv 2022
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.