UNIFERENCE: A Discrete Event Simulation Framework for Developing Distributed AI Models
Authors: Doğaç Eldenk, Stephen Xia
Organizations: Department of Computer Science, Northwestern University, Evanston, United States · Department of Computer Engineering, Northwestern University, Evanston, United States
Developing and evaluating distributed inference algorithms remains difficult due to the lack of standardized tools for modeling heterogeneous devices and networks. Existing studies often rely on ad-hoc testbeds or proprietary infrastructure, making results hard to reproduce and limiting exploration of hypothetical hardware or network configurations. We present UNIFERENCE, a discrete-event simulation (DES) framework designed for developing, benchmarking, and deploying distributed AI models within a unified environment. UNIFERENCE models device and network behavior through lightweight logical processes that synchronize only on communication primitives, eliminating rollbacks while preserving the causal order. It integrates seamlessly with PyTorch Distributed, enabling the same codebase to transition from simulation to real deployment. Our evaluation demonstrates that UNIFERENCE profiles runtime with up to 98.6% accuracy compared to real physical deployments across diverse backends and hardware setups. By bridging simulation and deployment, UNIFERENCE provides an accessible, reproducible platform for studying distributed inference algorithms and exploring future system designs, from high-performance clusters to edge-scale devices. The framework is open-sourced at https://github.com/Dogacel/Uniference.
Figures & tables
Framework
Use Case
Network Sim
Integration Complexity
Custom Algorithms
DistSim
Analyze training time and device activities in distributed DNN training
Profiled
N/A
N/A
SimAI
Simulator for AI large-scale training for analyzing performance data
Advanced
High
Not Supported
EdgeAISim
Simulating and modelling AI models in edge computing environments
Simple
Medium
Not Supported
Vidur
LLM inference simulator for deployment conf.
Profiled
Medium
Partial Support
HERMES
Event-driven, multi-stage inference pipeline simulation and optimization
Simple
Low
Partial Support
simPy
Discrete-event simulation
(ns.py)
Medium
Supported
TABLE I: Comparison of existing tools.
Operation
Hops
Total Bandwidth
Description
broadcast
N−1
(N−1)×M
One node sends data to all other nodes in the network
all_gather
N−1
N×(N−1)×M
Each node shares its data with all other nodes
all_reduce
2(N−1)
2(N−1)×M
Performs reduction operation (sum, max, etc.) across all nodes
send
1
M
Point-to-point transfer from one node to a specific destination
recv
1
M
Point-to-point receive from a specific source node
TABLE II: Uniference ’s Supported Distributed Communication Operations
Fig. 1: Discrete Event Simulation Overview
coef
std err
PI
R 2
Const
0.054
0.000
-
0.779
Max tokens
12.19
0.001
24.91
0.325
Yield count
0.191
0.000
0.000
0.999
Device count
0.000
0.000
0.000
0.999
TABLE III: Many token prediction experiment regression results
Hardware
Backend
R 2
MAPE
HPC
NCCL
0.9704
0.1783
HPC
Gloo
0.9999
0.0248
Jetson Orin
Gloo
1.0000
0.0091
TABLE IV: Model performance comparison (R 2 and MAPE).
Hardware
Backend
R 2
MAPE
HPC
NCCL
0.78-0.90
0.2065-0.3626
HPC
Gloo
0.9860
0.0536
TABLE V: Estimation accuracy of real time spent on network
Fig. 2: Uniference Trace Tool showing Voltage [ 16 ] (top) and the improved Kilovolts algorithm (bottom) with the communication overlap on xq .
Distributed inference has become an indispensable part of deploying medical models under practical latency, memory, and throughput constraints. Although modern frameworks improve serving efficiency through tensor parallelism, mixed precision, kernel fusion, and multi-device communication, they are generally assumed to preserve the behavior observed during centralized HuggingFace evaluation. This assumption creates an evaluation-deployment mismatch: a model may pass offline evaluation but produce a different output after the execution stack changes. To address this mismatch, we propose a testing framework and an improved, distributed-execution-sensitive medical-model benchmark that evaluates the same checkpoint and input under a centralized HuggingFace reference and matched distributed deployments. Extensive experiments across language, vision, and multimodal medical models show that execution changes can produce measurable output disagreements. Across supported visual settings, the test success rate ranges from 0.21 to 0.43 for single-modality models and from 0.32 to 0.98 for multimodal models. The benchmark is aimed at extending medical-model evaluation from capability and security to evaluation-deployment consistency.
Yifei Wang, Xiaohan Zhang, Youtao Ding +4
Shanghai Jiao Tong University, China · Beihang University, China · Nanyang Technological University, Singapore +1
Distributed machine learning (ML) is a key paradigm for today's large-scale artificial intelligence applications. As model inference arises as an important use case, faithful modeling of latency-sensitive collective communication has never been more important. Capturing the device architecture and modeling control and data paths at high fidelity is therefore a necessity today. Having a common, detailed representation for distributed ML infrastructure is also crucial. We revisit the promising open-source, community-driven simulator: ASTRA-sim. In this work, we identify limitations of the current ASTRA-sim simulator and augment it with new features. To this end, we enable fine-grained, high-fidelity simulation with a standardized infrastructure representation, opening new design space exploration opportunities. We propose the simulation at cache-line-sized load-store granularity, with a detailed graphics processing unit (GPU) execution model, to balance simulation scalability and fidelity. We also introduce InfraGraph, a standardized representation to capture distributed ML network infrastructure in detail. Using the updated ASTRA-sim 3.0 simulator, we showcase interesting design space explorations for designing optimized collective algorithms, network requirements, and GPU architectures.
William Won, Jinsun Yoo, Tuan Ta +16
AMD Research and Advanced Development · Georgia Institute of Technology · Keysight +1
Local execution of AI on edge devices is important for low latency and offline operation. However, deploying models on diverse hardware remains fragmented, often requiring model conversion or complete reimplementation outside the PyTorch ecosystem where the model was originally authored. We introduce ExecuTorch, a unified PyTorch-native deployment framework for edge AI. ExecuTorch enables seamless deployment of machine learning models across heterogeneous compute environments. It scales from embedded microcontrollers to complex system-on-chips (SoCs) with dedicated accelerators, powering devices ranging from wearables and smartphones to large compute clusters. ExecuTorch preserves PyTorch semantics while allowing customization, support for optimizations like quantization, and pluggable execution "backends". These features together enable fast experimentation, allowing researchers to validate deployment behavior entirely within PyTorch, bridging the gap between research and production.
Mergen Nachin, Digant Desai, Sicheng Stephen Jia +36