Tensor parallel spawns additional processes on GPU0 and uses additional memory

username: gbarello
cooked: <p>Hello, I am toying with tensor parallelism - it seems to be working great for the most part, but I noticed that it spawns additional processes on GPU0 which can take up a substantial amount of memory.</p>
<p>This seems to happen after model parallelization, before starting the training loop (though the memory overhead of those additional processes grows once data is fed through the model and gradients are computed).</p>
<p>Output of <code>nvidia-smi</code> after model instantiation but before parallelization:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/b/4b3d37a42fb22de33baf68ca201bbb55e31a6189.png" data-download-href="https://discuss.pytorch.org/uploads/default/4b3d37a42fb22de33baf68ca201bbb55e31a6189" title="Screenshot 2024-03-31 at 7.27.17 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/b/4b3d37a42fb22de33baf68ca201bbb55e31a6189_2_690x170.png" alt="Screenshot 2024-03-31 at 7.27.17 AM" data-base62-sha1="aJAY0pxojjBM25nNB9wH8KxtxBD" width="690" height="170" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/b/4b3d37a42fb22de33baf68ca201bbb55e31a6189_2_690x170.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/b/4b3d37a42fb22de33baf68ca201bbb55e31a6189_2_1035x255.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/b/4b3d37a42fb22de33baf68ca201bbb55e31a6189.png 2x" data-dominant-color="333333"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-03-31 at 7.27.17 AM</span><span class="informations">1114×276 62.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Output of <code>nvidia-smi</code> after after calls to <code>parallelize_module</code>:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/7/a7f29e3008399dfbd5650180ed7809b94a034426.png" data-download-href="https://discuss.pytorch.org/uploads/default/a7f29e3008399dfbd5650180ed7809b94a034426" title="Screenshot 2024-03-31 at 7.26.41 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a7f29e3008399dfbd5650180ed7809b94a034426_2_690x222.png" alt="Screenshot 2024-03-31 at 7.26.41 AM" data-base62-sha1="nXJxJnfaZbNBIUJdVEcLqnGLc6W" width="690" height="222" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a7f29e3008399dfbd5650180ed7809b94a034426_2_690x222.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a7f29e3008399dfbd5650180ed7809b94a034426_2_1035x333.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/7/a7f29e3008399dfbd5650180ed7809b94a034426.png 2x" data-dominant-color="363635"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-03-31 at 7.26.41 AM</span><span class="informations">1112×358 116 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Output of <code>nvidia-smi</code> during training loop:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/7/378f12923498fb287970d681b3751d79e91c1c80.png" data-download-href="https://discuss.pytorch.org/uploads/default/378f12923498fb287970d681b3751d79e91c1c80" title="Screenshot 2024-03-31 at 7.26.48 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/378f12923498fb287970d681b3751d79e91c1c80_2_690x223.png" alt="Screenshot 2024-03-31 at 7.26.48 AM" data-base62-sha1="7VuN7VATnlP8lfYaqeowNdeJ6Cs" width="690" height="223" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/378f12923498fb287970d681b3751d79e91c1c80_2_690x223.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/378f12923498fb287970d681b3751d79e91c1c80_2_1035x334.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/7/378f12923498fb287970d681b3751d79e91c1c80.png 2x" data-dominant-color="353535"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-03-31 at 7.26.48 AM</span><span class="informations">1120×362 116 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Here is a self-contained script that reproduces this:</p>
<pre><code class="lang-auto">import torch
import time
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
)

from torch.distributed import device_mesh, get_world_size, init_process_group, get_rank
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam
from tqdm.auto import tqdm
import numpy as np

class TestModule(torch.nn.Module):
    def __init__(self, Nin, Nmid, Nout):
        super().__init__()

        self.l1 = torch.nn.Linear(Nin, Nmid)
        self.l2 = torch.nn.Linear(Nmid, Nout)

        self.nonlin = torch.nn.Softmax()

    def forward(self, input):

        mid = self.l1(input)
        mid = self.nonlin(mid)

        out = self.l2(mid)
        out = self.nonlin(out)

        return out
    
class TestDataset(Dataset):
    def __init__(self, size_in, size_out, length):
        super().__init__()
        self.length = length
        self.size_in = size_in
        self.size_out = size_out

    def __len__(self):
        return self.length
    
    def __getitem__(self, idx):
        inp = np.float32(np.random.randn(self.size_in))
        out = np.float32(np.random.randn(self.size_out))

        return inp, out

init_process_group()

rank = get_rank()
torch.set_default_device(rank)

in_d = 1024
mid_d = 1024
out_d = 1024

d_len = 1000000

model = TestModule(in_d, mid_d, out_d)

mesh = device_mesh.init_device_mesh("cuda", (get_world_size(),))

parallelize_module(
    model.l1,
    device_mesh=mesh,
    parallelize_plan=ColwiseParallel(),
)
parallelize_module(
    model.l2,
    device_mesh=mesh,
    parallelize_plan=RowwiseParallel(),
)

dataset = TestDataset(in_d, out_d, length=d_len)
dataloader = DataLoader(dataset, batch_size = 128)

optimizer = Adam(model.parameters(), lr = .001)

for I, O in tqdm(dataloader):
    print(I.cuda().device)
    optimizer.zero_grad()
    out = model(I)
    loss = ((out - O)**2).sum()
    loss.backward()
    optimizer.step()
</code></pre>
<p>Run via <code>torchrun --nproc-per-node=4 simple_tp_test.py</code></p>

username: gbarello
cooked: <p>This is I think the same question from here <a href="https://discuss.pytorch.org/t/question-about-tensor-parallel-dtensor-parallelize-module/181739" class="inline-onebox">Question about tensor parallel (DTensor, parallelize_module)</a></p>

username: gbarello
cooked: <p>Solved my own issue, I guess this was because of a well known issue: <a href="https://discuss.pytorch.org/t/torch-not-able-to-utilize-gpu-ram-properly/83245/6" class="inline-onebox">Torch not able to utilize GPU ram properly - #6 by Tyan</a></p>