Compatibility with 3090/4090 Series

[alfredgu001324]

Hi there,

I hope you are doing well! Thanks for the great work and all the help so far. Since my last issue has been a while I just decided to open a new one to see if there is any new update from your/any users side.

I tested it again on a fresh new machine with a 3090 but it reached the same performance bottleneck as my 4090 (around 4k SPS). For details please see this issue: #418 (comment)

Just wondering if anyone has experienced the same issue before and hoping to get some suggestions from you all. I have tried changing the parameters but they are not really helping. The PufferDrive however can reach a 240k SPS performance, just wondering what is the main difference between the two numbers (4k vs. 240k). Thanks!

[mpragnay]

Hey, here are some commands and things to check to diagnose the slow SPS on your RTX 4090. High GPU usage can be misleading; it could be waiting on the CPU, overheating, or facing a power limit.
Have a look at these to help figure out what's going on:

• First, monitor for overheating and power throttling during the experiment:
  nvidia-smi -q -d TEMPERATURE,POWER
  In the output, check the Power Draw and compare it to the Power Limit. If they are very close, the card is power-throttling.
  Also, check the Thermal Slowdown or Gpu Idle reasons. If they are listed, thermal throttling is occurring.

• Next, check the actual clock speeds:
  nvidia-smi -q -d CLOCK
  Look at the Graphics clock speed (MHz). For a 4090, this should be high (often above 2500 MHz) under load. If it's much lower, something is limiting its performance.

• Check if the CPU is the bottleneck:
  Run nvtop in a terminal while the experiment is running.
  Keep an eye on the CPU usage, specifically if any single core is hitting 100%. SPS is a simulation step rate, so if the CPU can't keep up with calculations, it could be the bottleneck, even with high GPU usage. I don't think this is possible, as the Madrona engine is based on a GPU Megakernel architecture that reduces CPU-GPU communication. Just remember that if your CPU has very few cores, it could also be a bottleneck(data loading and training depend on this)

You can run this file: https://github.com/Emerge-Lab/gpudrive/blob/profiling/profiling/profile_madrona_gpudrive.py in the profiling branch, that would give better insights on whether SPS scales with num_envs

[alfredgu001324]

Thanks for the detailed suggestions! I truly appreciate it.

I quickly tried them, and I found that nvidia-smi -q -d TEMPERATURE,POWER and nvidia-smi -q -d CLOCK both show very normal values and everything is fine. For the CPU problem, as stated in the previous link as well, even though I have 64 cores, it seems that only one of them is max out and the rest are all idling. Seems like this is a single-threaded process.

I also tried the script that you mentioned, it seems that the SPS does scales with num_envs. So the simulator itself has no problem I think.

I later tried profiling it again using python -m cProfile -o ppo_profile.out baselines/ppo/ppo_pufferlib.py and observe the following

(gpudrive) aa@aa-System-Product-Name:~/project/gpudrive$ python -m pstats ppo_profile.out
Welcome to the profile statistics browser.
ppo_profile.out% sort time
ppo_profile.out% stats 20
Wed Sep 10 11:02:42 2025    ppo_profile.out

         10654875 function calls (10281258 primitive calls) in 348.713 seconds

   Ordered by: internal time
   List reduced from 23847 to 20 due to restriction <20>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     3135  268.084    0.086  268.480    0.086 /home/aa/project/gpudrive/gpudrive/env/env_torch.py:595(step_dynamics)
        1   24.637   24.637   24.637   24.637 /home/aa/project/gpudrive/gpudrive/env/base_env.py:161(_initialize_simulator)
      137   15.002    0.110   15.114    0.110 /home/aa/project/gpudrive/gpudrive/env/env_torch.py:392(reset)
     4159   14.355    0.003   14.355    0.003 {built-in method torch._C._cuda_synchronize}
      124   10.530    0.085   10.530    0.085 {method 'acquire' of '_thread.lock' objects}
     3135    1.605    0.001  288.278    0.092 /home/aa/project/gpudrive/gpudrive/env/env_puffer.py:235(step)
    32823    0.750    0.000    0.750    0.000 {built-in method torch._C._nn.linear}
      512    0.516    0.001    0.516    0.001 {method 'run_backward' of 'torch._C._EngineBase' objects}
     3409    0.466    0.000    0.483    0.000 /home/aa/project/gpudrive/gpudrive/datatypes/roadgraph.py:281(__init__)
     3647    0.347    0.000    0.347    0.000 {method 'logsumexp' of 'torch._C.TensorBase' objects}
     4558    0.314    0.000    0.314    0.000 {built-in method marshal.loads}
     3135    0.292    0.000    0.292    0.000 {built-in method torch.multinomial}
     3409    0.289    0.000    0.292    0.000 /home/aa/project/gpudrive/gpudrive/datatypes/observation.py:177(__init__)
  252/231    0.268    0.001    0.330    0.001 {built-in method _imp.create_dynamic}
    30842    0.265    0.000    0.265    0.000 {method 'sum' of 'torch._C.TensorBase' objects}
    13908    0.259    0.000    0.259    0.000 /home/aa/project/gpudrive/gpudrive/utils/geometry.py:15(normalize_min_max)
        1    0.257    0.257    0.257    0.257 {method 'read' of '_ssl._SSLSocket' objects}
        1    0.248    0.248  345.904  345.904 /home/aa/miniconda3/envs/gpudrive/lib/python3.11/site-packages/typer/main.py:674(wrapper)
    21945    0.242    0.000    0.242    0.000 {method 'cpu' of 'torch._C.TensorBase' objects}
     3135    0.231    0.000    0.238    0.000 /home/aa/project/gpudrive/gpudrive/integrations/puffer/ppo.py:618(<listcomp>)

Step_dynamics is taking most of the time. But this still would not explain the 4k SPS here vs. 230k SPS on PufferDrive. So I was wondering whether this is a pufferlib integration issue that slows down the training.

Uhmm maybe I should try the Sb3 integration to see if there is any issues there.

Thank you so much for your help again and really appreciate your time and help!