NVTX can't annotate backward process of AI model training

[JuiceLemonLemon]

Hello,
Good afternoon!

When I annotate my training process by NVTX, then show in Nsight system, I found the forward process is annotated successfully, but the backward process is not. I don't know why.

                    with nvtx.annotate("backward path", color="yellow"):
                        loss.backward()

Tasks

Preview Give feedback

No tasks being tracked yet.

[evanramos-nvidia]

A simple explanation of what's happening: The backward pass happens in a different thread, so the NVTX ranges generated by those operations are not seen as decedents of the scoped range that encases the call to .backward().