ModTrans: Translating Real-world Models for Distributed Training Simulator

Abstract: Large-scale distributed training has been a research hot spot in machine learning systems for industry and academia in recent years. However, conducting experiments without physical machines and corresponding resources is difficult. One solution is to leverage distributed training simulators, but current ones like ASTRA-sim do not support importing real-world developed models, which poses challenges for ML researchers seeking to use them. Based on this challenge, we developed ModTrans, a translator supporting format translation from any real-world model to the ASTRA-sim simulator's input, removing the barrier between machine learning experts and machine learning system researchers. The experiment results show that ModTrans's cost is negligible.

• The paper introduces ModTrans, a system that automatically translates real-world neural network models into simulator-compatible descriptions using ONNX, eliminating manual intervention.
• The paper demonstrates negligible translation overhead with sub-second conversion times for models like ResNet50 and VGG, ensuring efficiency and accuracy.
• The paper shows that decoupling model design from system simulation facilitates scalable, reproducible distributed training research by minimizing configuration errors.

ModTrans: ONNX-Based Real-World Model Translation for Distributed Training Simulation

Motivation and Problem Scope

Large-scale distributed training of neural networks imposes unique system and algorithmic challenges, primarily stemming from the formidable resource demands—network, memory, and computation—of contemporary deep models. Industrial platforms such as Google's PaLM have demonstrated training over thousands of accelerators interconnected at 81 Tbps, highlighting the infeasibility for most academic labs to conduct comparable experiments on physical clusters. Distributed training simulators (DTS) have emerged as a necessary abstraction layer, offering performance, communication, and scheduling what-if analysis on virtualized hardware topologies. However, frameworks like ASTRA-sim accept only hand-crafted DNN description files, tightly coupling ML experts and system designers. This rigid workflow and the requirement for manual input of layerwise compute/communication characteristics both hinder scalable experimentation and introduce potential configuration errors.

ModTrans System Overview

ModTrans is introduced as a layer-wise model translation system that bridges the gap between real-world neural network architectures (from PyTorch, TensorFlow, etc.) and the specific input format required by ASTRA-sim. ModTrans leverages ONNX as an IR to obtain a fully serialized and standardized representation of the network computation graph, circumventing the heterogeneity of native training libraries. The translation workflow consists of:

1. Deserialization of the ONNX representation to obtain the layer graph and associated tensor metadata.
2. Automated extraction of layerwise variable counts, data types, and model sizes.
3. Formatting of these data into ASTRA-sim DNN description files, including all parameters needed for compute/communication emulation.

Notably, ModTrans also provides interoperability with the ONNX Model Zoo for automatic retrieval and conversion of canonical architectures (e.g., ResNet50, VGG), further generalizing simulator input to a broad class of production models.

Figure 1: Schematic of the deep learning platform design space, encompassing both software and hardware axes, contextualizing the placement of simulators like ASTRA-sim.

Figure 2: Architecture of ASTRA-sim, showing the layered simulation approach across network, system, and DNN workload levels.

By removing manual intervention, the system decouples machine learning research from distributed systems engineering, enabling a modular, reproducible, and scalable system evaluation pipeline.

Workflow and Interface Integration

ModTrans interfaces with ONNX via Python bindings, exploiting the protobuf serialization to efficiently traverse model graphs. After parsing, computation times for each layer are modeled using SCALE-sim for accurate system simulation, while communication characteristics are inferred directly from tensor shapes and parameterization. This design aligns closely with ASTRA-sim’s workload model, which necessitates detailed layerwise timing and communication data for fidelity.

Figure 3: Example of an ASTRA-sim DNN description file, demonstrating how model topology and per-layer characteristics are encoded for simulation.

Figure 4: Visualization of a model in ONNX graph format, illustrating the IR as a sequence of computation and dataflow nodes.

ModTrans achieves a clear abstraction barrier—ML experts focus on model design and training, while system researchers consume the translated graph descriptions for distributed simulation without cross-domain friction.

Empirical Results and Runtime Analysis

The translation overhead introduced by ModTrans is empirically negligible: less than 1 second per model, irrespective of architecture size or layer count (e.g., 0.1s for ResNet50, 0.8s for VGG16/VGG19). This property ensures that ModTrans is not a bottleneck in the simulation workflow, allowing for batch conversion or online integration with systems research toolchains.

Figure 5: Execution time analysis for ModTrans over representative models, indicating sub-second translation latency even for large DNNs.

The output DNN descriptions match, byte-for-byte, those produced independently and manually for ASTRA-sim reference models, providing a strong self-consistency and validation signal for correctness. This equivalency enables immediate drop-in replacement in existing simulator stacks.

Theoretical and Practical Implications

The introduction of ModTrans has several key implications:

• Decoupling ML and System Design: ModTrans establishes a well-defined interface between machine learning model development (at the framework/ONNX/export level) and system performance research (at the simulator description file input level). This separation allows parallel, modular system evaluation and reproducibility.
• Scalability and Reproducibility: The system supports wide coverage over design and model spaces, enabling extensive ablation, system design, and hardware-software co-optimization studies. Researchers can efficiently benchmark new system scheduling and parallelization techniques with arbitrary models from the ONNX ecosystem.
• Simulator Generalization: The architecture of ModTrans is extensible; any DTS requiring layerwise workload input (not just ASTRA-sim) can integrate this workflow. The translation logic is not coupled to a specific simulator backend, motivating broader ecosystem adoption.
• Error Minimization: By automating model-to-description translation, the risk of manual calculation or configuration errors across hundreds of layers is essentially eliminated.

Extensions and Speculation for AI Research

As distributed training platforms and simulation tools proliferate, automated translation layers like ModTrans will become central to systematic co-design and benchmarking efforts. Several extensions are anticipated:

• Generalized Simulator Input Adapters: ModTrans can be adapted to target other simulators as long as their workload interfaces are compatible, expanding reach beyond ASTRA-sim.
• Automated Parameter Profiling: Integration with hardware profilers and simulators can be extended to infer dynamic compute/communication costs under real workload conditions, improving simulation fidelity for systems-centric ML research.
• Support for Model Variants and Non-Standard Layers: Expansion to cover non-standard or dynamic architectural features in ONNX (e.g., user-defined operators, conditional dataflow) is a natural next step.
• Integration into Continuous Evaluation Pipelines: ModTrans enables continuous evaluation frameworks for hardware-aware neural architecture search and large-scale model benchmarking under diverse system configurations.

Conclusion

ModTrans provides an efficient, correct, and generalizable mechanism for translating real-world ML models into distributed training simulator input formats, addressing a notably persistent system research bottleneck. The system dramatically streamlines the workflow across ML and systems domains, with strong runtime and correctness guarantees, and establishes a template for further automation in the AI hardware/software co-design stack. The long-term implications include improved reproducibility, modularity, and scalability of distributed training research (2604.01607).