Meta RecIR: Single Codebase for GPU and TPU Training of PyTorch Recommendation Models

Meta has released the RecIR compiler in alpha version, which allows PyTorch-based recommendation models to be trained on GPUs and TPUs
Training steps written in PyTorch are captured as FX graphs and converted for GPU (TorchRec, Inductor) or TPU (torchax, JAX/XLA) use
Currently, the only supported reference is the training of the DLRMv3 model in the generative-recommenders repository, and the inference path is experimental
A tool is provided to verify numerical parity between GPUs and TPUs, while performance metrics (throughput) are not disclosed

Technical Architecture and Features

The model remains a standard torch.nn.Module and the optimizer follows PyTorch standards, while RecIR handles capture, graph transformation, state management, distribution, and backend conversion
On TPUs, Pallas ring shift embedding lookups and row-wise sharding are used, while GPUs utilize FBGEMM TBE and Inductor
To address XLA's fixed-shape requirements and the variable-length UIH (User Interaction History) problem, a Token-budget batching approach is applied

Installation is done via source checkout on Python 3.12 or higher without a PyPI package, and DLRMv3 training requires a combination of PyTorch 2.10, TorchRec 1.4, and FBGEMM 1.5
Python data-dependent control flow cannot be automatically captured, and recompilation is required when batch size or table shapes change
It is unsuitable for teams aiming to immediately deploy arbitrary recommendation models to production due to API instability, but is well-suited for comparative training of large-scale HSTU/DLRMv3-family models