KOEN
Read Original
#AI

Optimizing ML Workload Network Efficiency (Part I): Feature Trimmer

Optimizing ML Workload Network Efficiency (Part I): Feature Trimmer
01

Summary

How Pinterest Slashed ML Network Traffic by 50% with 'Feature Trimmer'

Maximizing GPU utilization by moving beyond compression to a 'Send What You Use' strategy.

This article explores Pinterest's engineering journey to resolve network bottlenecks in high-scale ML serving environments. By aligning feature transmission with specific model requirements through automated signature synchronization, they unlocked massive infrastructure savings and improved compute efficiency.

  • 01Identified network bandwidth as the primary bottleneck limiting GPU compute power.
  • 02Achieved an immediate 20% network reduction by implementing LZ4 compression.
  • 03Developed the 'Feature Trimmer' module to prune redundant features before transmission.
  • 04Used model signatures (module_info.json) as the source of truth for required inputs.
  • 05Synchronized trimmer configurations with model rollouts using existing staged delivery rails.

RECOMMENDATION

For infrastructure engineers managing high-traffic ML platforms, check if network I/O is limiting your GPU throughput. Implementing model-aware data pruning is a highly effective way to reduce costs and scale more efficiently.

The Problem

In Pinterest's ML serving system using a root-leaf architecture, the network bandwidth between root and leaf became a performance bottleneck, preventing full utilization of GPU compute resources and forcing scaling based on network usage.

The Solution

The team implemented 'Feature Trimmer,' which filters and sends only the required features based on the model signature, and enabled LZ4 compression in the fbthrift RPC framework to minimize data transmission.

The Result

LZ4 compression reduced network usage by 20%, and the Feature Trimmer is expected to cut it by ~50%, allowing a migration from expensive network-optimized instances (m6in) to standard instances (m6i) for cost efficiency.

Trade-off

The optimization introduced a 5% increase in CPU usage and a 5ms (~10%) increase in p90 latency, along with increased complexity in synchronizing model signatures across deployment pipelines.

03

Key Concepts

Concept · 01

Root-Leaf Architecture

A distributed system pattern where a 'Root' node handles CPU-bound tasks like feature fetching and preprocessing, while 'Leaf' nodes focus on GPU-bound model inference.

  • Optimizes resource utilization by separating CPU and GPU workloads.
  • Minimizes feature store QPS by sharing a central feature cache in the root.
Concept · 02

Model Signature

Metadata that defines the input and output specifications for a model version, acting as the definitive list of data required for execution.

  • Stored as module_info.json and serves as the permit list for the Feature Trimmer.
  • Exported automatically during model training and consumed by deployment pipelines.
Concept · 03

Feature Trimmer

A specialized module that filters the global feature set into a model-specific allowlist before data is sent across the network.

  • Ensures only essential data is transmitted, reducing network usage to model-specific needs.
  • Maintains robustness through versioned lookups and fallback mechanisms during rolling updates.