This document is relevant for: Trn2, Trn3

Debugging model code#

Debugging with Torch Eager on CPU mode#

With VLLM_NEURON_CPU_MODE=1, and --enforce-eager set, print statements are supported for all models using valid CPU mode configurations. Furthermore, with VLLM_ENABLE_V1_MULTIPROCESSING=0, normal pdb support is also enabled for models using valid CPU mode configurations and world_size=1.

To use pdb for world_size > 1, install forked-pdb with pip install fpdb, and insert like:

__import__('fpdb').ForkedPdb().set_trace()

Original Source

These flags are recommended to be set during the CPU development phase.

Eager mode is not yet supported on Neuron, but will come soon.

Debugging with torch.compile#

Python Debugger (pdb)#

Usage#

This approach is useful when you want to inspect all variables defined before a breakpoint.

torch.compile uses dynamo to build FX graphs before generating HLOs and compiling NEFFs. Since dynamo runs in multiple processes, the Python debugger does not work out of the box. You must redirect I/O streams from child processes to the parent debugging process using our custom context original_stdio.

def forward(
    self,
    hidden_states: torch.Tensor,
    positions: torch.LongTensor | None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attn_metadata: object | None = None,
):

    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    with original_stdio():
        breakpoint()

(EngineCore_DP0 pid=217) 2026-01-06 18:29:14,026 - INFO - neuron_model_runner.py:575 - Starting model forward pass with input_ids.shape=torch.Size([19]), positions.shape=torch.Size([19])
> /workspace/src/NxDI/src/nxdi/model/llama3/model.py(616)torch_dynamo_resume_in_forward_at_616()
-> with original_stdio():
(Pdb) p hidden_states.shape
torch.Size([19, 4096])
(Pdb)

You can also print tensor values and retrieve actual local ranks of workers when in CPU mode:

(EngineCore_DP0 pid=459) (Worker_TP0 pid=465) 2026-01-06 18:45:40,349 - INFO - neuron_model_runner.py:575 - Starting model forward pass with input_ids.shape=torch.Size([19]), positions.shape=torch.Size([19])
> /workspace/src/NxDI/src/nxdi/model/llama3/model.py(616)torch_dynamo_resume_in_forward_at_616()
-> with original_stdio():
(Pdb)             breakpoint()
(Pdb) get_tensor_model_parallel_rank()
1
(Pdb) p hidden_states
tensor([[ 1.5411e-03, -1.1597e-02, -3.0151e-02,  ...,  3.5889e-02,
          3.0136e-04,  3.6011e-03],
        [ 4.9805e-02, -2.2949e-02, -4.4141e-01,  ..., -6.1035e-02,
          4.5898e-02, -1.0376e-03],
        [ 6.1523e-02, -3.4180e-02, -1.6699e-01,  ..., -5.4688e-02,
         -8.7280e-03,  4.6082e-03],
        ...,
        [ 1.9836e-03, -2.9297e-02,  1.4648e-01,  ..., -1.0132e-02,
          7.3853e-03, -5.9204e-03],
        [ 2.0874e-02,  1.7773e-01,  7.9102e-02,  ...,  1.4062e-01,
         -1.5503e-02, -2.3804e-02],
        [ 1.9836e-03, -2.9297e-02,  1.4648e-01,  ..., -1.0132e-02,
          7.3853e-03, -5.9204e-03]], dtype=torch.bfloat16)
(Pdb)

Current Limitations#

Stepping through code in debug mode is not supported:

> /workspace/src/NxDI/src/nxdi/model/llama3/model.py(623)torch_dynamo_resume_in_forward_at_616()
-> hidden_states = self.self_attn(
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) (Pdb)
ERROR 01-06 21:17:45 [multiproc_executor.py:671] WorkerProc hit an exception.
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671] Traceback (most recent call last):
...
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]   File "/workspace/src/NxDI/src/nxdi/model/llama3/model.py", line 616, in forward
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]     with original_stdio():
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]   File "/workspace/src/NxDI/src/nxdi/model/llama3/model.py", line 623, in torch_dynamo_resume_in_forward_at_616
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]     hidden_states = self.self_attn(
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]   File "/usr/lib/python3.10/bdb.py", line 90, in trace_dispatch
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]     return self.dispatch_line(frame)
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]   File "/usr/lib/python3.10/bdb.py", line 115, in dispatch_line
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671]     if self.quitting: raise BdbQuit
(EngineCore_DP0 pid=277) (Worker_TP1 pid=285) ERROR 01-06 21:17:45 [multiproc_executor.py:671] bdb.BdbQuit

When using torch.compile, Dynamo intercepts and modifies bytecode execution. Dynamo also breaks the compute graph when it encounters code it cannot trace. In the example above, Dynamo generates a new function torch_dynamo_resume_in_forward_at_616() that throws an exception, causing the Python debugger to trigger self.quitting.

Printing During Tracing#