Slow TDMPC training when doing multiple runs

[keyshavmor]

Hi Nicklas,

I hope you are well. I am writing to you about the TDMPC code that you authored.

I am currently working on a research project at the Institute of Neuroinformatics at ETH Zurich involving TD-MPC.

I have noticed that when running a single instance of TD-MPC, the training is quite fast. For reference, I have RTX3090 and H200 GPUS.

But when multiple runs for different tasks are executed, I notice a significant slow down (training takes half a day or some tasks like humanoid and dog take days) in the training process despite allocating sufficient memory in these large GPUs than can fit all 21 tasks in the RTX3090 for example, and also naturally on H200.

Do you have any idea why this might be happening? One of my guesses is these processes competing for the tasks of Mujoco that exclusively happen on CPU. So there is thread contention perhaps.

But, if you have any insights on how to make the training fast on something like H200, that would be really helpful and save us a lot of time.

We will give a big thanks to you in our paper if we ever reach that stage. :)

I look forward to hearing from you. Thank you!

Regards,
Keyshav Mor

[keyshavmor]

P.S: Similar issue was noticed for TDMPC2. The only caveat is RTX3090 seems faster than H200 but that doesn't make sense since H200 should be much much faster?

[nicklashansen]

Hi @keyshavmor , thanks for your interest in our work! The bottleneck is GPU utilization, not memory, so stacking multiple training runs will almost certainly slow down training. Additionally, unlike the RTX series, H200 is not designed for simulation / rendering so that may also be a factor.