Training Time Prediction for Mixed Precision-based Distributed Training

Abstract: Accurate prediction of training time in distributed deep learning is crucial for resource allocation, cost estimation, and job scheduling. We observe that the floating-point precision setting is a key determinant of training time, leading to training time variations of ~2.4x over its minimum. However, existing studies on distributed training time prediction rely on static model computation graphs that do not capture precision variations, including mixed precision. According to our experiments, training time prediction without considering precision results in significant prediction errors - reaching up to 147.85% in mean absolute percentage error (MAPE). To address this issue, we propose a precision-aware distributed training time predictor that achieves robust accuracy across diverse precision settings, including mixed precision, with 9.8% MAPE.

• The paper introduces a precision-aware training time predictor that significantly reduces error in mixed precision distributed training.
• It models both computation and communication overheads through dynamic operator-level profiling using torch.fx and torch.amp hooks.
• Experimental results show up to a 15× error reduction with an average MAPE of 9.8%, underscoring its practical value for efficient scheduling.

Precision-Aware Training Time Prediction for Distributed Deep Learning

Motivation and Problem Statement

Distributed training is the standard for scaling deep learning models, with parallelism strategies such as DP, TP, and PP. Accurate prediction of training time for single iterations is critical for resource allocation, scheduling, and cost estimation in production and research environments. The paper establishes that floating-point precision settings (FP32, FP16, mixed precision) are a significant determinant in training latency, producing up to a $\sim$2.4$\times$ difference in observed runtime. Existing approaches rely on static computation graphs that do not account for operator-level precision, particularly failing to adapt to mixed precision settings. Empirical analysis reveals mean absolute percentage error (MAPE) rates reaching 147.85% for models trained in mixed precision, indicating that static precision assumptions severely limit predictive utility.

Figure 1: Training time (ms) varies by precision settings, illustrating strong sensitivity and the necessity of precision-aware modeling.

Design of the Precision-Aware Predictor

The proposed predictor decomposes overall training time into computation graph execution and communication overheads, parameterized by precision and parallelism factors:

$$T(d, t, p) = T_{comp}(d, t, p) + T_{dp}(d) + T_{tp}(t) + T_{pp}(p)$$

• $T_{comp}$: Execution time of the distributed computation graph, where operator-level precision is profiled.
• $T_{dp}$/$T_{tp}$/$T_{pp}$: Communication overheads from data, tensor, and pipeline parallelisms.

Model partitioning is performed based on DP, TP, and PP degrees using torch.fx. Operator-level precision is determined dynamically, leveraging torch.amp hooks for mixed precision. Profiling is performed on both forward and backward operator executions to capture latency under actual batch size, precision, and device configurations. Communication overhead is modeled via aggregation and partitioning of gradient sizes according to precision and parallelism placement, divided by specified bandwidth.

Mixed Precision and Generalizability

Experimental results prove that mixed precision, now standard in LLM training, yields considerable runtime improvement and requires precision-aware prediction for generalizability. The predictor robustly identifies casted operator precisions and integrates them into latency estimation. The evaluation on a LLaMA-3.1–8B model across all precision/parallelism configurations demonstrates a strong numerical result: average MAPE falls to 9.8% for mixed precision, and 10.64% for FP16 alone, outperforming NeuSight and vTrain by an order of magnitude (approx. $15.08\times$ reduction in error).

Figure 2: Mixed precision error rates for prediction models, highlighting the superior accuracy of the proposed approach under realistic training conditions.

Implications and Future Developments

The empirical findings challenge the adequacy of static or precision-agnostic predictors, establishing that precision settings—especially in mixed precision—must be handled at operator granularity for predictive parity. Practically, this enables more precise resource assignment, improved cost forecasting, and better job scheduling policies in cluster and cloud environments. Theoretically, the model sets a precedent for high-fidelity runtime predictors capable of adapting to heterogeneous GPU architectures and evolving precision standards (e.g., BF16).

Future work, as proposed, aims to extend training time prediction to multi-node, heterogeneous GPU environments, where precision and communication dynamics are even more complex due to device variability and interconnect topologies. This direction is directly relevant for scaling trillion-parameter models and dynamic scheduler implementations.

Conclusion

This paper delivers a distributed training time predictor that accounts for arbitrary precision settings, including mixed precision, and multiple parallelism strategies. The approach partitions computation graphs, dynamically profiles operator precision, and integrates communication overheads, achieving 9.8% average MAPE and significantly outperforming previous state-of-the-art approaches. The implications extend to production-scale scheduling, cost estimation, and adaptive resource management. The methodology is positioned for further adaptation to heterogeneous and multi-node GPU clusters, aligning with the trajectory of scalable deep learning workloads.