Location Matters: LLM-Guided Joint Optimization of In-Network Aggregation Placement and Routing for DML Workloads

Abstract

In-network aggregation (INA) accelerates gradient aggregation in distributed machine learning (DML) by alleviating communication bottlenecks, but its effectiveness crucially depends on two location decisions: where to deploy INA functions and where to aggregate gradient flows. Most existing methods optimize INA placement and gradient flow routing independently, missing the advantages of joint optimization. This paper presents LLMINA, which leverages Large Language Models (LLMs) to automate the heuristic design for joint INA placement and gradient aggregation, aiming to minimize makespan (i.e., the total time required for all DML jobs to complete gradient aggregation). Directly using LLMs to generate end-to-end solutions is infeasible due to problem complexity and LLM limitations. Instead, LLMINA uses LLMs to generate heuristics for INA placement through an evolutionary process, and then applies an optimization-based heuristic for gradient routing that takes into account DML workload characteristics. Experiments across diverse network topologies and workloads show that LLMINA can significantly reduce makespan compared to state-of-the-art baselines. These results underscore that location matters for both INA deployment and aggregation, and highlight the potential of LLM-guided heuristic design for complex network resource optimization.