I’m experiencing lots of issues with sharing large amounts of data between several processes in a distributed training setup using pytorch on my local machine. When starting the program, everything works fine but after some time the actions involved in transfering the data via multiprocessing.Queues slow down drasticly (.put(..)
and .get(..)
).
I’ve been trying to fix this for weeks now, tried several different approaches with torch.multiprocessing.Queue instead of the standard package, as well as using Manager.Lists. I tried individual Queues between processes as well as just one with multiple producers. I doesn’t seem to make a difference. I even started to monitor my hardware because I worry that it might be a RAM issue or something, but I have no idea what else to try. I also cannot provide the source code of the project or a minimal example for reproduction sadly.
Here is some more information.
- CPU: Intel(R) Core(TM) i9-14900K
- GPU: NVIDIA GeForce RTX 4090 24GB, Driver version: DCH 552.22
- RAM: 192GB DDR5 5200MHz
- Python 3.12.3
- pip 24.0
- torch 2.3.0+cu121
The architecture looks like this:
- there are 20 Actor processes, each simulating 6 environments and gathering sample data
- the samples are compressed and transmitted via Queue to the Replay Buffer
- the Replay Buffer samples batches and provides these via Queue to 3 Uncompressor processes
- the Uncompressors provide the batches finally to the Learner
- the Learner sends the updates weights periodically to the Actors and updates the priorities in the Replay Buffer
I added in-memory compression to reduce I/O and memory foodprint in exchange for compute. Originally the compression was done in the Replay Buffer but this didn’t perform well since the bytes sent from the Actors were around x5 larger slowing the Queues down further. Also, the Replay Buffer wan’t able to uncompress the samples on its own to keep up with the Learner, that’s why I had to add the Uncompressors.
I added performance logging. Consider the following observations:
The duration for learning stays about constant, but fetching the training batches from the Uncompressor Queue takes much longer (several seconds for 1 batch) the longer the training goes on (because it doesn’t have enough batches in it apprently).
Here you can see the time it takes the 3 Uncompressor processes to get the data from the Queue, measured like this:
t_get = time.time()
item = self.q_samples_in.get(timeout=5)
t_get = time.time() - t_get
It increases steadily up to several seconds per batch, even though the size (bytes) of the data transfered stays constant with each batch. One might also notice that in the beginning there is no issue but it rather starts getting worse after the first 500 batches or so.
The duration for adding the samples to the buffer as well as sampling isn’t an issue here. The duration increase for sampling batches is expected since the amount of data is growing.
One can see the issue on the right charts where the buffer growth and training speed slows down (this is correct, due to planned Actor throttling, because not enough samples are being trained on, because the Uncompressors receive not enough data from the Replay Buffer). One can also see on the left that the amount of time required to receive data from the Actors increases as well, even though the amount of time taken is neglectable.
The really weird part is the 3rd chart highlighted in blue, where the time taken to add a sampled batch to the Queue for the Uncompressors drops to basically 0s at around step 500 in this case (i.e. the samples are added to the Queue instantly?) exactly around that time the Uncompressors start to struggle receiving the data within reasonable time. While on the other end the Uncompressors take ages to receive the data to the point where the Learner has to wait for it. Note: In this run this phenomenon happened at around step 500 but in other runs it might be 150000, I don’t see a pattern.
What’s causing this? Why does it work fast in the beginning and then suddently slows down the longer the training goes on?
I monitored my hardware (CPU utilization, frequence, Memory, page files, temperature etc.) – nothing unusual as far as I am aware. I monitored RAM usage for memory leaks but it’s steady as one would expect. I also deepcopy
all the items before sending and on receiving them, as well as deleting unused variables.
Sometimes it manages to recover, but often it doesn’t and it keeps on getting worse. What I noticed is the thing mentioned above that at (in this case) around step 500 on all processes the data transfers start to slow down simultaneously, that’s why I’m worrying it might be an issue with my RAM (even though the PC is built new and barely 2 months old). Also, when running any other applications the system seems to perform as usual, indicating to me that this has something todo with the python process itself. I know Queues are pretty complex under the hood but I couldn’t say what logic might cause such a behaviour.
Does someone know what causes this and how I can fix it?
I’m thankful for any input!
— T