PRIME: Pseudo-Random Integrated Multi-Part Entropy for Adaptive Packet Spraying in AI/ML Data centers

2025-07-30

AI Summary

In AI/ML distributed training, low-entropy, bursty long flows cause network load imbalance and severe tail latency escalation. To address this, we propose a topology-aware, congestion-driven adaptive packet spraying mechanism. Our approach innovatively integrates pseudo-random round-robin scheduling with multi-level entropy control, augmented by ACK de-aggregation optimization, decaying temporal modeling, and congestion-aware dynamic rescheduling—effectively mitigating bufferbloat and enhancing stability under heterogeneous or degraded network conditions. Evaluated on a production-grade simulator, our method achieves 15% higher throughput under permutation traffic and reduces tail latency by 27% in link-degradation scenarios, compared to state-of-the-art baselines. These results demonstrate substantial improvements in network robustness and resource utilization for large-scale distributed training.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Distributed Search

Application Category

• Responsible Web: Human-perceived consequences of algorithmic deployment on the web
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

Abstract

Large-scale distributed training in production data centers place significant demands on network infrastructure. In particular, significant load balancing challenges arise when processing AI/ML workloads, consisting of low-entropy, bursty and long-lived flows. Existing solutions designed for Ethernet, such as Equal-Cost Multi-Path (ECMP) struggle to maintain high network utilization. While major industry players (e.g., Ultra Ethernet Consortium) and parts of academia have proposed packet spraying to enhance AI/ML workload performance, we argue that existing packet spraying solutions lead to buffer inflation over time, negatively affecting network performance. Specifically, when ACK coalescing is used, these solutions lead to stale information, degrading network performance. Additionally, in asymmetric network conditions- such as mix of ordered an unordered traffic, or link degradation and failures- existing packet spraying solutions often lead to increased tail latency. In this paper, we present the design and evaluation of PRIME, a pseudo-randomized round-robin approach to packet spraying that considers the network topology to optimize load distribution and performance. PRIME uses congestion as an indicator to re-balance the load. To this extent, PRIME takes into account various congestion signals, accounting for congestion severity, and their decay times to avoid network hotspots. We extensively evaluated PRIME using large-scale production-level simulator. Our results indicate that, compared to existing solutions, PRIME leads to up to 15% improvement for permutation traffic and up to 27% improvement in network degradation scenarios

Problem

• Addresses load balancing challenges in AI/ML data centers
• Reduces buffer inflation and stale information in packet spraying
• Improves tail latency in asymmetric network conditions

Innovation

• Pseudo-randomized round-robin packet spraying
• Network topology-aware load distribution
• Congestion signal-based re-balancing