Minchul Kang Korea University [email protected]
Changyong Shin Korea University [email protected]
Jinwoo Jeong Korea University [email protected]
Hyunho Lee Korea University [email protected]
Younghun Go Korea University [email protected]
Gyeongmin Kim KT Corporation [email protected]
Gyeongsik Yang Korea University g [email protected]
Chuck Yoo Korea University [email protected]
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.4× 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.
I. I NTRODUCTION As deep learning (DL) models continue to grow in scale, training large-scale models on a single GPU becomes infeasible. For instance, training GPT-3 175B on a single NVIDIA A100 GPU would take about 14.8 years [1]. Consequently, distributed training leveraging multiple GPUs has emerged as the de facto standard. Accurate prediction of training time1 is essential for hardware resource planning [2], efficient job scheduling [3], [4], and cost estimation. We observe that training time increases by ∼2.4× depending on the floating-point precision setting (e.g., FP32, FP16, and mixed precision), implying that precision must be considered for accurate prediction. However, state-of-theart training time prediction studies assume a fixed precision setting [5], [6], leading to significant prediction errors (MAPE) of ∼147.85%. To this end, we propose a distributed training time predictor that supports arbitrary precision settings. By partitioning the model computation graph, the predictor automatically identifies operator-level precision, and incorporates communication overheads into the prediction. Our predictor achieves high accuracy, with 9.8% MAPE across various precision settings. II. BACKGROUND & M OTIVATION A. Background Parallelism strategies. Modern distributed training commonly employs a combination of data parallelism (DP), tensor 1 The training time in this work refers to single iteration time.
training time (ms)
arXiv:2604.16145v1 [cs.LG] 17 Apr 2026
Training Time Prediction for Mixed Precision-based Distributed Training
2000 1500
FP32 FP16 mixed
1000 500
OOM 0 (8,1,1) (4,1,2) (2,2,2) (1,8,1) (1,2,4) (4,2,1) (2,4,1) (2,1,4) (1,4,2) (1,1,8)
parallelism strategy (DP, TP, PP)
Fig. 1: Training time (ms) by precision settings. OOM denotes a GPU out-of-memory error.
parallelism (TP), and pipeline parallelism (PP). In DP [7], each GPU holds a full replica of the model and processes different input batches. TP [8] partitions the input dimensions of tensor operators within each model layer across multiple GPUs. PP [9] divides the model into sequential stages, each assigned to a different GPU, and processes micro-batches in a pipelined manner. The combination of these parallelisms significantly affects computation and communication time, as communication is known to be a bottleneck in distributed deep learning [10]. Floating-point precision. DL training uses various floatingpoint formats—FP32 for higher numerical accuracy and FP16 for faster computation. To balance FP32’s accuracy with FP16’s speed, mixed precision [11] is widely adopted, and recent large language models such as LLaMA [12], Qwen [13], and GPT [14] are trained under mixed precision. Mixed precision casts lower precision to compute-intensive operators (e.g., conv, matmul) and higher precision to operators requiring high numerical precision (e.g., softmax, reduction). Fig. 1 shows the training time of LLaMA 3.1–8B measured on eight NVIDIA H100 GPUs [15] interconnected with NVLink [16], under different parallelism strategies and floating-point precision settings. Training time varies by ∼2.4× across precision settings, highlighting its strong sensitivity to floating-point precision. This implies that accurate training-time prediction must account for the impact of floating-point precision.
III. D ESIGN Our approach is to predict the execution time of the model computation graph and communication overheads introduced by the precision setting and parallelism strategies (DP, TP, and PP) as formulated in Eq. 1. Here, Tcomp denotes the computation graph execution time of the given model, and Tdp , Ttp , and Tpp denote the communication overheads of DP, TP, and PP, respectively. T (d, t, p) = Tcomp (d, t, p) + Tdp (d) + Ttp (t) + Tpp (p) (1) Computation graph execution time. Given a model and a job config that specifies the precision setting, parallelism strategies, and hyperparameters (e.g., batch size), we first extract unique operators such as matmul, softmax from the model using torch.fx library [17]. Then, since the model computation graph can be distributed across multiple GPUs according to the given combination of parallelization strategies, we partition the graph into GPU-specific subgraphs based on the given DP, TP, and PP configurations, as detailed in Algorithm 1. Subsequently, we examine the precision of each operator within the subgraphs. If the training is configured for mixed precision, we hook the torch.amp [18] library to determine the casted precision of each operator. Otherwise, if the precision is fixed to FP32 or FP16, we simply use that predefined precision. Since the tensor shapes of operators vary depending on hyperparameters, such as batch size, we profile each operator’s execution time across both forward and backward computation graphs using the specific settings from
Ours
20
Ours
20 10
(8
,1 (4 ,1) ,1 (2 ,2) ,2 (1 ,2) ,8 (1 ,1) ,2 (4 ,4) ,2 (2 ,1) ,4 (2 ,1) ,1 (1 ,4) ,4 (1 ,2) ,1 ,8 )
0
,1 (4 ,1) ,1 (2 ,2) ,2 (1 ,2) ,8 (1 ,1) ,2 (4 ,4) ,2 (2 ,1) ,4 (2 ,1) ,1 (1 ,4) ,4 (1 ,2) ,1 ,8 )
0
vTrain
OOM
50 10
NeuSight
200 100
OOM
MAPE (%)
vTrain
MAPE (%)
NeuSight
100
(8
B. Motivation Prior work has proposed methods to predict the training time of DT. However, these methods rely heavily on static model computation graphs that do not capture precision variations such as mixed precision, which limits their applicability to general DT. NeuSight [5] decomposes the model computation graph of DL model into tiles, infers per-tile latency, and adds communication time for communication overheads. It assumes a fixed precision and, moreover, supports only a single parallelism strategy, not mixed parallelism (i.e., combinations of DP, TP, and PP). vTrain [6] converts the model computation graph to a CUDA graph, benchmarks kernel latencies, and incorporates communication overheads. However, like NeuSight, its prediction is based on a precision (FP32) and cannot cope with mixed or unseen precisions. We implement the prediction models of NeuSight and vTrain with FP32, and compare their predicted training times with the actual times observed under mixed precision and FP16 settings (Fig. 2a). The results demonstrate that both NeuSight and vTrain exhibit poor generalization to mixed and unseen precisions (i.e., FP16). Specifically, NeuSight’s prediction error increases ∼130.55%, while vTrain’s error reaches ∼147.85%. For unseen precisions, including mixed precision, input feature adjustments or reconstruction of the prediction model is inevitable.
(a) mixed precision.
(b) FP16.
Fig. 2: Prediction error (MAPE) of existing works and proposed approach. (X, Y, Z) on the x-axis represent the degrees of DP, TP, and PP. the job config. Finally, we aggregate the execution time of operators within each subgraph to obtain Tcomp . Algorithm 1 Partition model into GPU-specific subgraphs 1: Input: Model M , Job config J 2: Output: S – set of GPU-specific subgraphs 3: // Initialize subgraph set and layers per pipeline stage 4: S ← ∅; Lps ← ⌊layers(M )/J.P P ⌋ 5: for s = 0 to J.P P − 1 do // Iterate over pipeline stages 6: // Assign layers for stage s 7: Ls ← assign layers(M, s, Lps, J.P P ) 8: for d = 0 to J.DP − 1 do 9: for t = 0 to J.T P − 1 do // Create TP partition 10: L′s ← copy(Ls ) 11: // Slice weights if TP > 1 12: for all layer u in L′s do 13: for all weight w in u do 14: if J.T P > 1 and SliceNeeded(u, w)
then 15: 16:
L′s ← slice weight(w, t, J.T P ) S ← S ∪ {L′s }
17: return S
Communication overhead. In DP and TP, gradients are synchronized via all-reduce [19]. Accordingly, we derive Tdp and Ttp as the communication volume (i.e., total gradient size) divided by the link bandwidth Blink specified in the job config as formulated in Eq. 2 and Eq. 3. To determine Vdp , we aggregate the gradient sizes of all trainable parameters within the backward computation graph. Each size is precisely determined by the operator-level precision identified during the Tcomp prediction stage. Tdp = Vdp /Blink
(2)
In contrast, Vtp is determined by calculating the partial gradient sizes of operators partitioned in line 15 of Algorithm 1. Ttp = Vtp /Blink
(3)
The main communication overhead of PP (Tpp ) stems from the pipeline bubble, where stages wait for preceding computations to complete. We predict this overhead by scaling the predicted Tcomp by the PP degree specified in the job config, as formulated in Eq. 4.
Tpp = Tcomp × (P P − 1)
(4)
IV. E VALUATION We evaluate our predictor on the LLaMA 3.1-8B with C4 [20] dataset using eight NVIDIA H100 GPUs across all possible combinations of precision settings and parallelism strategies. In Fig. 2, the average MAPE values are 9.8% for mixed precision and 10.64% for unseen precision (FP16). Compared with existing approaches, our method achieves an ∼15.08× improvement, indicating robust and generalizable accuracy across diverse DT configurations. V. C ONCLUSION & F UTURE W ORK We propose a distributed training time predictor that accurately predicts training time for arbitrary precision settings. By partitioning the model computation graph and profiling operator execution with precision applied to the operator, our predictor achieves 9.8% average MAPE. We plan to extend this study to predict training time in multi-node heterogeneous GPU environments. VI. ACKNOWLEDGMENT This work was supported by KT(Korea Telecom)-Korea University AICT R&D Center. R EFERENCES [1] “Nvidia dgx nlp solution brief,” https://www.nvidia.com/content/dam/enzz/Solutions/gtcf22/dgx-pod/nvidia-dgx-nlp-solution-brief.pdf, 2022, accessed: 2026-03-01. [2] G. Yang, C. Shin, J. Lee, Y. Yoo, and C. Yoo, “Prediction of the resource consumption of distributed deep learning systems,” Proceedings of the ACM on Measurement and Analysis of Computing Systems, vol. 6, no. 2, pp. 1–25, 2022. [3] C. Shin, Y. Go, Y. Yoo, J. Jeong, J. Hwang, G. Yang, and C. Yoo, “Prediction-based gpu sharing for distributed training,” Future Generation Computer Systems, p. 108413, 2026. [4] Y. Go, C. Shin, M. Kang, J. Hwang, C. Yoo, and G. Yang, “Making sense of job preemption for distributed deep learning acceleration,” in 2026 63rd ACM/IEEE Design Automation Conference (DAC), 2026. [5] S. Lee, A. Phanishayee, and D. Mahajan, “Forecasting gpu performance for deep learning training and inference,” in Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, 2025, pp. 493–508. [6] J. Bang, Y. Choi, M. Kim, Y. Kim, and M. Rhu, “vtrain: A simulation framework for evaluating cost-effective and compute-optimal large language model training,” in 2024 57th IEEE/ACM International Symposium on Microarchitecture (MICRO). IEEE, 2024, pp. 153–167. [7] S. Li, Y. Zhao, R. Varma, O. Salpekar, P. Noordhuis, T. Li, A. Paszke, J. Smith, B. Vaughan, P. Damania, and S. Chintala, “Pytorch distributed: Experiences on accelerating data parallel training,” arXiv preprint arXiv:2006.15704, 2020. [Online]. Available: https://arxiv.org/abs/2006.15704 [8] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” arXiv preprint arXiv:1909.08053, 2019. [9] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and z. Chen, “Gpipe: Efficient training of giant neural networks using pipeline parallelism,” Advances in neural information processing systems, vol. 32, 2019. [10] Y. Yoo, G. Yang, C. Shin, H. Cho, W. Choi, Z. Niu, and C. Yoo, “Revisiting traffic splitting for software switch in datacenter,” Proceedings of the ACM on Measurement and Analysis of Computing Systems, vol. 9, no. 2, pp. 1–26, 2025.
[11] P. Micikevicius, S. Narang, J. Alben, G. Diamos, E. Elsen, D. Garcia, B. Ginsburg, M. Houston, O. Kuchaiev, G. Venkatesh, and H. Wu, “Mixed precision training,” 2018. [Online]. Available: https://arxiv.org/abs/1710.03740 [12] S. R. Cunningham, D. Archambault, and A. Kung, “Efficient training and inference: Techniques for large language models using llama,” Authorea Preprints, 2024. [13] J. Bai, S. Bai, Y. Chu, Z. Cui, K. Dang, X. Deng, Y. Fan, W. Ge, Y. Han, F. Huang et al., “Qwen technical report,” arXiv preprint arXiv:2309.16609, 2023. [14] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell et al., “Language models are few-shot learners,” Advances in neural information processing systems, vol. 33, pp. 1877–1901, 2020. [15] NVIDIA Corporation, “NVIDIA H100 Tensor Core GPU,” https://www.nvidia.com/en-us/data-center/h100/, 2024, accessed: 2025-05-30. [16] “Nvidia nvlink: High-speed gpu interconnect,” https://www.nvidia.com/en-us/data-center/nvlink/, accessed: 202602-27. [17] “torch.fx,” https://docs.pytorch.org/docs/stable/fx.html, accessed: 202603-01. [18] “Automatic mixed precision (amp),” https://docs.pytorch.org/docs/stable/amp.html, accessed: 2026-02-28. [19] P. Patarasuk and X. Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations,” Journal of Parallel and Distributed Computing, vol. 69, no. 2, pp. 117–124, 2009. [20] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu, “Exploring the limits of transfer learning with a unified text-to-text transformer,” Journal of machine learning research, vol. 21, no. 140, pp. 1–67, 2020.