tvis

tvis: A Step-Through Debugger for PyTorch Training

tvis is an advanced debugger designed specifically for PyTorch training, allowing users to gain in-depth insights into their training process without needing to modify existing training scripts. By running your training script through tvis, critical information is captured for careful examination, helping to identify issues or optimize performance effectively.

Key Features:

• Comprehensive Recording:

  • Capture raw images or text strings, all transformations, tokenizers, and collated batches during the training process.
  • Track calls to every nn.Module and function, detailed in a call order that links back to the source line.
  • Access detailed inputs and outputs for each function call, including shapes, data types, statistics, and full values.
  • Monitor gradients of activations and their respective parameter updates to track model changes.
  • Analyze per-sample loss and observe how training progresses at each optimizer step.

• Detailed Observations:

  • Record the shapes produced by each line in your functions to understand data flow.
  • Visualize output decoded into class names or text through tokenization.
  • Gain insights into forward and backward call times measured accurately on the GPU, ensuring any overhead from tvis is excluded.

Usage Example:

Simply run your training script like usual, and switch to debugging with tvis:

python train.py --lr 3e-4                       # Regular training call without tvis

tvis run --steps 3 train.py --lr 3e-4           # Debug run capturing the first 3 steps

tvis open mycluster:~/proj/tvis_runs            # Open the results on a local browser

Visual Analysis:

• Interactive Replay Viewer: Visualize the training steps as they unfold, examining activations, losses, and parameter updates with ease.
• Layer Insights: Access detailed information on activations and gradients for each layer, adjust views to inspect individual neurons, and visualize loss metrics across the training run.
• Grad-CAM Integration: Generate true Grad-CAM visualizations to understand model behavior for specific classes.

Accurate and Reliable:

• Observations taken during training do not alter the outcomes. All captured data is consistent with what would be seen in a non-debugging run.
• The performance metrics are precise, allowing for confident insights and decisions based on the gathered data.

Limitations:

Currently, tvis is limited to a single process and device environment. Support for distributed training and specific features, such as torch.compile integration and certain tensor limits, are part of future developments.

For a comprehensive understanding, refer to the additional documentation files, including SCOPE.md and AGENTS.md. Whether debugging or optimizing training, tvis serves as a powerful tool for PyTorch developers.