Source-linked AI summary
Component Benchmark: Hierarchical Model Profiling for Large-scale Recommendation Systems
Dharak Kharod, Yuzhen Huang, Zhou Wang, Jackie Xu, Fuzail Khan, Jacky Zhou, Hao Yan, Lidong Zhao, Xizhou Feng, Yvonne Liu, Karthik Jayaraman, Praveen Ramachandran, Vishwa Karia, Yashasvi Makin
TL;DR
Recommendation models have component-level profiling needs that standard end-to-end and operator-level tools do not address. CB independently profiles submodules and presents results hierarchically; in practice, component analysis supported optimizations including a reported 3% net training and serving throughput gain.
Problem
Recommendation architectures contain varied operators and inefficiencies, while standard profiling tools do not attribute performance to the submodules practitioners use to reason about models.
Source: §I · p. 1Method
CB captures submodule inputs, profiles components independently, and organizes results in a hierarchical, plugin-extensible system.
Source: §I · pp. 1–2Results
3% net training and serving throughput gain followed re-authoring certain memory-bound GEMM/GEMV operations as elementwise kernels, enlarging the PT2 fusion region and improving memory access patterns.
Source: §V · p. 7Takeaways & Limitations
CB's component-level analysis can identify optimization opportunities that whole-model traces cannot attribute to a specific submodule.
Source: §I · p. 2
Abstract
from arXiv · showhide
Large-scale recommendation models pose distinct, under-explored profiling challenges. Most recommendation model architectures are structurally heterogeneous, intermixing memory-bandwidth-bound operations, small compute-bound dense layers, dynamic shapes from jagged categorical features, and low-arithmetic-intensity operations. Recommendation models evolve rapidly as modeling engineers experiment with compositions, often written without visibility into hardware execution characteristics. Standard profiling tools offer either end-to-end throughput or operator-level traces, but cannot attribute performance to the submodules that practitioners reason about. We present Component Benchmark (CB), a profiling system that independently characterizes each submodule performance in a hierarchical manner, providing a tree-structured, interactive visualization that brings performance clarity to ML practitioners. At its core, CB provides a simple yet extensible, submodule-based benchmarking framework with a plugin architecture that enables hierarchical performance analysis. These large-scale recommendation models are TB-scale, run on thousands of GPUs and ingest 100B examples per day. We demonstrate CB's effectiveness on common open sourced models and discuss how CB has been leveraged to accelerate modern recommendation model performance analysis and optimization.
I. INTRODUCTION
Recommendation models combine diverse, rapidly changing components, while existing profiling outputs do not readily connect performance to model components. Component Benchmark (CB) captures and profiles submodules independently, then presents their performance hierarchically.
I. INTRODUCTION
- Recommendation architectures remain an active design space, with inefficiencies from small operator dimensions, redundant compute, and repetitive kernels.The paper names FM, DCN, DHEN, Interformer, and HSTU among the architectures.Source: §I · p. 1
- Meaningful profiling must account for real data and complex training stacks, including categorical-feature sparsity and platform and framework differences.The paper also notes sequence-length distributions as relevant to large language model profiling.Source: §I · p. 1
- Existing tools offer end-to-end throughput or low-level operator traces, but often lack component semantics for attributing hotspots to model modules.Practitioners also want to break down metrics such as MFU by module.Source: §I · p. 1
- CB organizes profiling data around model components to support hierarchical performance analysis at module and component levels.This organization is intended to bring performance clarity to practitioners.Source: §I · p. 1
- CB uses PyTorch module hooks to capture submodule inputs and profile each component independently, reconstructing a tree-structured view for locating hotspots.The system is described as lightweight, extensible, and runnable on a single GPU.Source: §I · p. 1
- CB's plugin framework supports customized training stacks, input preprocessing, runtime profiling, and result visualization in an interactive hierarchical format.The paper discusses CB's analysis of representative DLRM-like models and its use in practical bottleneck analysis and optimization.Source: §I · p. 2
- Industrial recommendation models can be terabyte-scale, train on thousands of GPUs, and ingest about 10^11 examples daily, making component-level optimization consequential at fleet scale.The paper says even single-digit throughput improvements compound into fleet-wide capacity efficiencies.Source: §I · p. 2
II. RELATED WORK
Related systems provide GPU traces, end-to-end benchmarks, or recommendation serving measurements, but the paper identifies a gap in hierarchical component-level profiling. CB combines component extraction, profiling, and result visualization in a workflow tailored to heterogeneous training stacks.
II. RELATED WORK
- GPU profilers provide kernel- or operator-level traces, while MLPerf measures end-to-end benchmarks and DeepSpeed reports per-module FLOP counts rather than hierarchical multi-metric analysis.The paper says PyTorch 2 compilation can erase module boundaries, making operator-to-module mapping unreliable; CB independently profiles extracted submodules.Source: §II · p. 2
- DeepRecSys measures recommendation inference latency under serving conditions, whereas CB targets training-time, submodule-level profiling for architecture and kernel optimization.The paper states that, to its knowledge, CB is the first system designed for hierarchical component-level analysis of recommendation architectures.Source: §II · p. 2
A. Core Design
- CB's workflow connects framework adaptation, hook-based submodule and input extraction, preprocessing, component profiling, and result visualization.The visualizer aggregates, post-processes, and presents profiling results.Source: A. Core Design · p. 2
1) Leveraging PyTorch Hooks:
CB captures submodules with PyTorch hooks, profiles them recursively, and provides plugin interfaces for adapting and extending the profiling pipeline.
1) Leveraging PyTorch Hooks:
- CB uses PyTorch hooks to capture submodule inputs at selected component boundaries and locally profile components, avoiding unreliable attempts to recover module execution information from compiler traces.Hooks on submodules of interest capture inputs for reconstructing forward and backward behavior.Source: 1) Leveraging PyTorch Hooks · pp. 2–3
2) Recursive profiling:
- CB orders submodules by depth and profiles breadth-first, with configurable depth and filtering; when a parent cannot be profiled, it can continue with children and estimate parent behavior from their results.The fallback is described for cases such as a high-level module too large to fit on one GPU.Source: 2) Recursive profiling · p. 3
3) Plugin Design:
- CB's plugin design supports framework-specific model and input extraction, preprocessing before profiling, profiler plugins for collecting artifacts, and result plugins for post-processing and visualization.Built-in profiler plugins include trace collection and analysis, memory snapshots, and kernel-level analysis.Source: 3) Plugin Design · p. 3
1) Components Management:
CB adapts component extraction to training stacks and adjusts captured inputs to support profiling large models locally.
1) Components Management:
- Component Providers retrieve a training stack's top-level module and corresponding inputs by examining its data-loading and model-access logic.The paper reports integrating CB into recommendation training stacks by customizing these providers.Source: 1) Components Management · p. 4
- In most cases, one batch loaded through the training stack is representative; providers can also be extended to supply multiple batches for dynamic input shapes..Source: 1) Components Management · p. 4
- For recommendation models, the embedding hash size may be reduced so that lookups requiring several terabytes of memory can be profiled on one GPU..Source: 1) Components Management · p. 4
- CB extracts inputs using a smaller batch to reduce memory use, then replicates captured tensors to the original batch size during preprocessing.The paper says this preserves the input distribution and reduces out-of-memory errors when profiling large models locally.Source: 1) Components Management · p. 4
2) Profiles Management and Visualization:
Component Benchmark turns profiling data into interactive, hierarchy-organized reports, then demonstrates the system on DLRM and NanoGPT on one Nvidia H100 GPU.
2) Profiles Management and Visualization:
- Reports aggregate computation time, memory usage, and FLOPs by model hierarchy, with tables and visualizations that highlight performance hotspots.Users can filter, explore, and export results for further analysis or sharing.Source: 2) Profiles Management and Visualization · p. 4
- The reports support collaboration and model optimization by letting users explore and share profiling results.Results can be filtered and exported.Source: 2) Profiles Management and Visualization · p. 4
IV. EXPERIMENTAL ANALYSIS
- The demonstration benchmarks open-source DLRM and NanoGPT models on one Nvidia H100 GPU using MFU, Kineto trace, memory snapshot, visualization, and report-generation plugins.The experiment is presented as a demonstration of CB’s functionality and its use in recommendation-model performance analysis.Source: §IV · p. 4
A. Benchmark Results
The result table organizes profiling information by submodule and hierarchical depth, linking model structure to performance metrics and deeper profiling tools.
1) Result table:
- Each table row represents a submodule at a given depth, identified by a hierarchical path that traces the model from its top level to leaf submodules.This recursive organization supports performance analysis at each level and helps pinpoint submodule bottlenecks or inefficiencies.Source: 1) Result table · p. 4
- Component names show each module’s position in the model, while links provide access to traces, GPU snapshots, Perfetto, kernel profiles, and logs..Source: 1) Result table · pp. 4–5
- The table reports latency, static model size, TFLOPS/sec, MFU, CPU and GPU time, and estimated average bandwidth and utilization for each submodule.Latency is measured in milliseconds and model size in megabytes.Source: 1) Result table · pp. 4–5
2) Icicle View Visualization:
The icicle view maps model hierarchy to submodule performance, enabling users to inspect architecture and compare bottlenecks at multiple levels.
2) Icicle View Visualization:
- The interactive icicle represents submodules as nodes in a tree from the root module to leaf submodules, and supports flexible zooming.Figure 4 shows the DLRM benchmark view.Source: 2) Icicle View Visualization · p. 5
- The DLRM view summarizes four modules: sparse and dense arch, inter arch, and over arch, tracing embedding lookups and float features through interaction to a probability output..Source: 2) Icicle View Visualization · p. 5
- The sparse arch is almost entirely memory-bandwidth-bound, with low FLOPs and high bandwidth utilization per module.The icicle view color-codes submodules by the selected metric, MFU in this example, and users can click modules for detailed metrics.Source: 2) Icicle View Visualization · p. 5
- Figure 5 zooms into the inter-arch module to support inspection across its hierarchy..Source: 2) Icicle View Visualization · p. 5, Figure 5 · p. 6
- The inter-arch deep-fm and fm siblings differ sharply: deep-fm has 27.06% MFU and is compute-bound, while fm has 0.06% MFU and is bandwidth-bound.Deep-fm performs a large matrix multiplication; fm computes the factorization-machine cross term through three passes of small element-wise and reduction operations with large memory transfer.Source: 2) Icicle View Visualization · p. 5
- The inter-arch deep-fm.dense module is a container with no compute specific to itself; its performance is the latency-weighted mean over measured leaves..Source: 2) Icicle View Visualization · p. 5
- The dense arch is an MLP stack alternating compute-bound Linear layers at 25.91% and 20.74% MFU with bandwidth-bound ReLU activations at 0% MFU.The Linear-layer difference is attributed to shape: module.0 multiplies 1024x4096, while module.2 multiplies the narrower 4096x512.Source: 2) Icicle View Visualization · p. 6
3) Extensibility to LLMs:
CB’s provider abstraction is not specific to recommendation models: the authors applied it unchanged to NanoGPT, where the hierarchical view exposed contrasts among transformer components.
3) Extensibility to LLMs:
- CB requires only a top-level PyTorch nn.Module and a representative input batch, and its traversal, input capture, and plugins operate on the PyTorch module tree..Source: 3) Extensibility to LLMs · p. 6
- NanoGPT’s icicle shows 12 repeated transformer blocks accounting for most runtime, while the lm head carries a large share of FLOPs in a narrower width.The lm head is the model’s only green MFU cell, a visual signature of high arithmetic intensity.Source: 3) Extensibility to LLMs · p. 6
- Within a transformer block, attention has 10% MFU versus 20% for the MLP, motivating investigation of the attention kernel and its input Q,K,V dimensions.The per-submodule Perfetto trace could help confirm the kernel hypothesis.Source: 3) Extensibility to LLMs · p. 6
- Layer-norm, GeLU, and dropout blocks have 0–1% MFU; the authors suggest fusion might increase overall MFU and compute efficiency.The observations are preliminary and intended as demonstrations on open-source models.Source: 3) Extensibility to LLMs · pp. 6–7
V. RESULTS IN PRACTICE
Case studies show how CB's submodule-level views support practical performance analysis and optimization, including cases where whole-model traces cannot localize or size opportunities.
V. RESULTS IN PRACTICE
- 3% training and serving throughput gain followed re-authoring memory-bound GEMM/GEMV operations as elementwise kernels, enlarging PT2 fusion and improving memory access.The elementwise kernels were individually slower than vendor-library GEMM kernels.Source: §V · p. 7
- 83 ms to 50 ms model-isolation iteration time sized PT2's opportunity before integration and justified work that delivered a 34% QPS increase.Whole-model methods could not attribute compiler benefit to an uncompiled submodule.Source: §V · p. 7
- 2.8× memory and 39% more FLOPs than its predecessor were localized to a heavier event model and enlarged set-transformer blocks; restoring baseline dimensions recovered 10% QPS and 22% memory savings.Source: §V · p. 7
- Scaling down two low-performing submodules identified by CB's joint TFLOPS-and-latency view almost entirely eliminated the compute stall.Source: §V · pp. 7–8
VI. PRACTICAL LESSONS LEARNED
The practical lessons emphasize submodule-level profiling against real data, approachable visualization, and an extensible design; the conclusion presents CB as a tool for analyzing and optimizing large-scale recommendation-model training.
VI. PRACTICAL LESSONS LEARNED
- Profiling at the submodule level connects performance to actionable PyTorch modules, while real data and single-GPU analysis support profiling in large data systems.Source: §VI · p. 8
- CB's extensible design also supports numeric-equivalence checks and benchmarking submodules on new hardware to identify hotspots.Source: §VI · p. 8
- Engineers prefer illustrative, intuitive profiling experiences, and the hierarchical visualization has been widely welcomed.Source: §VI · p. 8
VII. CONCLUSION
- The conclusion presents CB as a simple, extensible system with recursive submodule-level profiling for recommendation-model performance analysis.It highlights the plugin architecture, core implementation, and result visualization as parts of the system.Source: §VII · p. 8