Optimus: A Generic Operator-Level PyTorch Model Transformation Framework

2026-07-03

AI Summary

Industrial-scale recommendation and ranking models feature highly complex and continuously evolving architectures, rendering traditional optimization approaches—based on manual intervention or module-level rules—difficult to scale. This work proposes the first extensible and customizable operator-level automatic transformation framework integrated into PyTorch 2.x. By leveraging FX intermediate representation, the PT2 compiler, predefined pattern matching, and a greedy search algorithm, the framework achieves general-purpose model optimizations while strictly preserving computational semantics. Evaluated on real-world industrial recommendation models, the approach delivers up to 63% inference speedup, a 6% reduction in peak memory usage, and over 400 seconds of compilation time savings. The implementation has been open-sourced as part of PyTorch 2.x.

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.

Problem

• model transformation
• PyTorch
• graph optimization
• scalability
• deep learning

Innovation

• operator-level transformation
• PyTorch FX
• pattern matching
• model optimization
• machine learning compiler