Memory leak when using RPC for pipeline parallelism

HuYang719: <p>Hi, when I use torch.distributed.rpc to implement pipeline parallelism for transformer-based inference, the memory consumption increases with each forward pass.  The code for 2 nodes is like this,</p>
<p>First, I define two classes for transformer shard</p>
<pre><code class="lang-auto">import os
import sys
import threading
import time
import torch
import torch.nn as nn
import torch.distributed.rpc as rpc
from torch.distributed.rpc import RRef
from transformers import ViTFeatureExtractor, ViTForImageClassification
#####################################
#               Define Transformer Shard             #
####################################


class TransformerShard1(nn.Module):
    def __init__(self, device, config,num_layers):
        super().__init__()
        # self.config = config
        self.model = ViTForImageClassification.from_pretrained(config)
        self.model.vit.encoder.layer = nn.Sequential(*[self.model.vit.encoder.layer[i] for i in range(num_layers)])
        self.model.vit= nn.Sequential(*list(self.model.vit.children())[:-1])
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        self._lock = threading.Lock()
        self.device = device

    def forward_kernel(self, x):
        x = self.model(x).to_tuple()[0]
        end = time.time()
        return x

    @torch.no_grad()
    def forward(self, pixel_values=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        x = pixel_values.to_here().to(self.device)
        with self._lock:
            x = self.forward_kernel(x)
        return x.cpu()

class TransformerShard2(nn.Module):
    def __init__(self,device, config, num_layers):
        super().__init__()
        self.model = ViTForImageClassification.from_pretrained(config)
        self.model.vit.encoder.layer = nn.Sequential(*[self.model.vit.encoder.layer[i] for i in range(num_layers, 2*num_layers)])
        self.model.vit= nn.Sequential(*list(self.model.vit.children())[1:-1])
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        self._lock = threading.Lock()
        self.device = device

    def forward_kernel(self, x):
        x = self.model(x)[0]
        return x

    @torch.no_grad()
    def forward(self, x_rref=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            x = self.forward_kernel(x)
        return x.cpu()
</code></pre>
<p>Then I stitch them into one class for forwarding:</p>
<pre><code class="lang-auto">class DistViT(nn.Module):
    def __init__(
        self, 
        split_size,
        workers,
        config,
        num_layers,
        *args, **kwargs
    ):
        super().__init__()


        self.split_size = split_size   # for microbatch
        self.num_layers = num_layers 
        self.p1_rref = rpc.remote(
            workers[0],
            TransformerShard1,
            args = ("cpu", config, int(num_layers/2)) + args,
            kwargs = kwargs
        )

        self.p2_rref = rpc.remote(
            workers[1],
            TransformerShard2,
            args = ("cpu", config, int(num_layers/2)) + args,
            kwargs = kwargs
        )
   

    def forward(self, pixel_values=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        out_futures = []
        start = time.time()
        id = 0
        for x in iter(pixel_values.split(self.split_size, dim=0)):
            x_rref = RRef(x)
            y_rref = self.p1_rref.remote().forward(x_rref)
            z_fut = self.p2_rref.rpc_async().forward(y_rref) 
            out_futures.append(z_fut)
        torch.futures.wait_all(out_futures)
        return out_futures

    def parameter_rrefs(self):
        remote_params = []
        remote_params.extend(self.p1_rref.remote().parameter_rrefs().to_here())
        remote_params.extend(self.p2_rref.remote().parameter_rrefs().to_here())
        return remote_params
</code></pre>
<p>Then run RPC process like this:</p>
<pre><code class="lang-auto">######################################
#                   Run RPC Processes                   #
######################################
config = 'google/vit-base-patch16-224'
num_layers = 12

num_batches = 1
batch_size = 256

img = torch.randn(3, 384, 384)
imgs = [img for i in range(batch_size)]

feature_extractor = ViTFeatureExtractor.from_pretrained(config)
def run_master(split_size):
    # put the two model parts on worker1 and worker2 respectively
    print("Run mastering \n")
    for si in range(len(split_size)):
        model = DistViT(split_size[si], ["worker0", "worker1"], config, num_layers)
        inputs = feature_extractor(images=imgs, return_tensors="pt")
        for i in range(num_batches):
            # generate random inputs and labels       
            outputs = model(**inputs)


def run_worker(rank, world_size, num_split):
     # run on local host
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29501'

    # Higher timeout is added to accommodate for kernel compilation time in case of ROCm.
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256,rpc_timeout=3000)

    if rank == 0:
        rpc.init_rpc(
            "worker0",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        run_master(num_split)
    else:
        rpc.init_rpc(
            "worker1",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        pass

    rpc.shutdown()

</code></pre>
<p>The main function:</p>
<pre><code class="lang-auto">if __name__=="__main__":
    world_size = 2
    rank=int(sys.argv[1])
    num_split=[8]

    print(f"{config}, {num_layers}, {num_split}")
    
    tik = time.time()
    run_worker(rank, world_size, num_split)
    tok = time.time()
    print(f"Total program execution time = {tok - tik}")
</code></pre>
<p>It needs to import <a href="https://huggingface.co/transformers/model_doc/vit.html?highlight=vit" rel="noopener nofollow ugc">transformers</a> package with the command:</p>
<pre><code class="lang-auto">pip install transformers
</code></pre>
<p>The whole script is uploaded in <a href="https://pastebin.ubuntu.com/p/XVfsjHB5W9/" rel="noopener nofollow ugc">ubuntu pastebin</a></p>
<p>I run it on my macOS with PyTorch 1.8.1 CPU-only, the running command is like this:<br>
on the first terminal:</p>
<pre><code class="lang-auto">python pipeline_parallelism.py 0
</code></pre>
<p>on the second terminal:</p>
<pre><code class="lang-auto">python pipeline_parallelism.py 1
</code></pre>
<p>I use top command to check the memory usage, and after each forward, the memory increases about 3MB until OOM or the RPC shutdown. Could anyone help me to fix or find the problem? Thank you very much!</p>

wayi: <p>This looks interesting!</p>
<p>I think firs you need to limit the debugging scope. For example, if you change the microbatch size to 1, will it still cause memory leak? If so, then the issue at least is not microbatching in pipeline parallelism.</p>
<p>On the other hand, have you tried PyTorch native pipeline parallelism API <a href="https://pytorch.org/docs/stable/pipeline.html" rel="noopener nofollow ugc">here</a>?</p>

HuYang719: <p>Hi Yi Wang,</p>
<p>Thanks for your reply! When I limit the microbatch size to 1 and change the num_batches to 10, the memory leak problem exists. I will continue to debug this, any suggestions?</p>
<p>Thank you for recommending this API But PyTorch native pipeline parallelism API seems not to work for multiple distributed machines, right? Currently we want to build pipeline parallelism for multiple nodes.</p>
<p>Best,<br>
Yang</p>

wayi: <blockquote>
<p>I will continue to debug this, any suggestions?</p>
</blockquote>
<p>Then the bug is unlikely to be relevant to the pipelining logic. One thing I would check is all the RRef objects. Have you used any of the RRef objects or its field temporarily? If any, you can try to delete it after the usage. For example, you probably can try <code>del pixel_values</code> after the line <code>x = pixel_values.to_here().to(self.device)</code>.</p>

HuYang719: <p>Hi Yi Wang,</p>
<p>Thanks for your reply. I edit the code according to your suggestions like this (del pixel_values, del x_rref):</p>
<pre><code class="lang-auto">class TransformerShard1(nn.Module):
    def __init__(self, device, config,num_layers):
        super().__init__()
        # self.config = config
        self.model = ViTForImageClassification.from_pretrained(config)
        self.model.vit.encoder.layer = nn.Sequential(*[self.model.vit.encoder.layer[i] for i in range(num_layers)])
        self.model.vit= nn.Sequential(*list(self.model.vit.children())[:-1])
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        self._lock = threading.Lock()
        self.device = device

    def forward_kernel(self, x):
        x = self.model(x).to_tuple()[0]
        end = time.time()
        return x

    @torch.no_grad()
    def forward(self, pixel_values=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        x = pixel_values.to_here().to(self.device)
        del pixel_values
        gc.collect()
        with self._lock:
            x = self.forward_kernel(x)
        print("Shard1 finish its microbatch ")
        return x.cpu()

class TransformerShard2(nn.Module):
    def __init__(self,device, config, num_layers):
        super().__init__()
        self.model = ViTForImageClassification.from_pretrained(config)
        self.model.vit.encoder.layer = nn.Sequential(*[self.model.vit.encoder.layer[i] for i in range(num_layers, 2*num_layers)])
        self.model.vit= nn.Sequential(*list(self.model.vit.children())[1:-1])
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        self._lock = threading.Lock()
        self.device = device

    def forward_kernel(self, x):
        x = self.model(x)[0]
        return x

    @torch.no_grad()
    def forward(self, x_rref=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        x = x_rref.to_here().to(self.device)
        del x_rref
        gc.collect()
        with self._lock:
            x = self.forward_kernel(x)
        print("Shard1 finish its microbatch ")
        return x.cpu()
</code></pre>
<p>The overall memory usage is reduced (for the microbatch is 256, the peak memory usage is reduced from ~800MB to ~500MB). However, after each forward pass, the memory still keeps growing. It even does not be reduced after finishing one batch, which is very strange.</p>

wayi: <p>I only provided one example, and it’s expected that my suggestions is not exhaustive.</p>
<p>Memory leak is a really tricky issue. You can try reverting some code and checking when the memory leak appears step by step.</p>
<p>On the other hand, check out this question: <a href="https://discuss.pytorch.org/t/how-to-debug-causes-of-gpu-memory-leaks/6741/3" class="inline-onebox">How to debug causes of GPU memory leaks? - #3 by smth</a></p>

HuYang719: <p>Hi Yi Wang,</p>
<p>Thanks for your reply and this useful link! I found this memory leak problem is related to <code>num_worker_threads=256</code> in <code>options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256)</code>. According to your suggestion and <a href="https://discuss.pytorch.org/t/how-to-debug-causes-of-gpu-memory-leaks/6741/3">the link</a>, I first add:</p>
<pre><code class="lang-auto">for obj in gc.get_objects():
            try:
                if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
                    print(f"After forward: {type(obj)}, {obj.size()}")
            except: pass
</code></pre>
<p>to trace the tensor and add <code>del ...</code> to make sure there are the same amount of tensors between different forward passes. But the memory still keeps growing.</p>
<p>However, when I reduce the  <code>num_worker_threads=256</code> in <code>options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256)</code> to  <code>num_worker_threads=2</code>, the memory stops growing. Further, the memory shows significant growing only when I set num_worker_threads &gt; 6.</p>
<p>I could only find little information about this argument in the PyTorch <a href="https://pytorch.org/docs/stable/rpc.html" rel="noopener nofollow ugc">document</a>:</p>
<blockquote>
<ul>
<li><strong>num_worker_threads</strong> (<a href="https://docs.python.org/3/library/functions.html#int" rel="noopener nofollow ugc"><em>int</em></a><em>,</em> <em>optional</em>) – The number of threads in the thread-pool used by <code>TensorPipeAgent</code> to execute requests (default: 16).</li>
</ul>
</blockquote>
<p>And <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html?highlight=distributed%20pipeline" rel="noopener nofollow ugc">the tutorial for pipeline parallelism</a> set it to 128, which I thought was related to the communication/request speed.  I am not familiar with it, and wonder:</p>
<ol>
<li>Why does larger num_worker_threads cause the memory leak problem?</li>
<li>When do we need a larger num_worker_threads?</li>
</ol>
<p>Do you have any related experience or suggestions? I need your help! Thank you very much!</p>
<p>Best,<br>
Yang</p>

wayi: <blockquote>
<ol>
<li>Why does larger num_worker_threads cause the memory leak problem?</li>
<li>When do we need a larger num_worker_threads?</li>
</ol>
</blockquote>
<p>Good question!  cc: <a class="mention" href="/u/lcw">@lcw</a> (TensorPipe expert).</p>
<p>On the other hand, can you try <code>ProcessGroupRpcBackendOptions</code> to see if there is memory leak? This backend will be deprecated soon, as it is much slower on GPUs, but you should still be able to use it now.</p>
<p>Example code:</p>
<pre><code class="lang-auto">rpc.init_rpc(
    "worker1",
     rank=0,
     world_size=2,
     backend=rpc.BackendType.PROCESS_GROUP,
     rpc_backend_options=rpc.ProcessGroupRpcBackendOptions(
         num_send_recv_threads=16,
         rpc_timeout=20 # 20 second timeout
     )
)
</code></pre>

HuYang719: <p>Hi Yi Wang,</p>
<p>When I change to <code>ProcessGroupRpcBackendOptions</code> and <code>set num_send_recv_threads=16</code>, the memory also keeps growing with each forward pass. I observe it likes this: the memory first is reduced ~ 1M after one forward pass and then be increased &gt;3M after several passes, which is similar to the <code>TensorPipe</code> backend.</p>
<p>In addition, <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html?highlight=distributed%20pipeline" rel="noopener nofollow ugc">the tutorial</a> with <code>mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True)</code> and <code> options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256, rpc_timeout=300)</code> do not cause memory leak problem on single machine. However, if I change it for 2 nodes like the following, it shows memory leak problem (the modified lines are marked with comment <code># original:</code>):</p>
<pre><code class="lang-auto">import os
import sys
import threading
import time
from functools import wraps

import torch
import torch.nn as nn
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
import torch.optim as optim
from torch.distributed.optim import DistributedOptimizer
from torch.distributed.rpc import RRef

from torchvision.models.resnet import Bottleneck


#########################################################
#           Define Model Parallel ResNet50              #
#########################################################

# In order to split the ResNet50 and place it on two different workers, we
# implement it in two model shards. The ResNetBase class defines common
# attributes and methods shared by two shards. ResNetShard1 and ResNetShard2
# contain two partitions of the model layers respectively.


num_classes = 1000


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class ResNetBase(nn.Module):
    def __init__(self, block, inplanes, num_classes=1000,
                 groups=1, width_per_group=64, norm_layer=None):
        super(ResNetBase, self).__init__()

        self._lock = threading.Lock()
        self._block = block
        self._norm_layer = nn.BatchNorm2d
        self.inplanes = inplanes
        self.dilation = 1
        self.groups = groups
        self.base_width = width_per_group

    def _make_layer(self, planes, blocks, stride=1):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if stride != 1 or self.inplanes != planes * self._block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * self._block.expansion, stride),
                norm_layer(planes * self._block.expansion),
            )

        layers = []
        layers.append(self._block(self.inplanes, planes, stride, downsample, self.groups,
                                  self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * self._block.expansion
        for _ in range(1, blocks):
            layers.append(self._block(self.inplanes, planes, groups=self.groups,
                                      base_width=self.base_width, dilation=self.dilation,
                                      norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def parameter_rrefs(self):
        r"""
        Create one RRef for each parameter in the given local module, and return a
        list of RRefs.
        """
        return [RRef(p) for p in self.parameters()]


class ResNetShard1(ResNetBase):
    """
    The first part of ResNet.
    """
    def __init__(self, device, *args, **kwargs):
        super(ResNetShard1, self).__init__(
            Bottleneck, 64, num_classes=num_classes, *args, **kwargs)

        self.device = device
        self.seq = nn.Sequential(
            nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False),
            self._norm_layer(self.inplanes),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            self._make_layer(64, 3),
            self._make_layer(128, 4, stride=2)
        ).to(self.device)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            out =  self.seq(x)
            print("Shard1 finish 1 microbatch")
        return out.cpu()


class ResNetShard2(ResNetBase):
    """
    The second part of ResNet.
    """
    def __init__(self, device, *args, **kwargs):
        super(ResNetShard2, self).__init__(
            Bottleneck, 512, num_classes=num_classes, *args, **kwargs)

        self.device = device
        self.seq = nn.Sequential(
            self._make_layer(256, 6, stride=2),
            self._make_layer(512, 3, stride=2),
            nn.AdaptiveAvgPool2d((1, 1)),
        ).to(self.device)

        self.fc =  nn.Linear(512 * self._block.expansion, num_classes).to(self.device)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            out = self.fc(torch.flatten(self.seq(x), 1))
            print("Shard2 finish 1 microbatch")
        return out.cpu()


class DistResNet50(nn.Module):
    """
    Assemble two parts as an nn.Module and define pipelining logic
    """
    def __init__(self, split_size, workers, *args, **kwargs):
        super(DistResNet50, self).__init__()

        self.split_size = split_size

        # Put the first part of the ResNet50 on workers[0]
        self.p1_rref = rpc.remote(
            workers[0],
            ResNetShard1,
            # original: args = ("cuda:0",) + args,
            args = ("cpu",) + args,
            kwargs = kwargs
        )

        # Put the second part of the ResNet50 on workers[1]
        self.p2_rref = rpc.remote(
            workers[1],
            ResNetShard2,
            # original: args = ("cuda:1",) + args,
            args = ("cpu",) + args,
            kwargs = kwargs
        )

    def forward(self, xs):
        # Split the input batch xs into micro-batches, and collect async RPC
        # futures into a list
        out_futures = []
        for x in iter(xs.split(self.split_size, dim=0)):
            x_rref = RRef(x)
            y_rref = self.p1_rref.remote().forward(x_rref)
            z_fut = self.p2_rref.rpc_async().forward(y_rref)
            out_futures.append(z_fut)

        # collect and cat all output tensors into one tensor.
        return torch.cat(torch.futures.wait_all(out_futures))

    def parameter_rrefs(self):
        remote_params = []
        remote_params.extend(self.p1_rref.remote().parameter_rrefs().to_here())
        remote_params.extend(self.p2_rref.remote().parameter_rrefs().to_here())
        return remote_params


#########################################################
#                   Run RPC Processes                   #
#########################################################

#original: num_batches = 3
num_batches = 10
batch_size = 120
image_w = 128
image_h = 128


def run_master(split_size):

    # put the two model parts on worker1 and worker2 respectively
    model = DistResNet50(split_size, ["worker0", "worker1"])
    loss_fn = nn.MSELoss()
    opt = DistributedOptimizer(
        optim.SGD,
        model.parameter_rrefs(),
        lr=0.05,
    )

    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)

    for i in range(num_batches):
        print(f"Processing batch {i}")
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)

        # The distributed autograd context is the dedicated scope for the
        # distributed backward pass to store gradients, which can later be
        # retrieved using the context_id by the distributed optimizer.
        with dist_autograd.context() as context_id:
            outputs = model(inputs)
            dist_autograd.backward(context_id, [loss_fn(outputs, labels)])
            opt.step(context_id)


def run_worker(rank, world_size, num_split):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'


    # Higher timeout is added to accommodate for kernel compilation time in case of ROCm.
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=128, rpc_timeout=300)

    if rank == 0:
        rpc.init_rpc(
             # original "master"
            "worker0",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        run_master(num_split)
    else:
        rpc.init_rpc(
            f"worker{rank}",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        pass

    # block until all rpcs finish
    rpc.shutdown()


if __name__=="__main__":
    # original: world_size = 3
    world_size = 2
    rank=int(sys.argv[1])
    # original:  for num_split in [1, 2, 4, 8]:
    for num_split in [8]:
        # original: mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True)
        run_worker(rank, world_size, num_split)
</code></pre>
<p>Do I have the wrong usage for RPC for multiple nodes? Or do you know any referencess/examples for multiple nodes? Thanks for your time and your suggestion!</p>
<p>Best,<br>
Yang</p>

wayi: <p>As I understand, originally the example uses 3 processes: 1 master, and worker1 + worker2, where both workers are assigned to two CUDA devices. Now you only use 2 processes: worker0 (also as master) and worker1, both are assigned to CPU, so basically worker0 runs RPC on itself.</p>
<p>Here are my suggestions:</p>
<ol>
<li>I don’t know if a process runs RPC on itself will lead to 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 <code>ResNetShard1</code> locally.</li>
<li>I think you should still use <code>mp.spawn</code> in the main function. I haven’t tried multi-node training by myself. Does <code>mp.spawn</code> fail on multiple nodes?</li>
<li>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 memory leak as well. If so, the problem could be TensorPipe + CPU.</li>
</ol>

HuYang719: <p>Hi Yi Wang,</p>
<p>Thanks for your good suggestions. First, I need to <strong>correct</strong> a mistake in my observation:<br>
use <code>mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True)</code> also shows memory leak problem locally, same with use <code>run_worker(rank, world_size, num_split)</code>. I only observe the master process, which does not shows growth. But the worker part shows growth, which has the same result with using <code>run_worker(rank, world_size, num_split)</code> and also could be observed in <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html?highlight=distributed%20pipeline" rel="noopener nofollow ugc">the tutorial for the ResNet</a>.</p>
<p>In my script for transformer models, both <code>mp.spawn()</code> method and <code>run_worker()</code> have similar result :</p>
<ul>
<li>Set  the <code>num_worker_threads=2</code>, memory does not increase</li>
<li>Set the  <code>num_worker_threads=16</code>, memory first increases to a certain level and stops growing</li>
<li>Set the <code>num_worker_threads=128</code>, memory keeps growing after every forward pass</li>
</ul>
<p>For the tutorial for the ResNet, seems ResNet has more activations for transmission:</p>
<ul>
<li>Set the <code>num_worker_threads=2</code>, need to run for more than set timeout (60000 ms)</li>
<li>Set the <code>num_worker_threads=16</code>, need to run for more than set timeout (60000 ms)</li>
<li>Set the <code>num_worker_threads=64</code>, 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.</li>
<li>Set the <code>num_worker_threads=128</code>, similar to 64. But the peak memory is larger.</li>
</ul>
<p>And for your suggestions:</p>
<aside class="quote no-group quote-modified" data-username="wayi" data-post="10" data-topic="127236">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<ul>
<li>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 <code>ResNetShard1</code> locally.</li>
</ul>
</blockquote>
</aside>
<p>I think it is not related to RPC in itself since only the memory for the master process does not keep growing when using <code>mp.spawn()</code>. 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.</p>
<aside class="quote no-group" data-username="wayi" data-post="10" data-topic="127236">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<p>I think you should still use <code>mp.spawn</code> in the main function. I haven’t tried multi-node training by myself. Does <code>mp.spawn</code> fail on multiple nodes?</p>
</blockquote>
</aside>
<p>I do not use <code>mp.spawn()</code> is due to cannot give a certain rank in <code>mp.spawn()</code> for multiple nodes. Or any suggestions or suitable usage for this API?</p>
<aside class="quote no-group quote-modified" data-username="wayi" data-post="10" data-topic="127236">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<p>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.</p>
</blockquote>
</aside>
<p>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.</p>
<p>Thanks for your reply and suggestions! Hope to hear more of your thoughts</p>
<p>Best,<br>
YANG</p>

wayi: <aside class="quote no-group" data-username="HuYang719" data-post="11" data-topic="127236">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/huyang719/48/40462_2.png" class="avatar"> HuYang719:</div>
<blockquote>
<p>For the tutorial for the ResNet, seems ResNet has more activations for transmission:</p>
</blockquote>
</aside>
<p>Thanks for all the detailed comments! Just want to confirm my understanding:</p>
<p>Are you saying that, even in the <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html?highlight=distributed%20pipeline" rel="noopener nofollow ugc">tutorial for the ResNet</a> instead of your own code, you can still observe the memory leak on devices cuda:0 and cuda:1, just not on the master process.</p>
<p>So far, such memory leak is irrelevant to 1) RPC backend (TensorPipe or ProcessGroup), 2) worker’s device (GPU or CPU), or 3) how we launch distributed training (<code>mp.spawn</code> or just <code>run_worker</code>, even within a single host).</p>
<p>The only factor that is relevant to the memory leak is <code>num_worker_threads</code>:</p>
<ul>
<li>Set the <code>num_worker_threads=2</code>, memory does not increase</li>
<li>Set the <code>num_worker_threads=16</code>, memory first increases to a certain level and stops growing</li>
<li>Set the <code>num_worker_threads=128</code>, memory keeps growing after every forward pass</li>
</ul>
<p><a class="mention" href="/u/mrshenli">@mrshenli</a></p>

HuYang719: <p>Hi Yi Wang,</p>
<p>Yes, I agree with all of you said except I do not test it on GPU (cuda:0, cuda:1) since without enough GPU devices.</p>

wayi: <blockquote>
<p>Yes, I agree with all of you said except I do not test it on GPU (cuda:0, cuda:1) since without enough GPU devices.</p>
</blockquote>
<p>So basically you mean the memory leak is reproducible as long as “cuda” is replaced by “cpu” in the <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html?highlight=distributed%20pipeline" rel="noopener nofollow ugc">tutorial for the ResNet</a>.</p>
<p>If so, can you file a bug <a href="https://github.com/pytorch/pytorch/issues/new?assignees=&amp;labels=&amp;template=bug-report.md" rel="noopener nofollow ugc">here</a>. I will let RPC developers aware of this bug. Thanks!</p>

HuYang719: <p>Hi Yi Wang,</p>
<p>I have opened an <a href="https://github.com/pytorch/pytorch/issues/61920" rel="noopener nofollow ugc">issue</a> 3 days ago but current no response. That will be very helpful if you could let them know this! Thank you very much!</p>
<p>Best,<br>
Yang</p>

wayi: <p>I updated that bug thread with a summary of the debugging efforts we have so far. If no one answers that bug thread, I will bring it up in a meeting next week. Thanks!</p>

HuYang719: <p>Thanks for your time and effort! I will test it one more time on GPU tomorrow in case of any other problems before I close this thread. If no more other problems, I will close it tomorrow. Thanks!</p>

UnbearableFate: <p>I encountered a similar issue with the current version, Torch 2.6 + CUDA 12.4. I am running experiments on multiple nodes using mpirun with the <code>--map-by ppr:1:node</code> setting. Even after reducing <code>num_worker_threads</code> to 6, the issue remains unresolved. Are there any other possible solutions?</p>