Validate correctness of forward_pre_hook + tensor backward hook for OOM attribution

[abhinavsriva]

A lightweight execution-entry tracking mechanism was recently added to help attribute CUDA OOMs and runtime failures to the module execution context active at failure time.

Code: https://github.com/traceopt-ai/traceml/blob/main/src/traceml/utils/entry_hook.py

The current implementation relies on:

1. forward_pre_hook on leaf modules
2. Tensor-level backward hooks (Tensor.register_hook)
3. A shared execution-state pointer (EXECUTION_LAYER.current)

What to check if in a normal training run, does the reported:

1. module name
2. phase (forward / backward)
3. roughly match expectations when an OOM occurs?

Does it work in at least on a simple PyTorch training loop

Does it behave sensible when gradient accumulation is enabled (e.g. multiple forward/backward passes before optimizer step)?

Exact attribution is not required, the goal is to identify a useful execution context, not exact tensor causality.

[Ramshankar07]

I'll work in this

[Vishalkulkarni45]

@abhinavsriva can i work on this ?

[Vishalkulkarni45]

Test Environment:

• GPU: NVIDIA GeForce GTX 1650 (4GB VRAM)
• PyTorch: 2.5.1+cu121
• TraceML: 0.2.2 (from PyPI)

Test Setup:
8-layer fully connected network (fc1→fc8) with intentional OOM after layer 5.

import torch
import torch.nn as nn
from traceml.decorators import trace_model_instance, trace_step


class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 256)
        self.fc2 = nn.Linear(256, 512)
        self.fc3 = nn.Linear(512, 1024)
        self.fc4 = nn.Linear(1024, 512)
        self.fc5 = nn.Linear(512, 256)
        self.fc6 = nn.Linear(256, 128)  # OOM will trigger here
        self.fc7 = nn.Linear(128, 64)
        self.fc8 = nn.Linear(64, 10)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = torch.relu(self.fc5(x))
        _ = torch.zeros(15000, 15000, 15000, device=x.device)  # OOM
        x = torch.relu(self.fc6(x))
        x = torch.relu(self.fc7(x))
        return self.fc8(x)


if __name__ == "__main__":
    if not torch.cuda.is_available():
        exit(0)
    
    device = torch.device("cuda")
    model = TestModel().to(device)
    trace_model_instance(model, trace_execution=True)
    
    inputs = torch.randn(32, 128, device=device)
    with trace_step(model):
        model(inputs)

Result:

--- script crashed here ---
[TraceML] Last execution point: forward_fc5

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 12572.86 GiB. GPU 0 has a total capacity of 4.00 GiB of which 3.18 GiB is free. Of the allocated memory 13.75 MiB is allocated by PyTorch, and 8.25 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

@abhinavsriva

[abhinavsriva]

@Vishalkulkarni45
This checks for forward pass. Any idea on how would someone check in backward pass, somehow injecting error in compiutation graph

[passionworkeer]

For validation: 1) Test with a simple ResNet/VGG training loop to verify hook firing order, 2) For gradient accumulation, ensure EXECUTION_LAYER.current is saved/restored between accumulation steps, 3) Check if backward hooks fire in reverse order - you may need to track multiple modules. Use torch.cuda.memory_allocated() before/after each hook to correlate OOM with specific modules.

[Vishalkulkarni45]

i used below code to test backward hook

import torch
import torch.nn as nn
from traceml.decorators import trace_model_instance, trace_step


class DeepModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 256)
        self.fc2 = nn.Linear(256, 512)
        self.fc3 = nn.Linear(512, 1024)
        self.fc4 = nn.Linear(1024, 2048)
        self.fc5 = nn.Linear(2048, 1024)
        self.fc6 = nn.Linear(1024, 512)
        self.fc7 = nn.Linear(512, 256)
        self.fc8 = nn.Linear(256, 10)
    
    def _trigger_oom_during_backward(self, grad):
        # This fires DURING backward pass of fc5
        # Should show: backward_fc5
        # Actually shows: forward_fc8
        huge_tensor = torch.empty(1000000, 1000000, device='cuda')
        return grad
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        
        # Register hook on FC5 output
        x = torch.relu(self.fc5(x))
        x.register_hook(self._trigger_oom_during_backward)
        
        x = torch.relu(self.fc6(x))
        x = torch.relu(self.fc7(x))
        x = self.fc8(x)
        return x


if __name__ == "__main__":
    if not torch.cuda.is_available():
        exit(0)
    
    device = torch.device("cuda")
    model = DeepModel().to(device)
    trace_model_instance(model, trace_execution=True)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    
    inputs = torch.randn(32, 128, device=device)
    labels = torch.randint(0, 10, (32,), device=device)
    
    with trace_step(model):
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()  # OOM will happen HERE
    
    optimizer.step()

Crash occurs at: backward_fc5, but i am getting

[TraceML] Last execution point: forward_fc8
.
.
.
 huge_tensor = torch.empty(1000000, 1000000, device='cuda')
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3725.29 GiB. GPU 0 has a total capacity of 4.00 GiB of which 3.15 GiB is free. Including non-PyTorch memory, this process has 17179869184.00 GiB memory in use. Of the allocated memory 40.66 MiB is allocated by PyTorch, and 3.34 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management 

I looked into this and reason i think is PyTorch runs backward gradient hooks in a separate thread, but TraceML uses threading.local() which isolates data per-thread. The backward hook updates the execution layer in one thread, but the crash handler reads from a different thread's storage.

@abhinavsriva

[Vishalkulkarni45]

1. Test with a simple ResNet/VGG training loop to verify hook firing order,

TraceML does not print hook firing information during normal/successfull execution,

2. For gradient accumulation, ensure EXECUTION_LAYER.current is saved/restored between accumulation steps,

saved/restore is not implemented. EXECUTION_LAYER.current simply persists between accumulation steps and gets overwritten naturally by each forward pass . In the current implementation it can not tell at which layer of particular accumulator the error/crash occurred it only tells us at which layer crash happened

[abhinavsriva]

@Vishalkulkarni45 quite feasible that is it a thread issue (but by default the number of threads are set to 1 in Pytorch). Also you are regitering a hook which might be interferiong with traceml's hook.