Lightning-AI / Lightning-AI/lightning-thunder

Provide debugging traces and options as a ENV variable or JIT option

Open
#304 2 comments 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

debugging enhancement
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

🚀 Feature

An environment variable that dumps out the various Thunder provided debug traces to a log file. This can have variable levels like
export THUNDER_DEBUG=<option>

0/'' : Disable
1/'trace' : Enable and dump Thunder generated trace. Can be limited to the trace after delete last used
2/'nvfuser_region' : Enable and dump nvFuser captured regions in addition to 1
3/'nvfuser_code' : Enable and dump nvFuser generated CUDA kernel code in addition to 1 and 2
4/'torch_compile_debug' : Enable the torch.compile debug logging (TORCH_COMPILE_DEBUG=1)

This is a narrow example of the possible debug log levels. Each of these logs can be in a different log file.

Motivation

To get the trace and other debugging information today, we need to add code that captures the trace and prints it after running a model iteration with the inputs.

  • This is cumbersome as the training code needs to be edited to enable tracing and re-edited when finished.
  • The ability to find when an iteration has finished and add the tracing code at the appropriate location may not always be possible as Thunder aims to compile more and more convoluted set of repositories. For example, when using libraries like
    Lightning Trainer, the user may want to just call model.train() but editing the iteration loop can be difficult.

cc - @mruberry

cc @carmocca @apaz-cli

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating the existing trace capture and printing flow, then review how Thunder, nvFuser, and torch.compile debugging are currently enabled. Define the supported debug levels and log destinations; done should allow tracing from an environment variable or JIT option without editing a training loop.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, devtools
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.