Hi Yi Wang,
Thanks for your good suggestions. First, I need to correct a mistake in my observation:
use mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True) also shows memory leak problem locally, same with use run_worker(rank, world_size, num_split). I only observe the master process, which does not shows growth. But the worker part shows growth, which has the same result with using run_worker(rank, world_size, num_split) and also could be observed in the tutorial for the ResNet.
In my script for transformer models, both mp.spawn() method and run_worker() have similar result :

Set  the num_worker_threads=2, memory does not increase
Set the  num_worker_threads=16, memory first increases to a certain level and stops growing
Set the num_worker_threads=128, memory keeps growing after every forward pass

For the tutorial for the ResNet, seems ResNet has more activations for transmission:

Set the num_worker_threads=2, need to run for more than set timeout (60000 ms)
Set the num_worker_threads=16, need to run for more than set timeout (60000 ms)
Set the num_worker_threads=64, memory keeps growing after every forward pass in one batch and is reduced after finishing one batch, and the next batch goes to a larger value. Eg, increase from ~230M to peak memory ~2090M in the first batch; then is reduced to  ~235M and goes to peak ~ 2250M in the second batch. But the growth trend is slowing down.
Set the num_worker_threads=128, similar to 64. But the peak memory is larger.

And for your suggestions:



 wayi:


I don’t know if a process runs RPC on itself will lead to a memory leak or not. Ideally it should work as doing the work locally, but I am not sure if such an unreasonable corner case is supported well. You can update the code to train ResNetShard1 locally.



I think it is not related to RPC in itself since only the memory for the master process does not keep growing when using mp.spawn(). But I am not sure whether it is a normal phenomenon since the growth speed for memory seems to show slowing down after multiple batches in the original tutorial.  I will increase the batch size for my script.



 wayi:

I think you should still use mp.spawn in the main function. I haven’t tried multi-node training by myself. Does mp.spawn fail on multiple nodes?


I do not use mp.spawn() is due to cannot give a certain rank in mp.spawn() for multiple nodes. Or any suggestions or suitable usage for this API?



 wayi:

Out of curiosity, since I don’t have experience on running the computation on CPU cores via RPC, I wonder if you can run the same code on 2 CPU cores of the same machine, by changing “cuda:0” and “cuda:1” to “cpu:0” and “cpu:1”, respectively. We can check if this will also cause a memory leak as well. If so, the problem could be TensorPipe + CPU.


Yes, I could change “cuda:0” and “cuda:1” to “cpu:0” and “cpu:1”, and the code runs successfully. But it also shows a memory leak problem.
Thanks for your reply and suggestions! Hope to hear more of your thoughts
Best,
YANG