Table of Contents
Fetching ...

TTrace: Lightweight Error Checking and Diagnosis for Distributed Training

Haitian Jiang, Shaowei Zhu, Zhen Zhang, Zhenyu Song, Xinwei Fu, Zhen Jia, Yida Wang, Jinyang Li

TL;DR

TTrace introduces a principled differential testing workflow to detect and localize silent bugs in distributed training. By aligning intermediate tensors with a trusted reference through a canonical mapping and by estimating dynamic numerical tolerances via perturbation-based error analysis, it distinguishes bug-induced deviations from inherent FP round-off noise. The system comprises Canonical Mapping, Consistent Tensor Generator, Trace Collector, Equivalence Checker, and Perturbation-Based Error Estimation, enabling fine-grained localization to specific modules. Validation on Megatron-LM/TransformerEngine demonstrates detection of 14 silent bugs (11 existing, 3 new) with high efficiency, including support for BF16 and FP8, and real-world adoption in open-source frameworks. The work provides a practical, scalable approach to robustly verify distributed training correctness with minimal integration overhead.

Abstract

Distributed training is essential for scaling the training of large neural network models, such as large language models (LLMs), across thousands of GPUs. However, the complexity of distributed training programs makes them particularly prone to silent bugs, which do not produce explicit error signals but lead to incorrect training outcomes. Effectively detecting and localizing such silent bugs in distributed training is challenging. Common debugging practices based on monitoring training loss or gradient norm curves are indirect, inefficient, and provide no way to localize bugs. To address those challenges, we design and implement TTrace, the first systematic differential testing system for detecting and localizing silent bugs in distributed training. TTrace aligns intermediate tensors from distributed training with those from a trusted reference implementation. To properly compare the floating-point values in the corresponding tensors, we propose a novel mathematical analysis that provides a guideline for setting tolerances, enabling TTrace to distinguish bug-induced errors from numerical errors. Experimental results demonstrate that TTrace effectively detects 11 existing bugs and 3 new bugs in the widely used Megatron-LM framework, while requiring fewer than 10 lines of code changes. TTrace is effective in various training recipes, including low-precision recipes involving BF16 and FP8. Notably, a popular open-source training framework has already adopted the method proposed by TTrace in its development workflow.

TTrace: Lightweight Error Checking and Diagnosis for Distributed Training

TL;DR

TTrace introduces a principled differential testing workflow to detect and localize silent bugs in distributed training. By aligning intermediate tensors with a trusted reference through a canonical mapping and by estimating dynamic numerical tolerances via perturbation-based error analysis, it distinguishes bug-induced deviations from inherent FP round-off noise. The system comprises Canonical Mapping, Consistent Tensor Generator, Trace Collector, Equivalence Checker, and Perturbation-Based Error Estimation, enabling fine-grained localization to specific modules. Validation on Megatron-LM/TransformerEngine demonstrates detection of 14 silent bugs (11 existing, 3 new) with high efficiency, including support for BF16 and FP8, and real-world adoption in open-source frameworks. The work provides a practical, scalable approach to robustly verify distributed training correctness with minimal integration overhead.

Abstract

Distributed training is essential for scaling the training of large neural network models, such as large language models (LLMs), across thousands of GPUs. However, the complexity of distributed training programs makes them particularly prone to silent bugs, which do not produce explicit error signals but lead to incorrect training outcomes. Effectively detecting and localizing such silent bugs in distributed training is challenging. Common debugging practices based on monitoring training loss or gradient norm curves are indirect, inefficient, and provide no way to localize bugs. To address those challenges, we design and implement TTrace, the first systematic differential testing system for detecting and localizing silent bugs in distributed training. TTrace aligns intermediate tensors from distributed training with those from a trusted reference implementation. To properly compare the floating-point values in the corresponding tensors, we propose a novel mathematical analysis that provides a guideline for setting tolerances, enabling TTrace to distinguish bug-induced errors from numerical errors. Experimental results demonstrate that TTrace effectively detects 11 existing bugs and 3 new bugs in the widely used Megatron-LM framework, while requiring fewer than 10 lines of code changes. TTrace is effective in various training recipes, including low-precision recipes involving BF16 and FP8. Notably, a popular open-source training framework has already adopted the method proposed by TTrace in its development workflow.

Paper Structure

This paper contains 40 sections, 3 theorems, 3 equations, 9 figures, 3 tables.

Key Result

Theorem 1

A Pre-LN transformer layer $f_l$ is well-conditioned with high probability under standard initialization: its Lipschitz constant scales as $C_l \approx 1 + \mathcal{O}(l^{-\frac{1}{2}})$.

Figures (9)

  • Figure 1: The loss and gradient norm curve of a correct and buggy training. The bug can only be observed after 4k iterations from training loss curve.
  • Figure 2: The system overview of TTrace.
  • Figure 3: Example annotation clips for a GPT model. Module annotations specify which input/output tensors to trace. The numbers specify the tensor dimensions being partitioned by the particular parallelism.
  • Figure 4: The root cause of the bug in the case study.
  • Figure 5: Illustration of the layer index mapping in pipeline parallelism and virtual pipeline parallelism for canonical module name. Each colored box is one layer. The purple example means layer 0 in the 2nd virtual pipeline of the 1st pipeline stage maps to layer 4 in the reference.
  • ...and 4 more figures

Theorems & Definitions (3)

  • Theorem 1: Smoothness of Transformer Layers at Initialization
  • Theorem 2: Expected Activation Difference Bound
  • Theorem 3: Expected Gradient Difference Bound