pytorch / pytorch/rl

Long GPU idle times in loss forward pass

Open
#1,954 1 comment 0 reactions 1 assignee View on GitHub

@vmoens is already working on this.

Since Feb 22, 2024.

  • #1958 by @vmoens — open
bug
Dominant language
Python
Stars
3.6k
Forks
487
Avg merge
1d 1h
Merged PRs (30d)
207

Description

I have just implemented an RL agent for a custom environment (wrapped into a TorchRL env). I am trying to reimplement the RAPS algorithm using SAC and for that I have used the SACLoss provided by TorchRL.
Here, I mainly stuck to the examples/sac for structuring my code and setting everything up.

However, training the agent, I experienced bad GPU utilization. Profiling, I found that what takes most time is the SACLoss forward pass. I then proceeded using nsys profile in order to investigate further into this forward pass. In the screenshot attached, I have recorded a single representative forward pass through the SACLoss (after some warmup iterations).
You can see that the GPU is only utilized for short times at the start and end of the forward pass and a sligthly longer period in the middle. Is this behavior expected?
I also notice the CPU process where Python is running to be at 100. I am not sure what is causing this as all my networks are on GPU and there shouldn't be much also running during the Loss forward pass, right?

If all this not expected, how can I proceed in order to increase utilization (or first find out what is causing low utilization)?

Screenshots

Nvidia Nsight Systems Screenshot

Environment:

  • GPU: RTX 3070 Mobile
  • Python inside Conda Environment
    • Pytorch 2.1.0
    • torchrl 0.3.0

Contributor guide

Open the contributing guide

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.