os.environ[CUDA_VISIBLE_DEVICES] does not work well

haoran-hash: <p>The code is below.</p>
<pre><code class="lang-auto">import torch
from torch import nn
import torch.distributed as dist
import torch.multiprocessing as mp
import os


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.attr1 = nn.Parameter(torch.tensor([1., 2., 3.]))
        self.register_buffer('attr2', torch.tensor([4., 5., 6.]))
        self.attr3 = torch.tensor([7., 8., 9.])
    
    def forward(self, x, rank):
        hd = x * self.attr1
        self.attr2 = self.attr2 / (rank + 1)
        hd = hd * self.attr2
        self.attr3 = self.attr3.to(rank)
        self.attr3 = self.attr3 / (rank + 1)
        y = hd * self.attr3
        y = y.mean()

        return y


def run(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    # torch.cuda.set_device(rank)
    os.environ['CUDA_VISIBLE_DEVICES'] = f'{rank}'

    my_model = MyModel().to(rank)
    my_model = nn.parallel.DistributedDataParallel(my_model, device_ids=[rank], output_device=rank)
    optimizer = torch.optim.SGD(my_model.parameters(), lr=0.001, momentum=0.9)
    input = torch.tensor([1., 2., 3.]) * (rank + 1)

    optimizer.zero_grad()
    output = my_model(input, rank)
    output.backward()
    if rank == 0:
        print(my_model.module.attr1.grad)
    optimizer.step()

    if rank == 0:
        print(my_model.module.attr1)
        print(my_model.module.attr2)
        print(my_model.module.attr3)


if __name__ == '__main__':
    world_size = 2
    mp.spawn(run, args=(world_size, ), nprocs=2)

    print('执行完毕')
</code></pre>
<p>Initially, I write this code in order to see the synchronization mechanism of parameter and buffer in multi GPU training.<br>
Finally, I find <code>torch.cuda.set_device(rank)</code> work well, but <code>os.environ['CUDA_VISIBLE_DEVICES']</code> not work well. The latter will report an error.<br>
The error information is below.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/9/8930621f9687136f77ce37817a552c4e6306452e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/8930621f9687136f77ce37817a552c4e6306452e" title="屏幕截图 2021-09-21 174304"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/8930621f9687136f77ce37817a552c4e6306452e_2_690x268.jpeg" alt="屏幕截图 2021-09-21 174304" data-base62-sha1="jzD4B9CDbchbb1GbmgF17uaaxAy" width="690" height="268" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/8930621f9687136f77ce37817a552c4e6306452e_2_690x268.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/8930621f9687136f77ce37817a552c4e6306452e_2_1035x402.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/8930621f9687136f77ce37817a552c4e6306452e_2_1380x536.jpeg 2x" data-dominant-color="111112"><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">屏幕截图 2021-09-21 174304</span><span class="informations">1453×565 224 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>Hope someone can tell me why.</p>

JuanFMontesinos: <p>you have to set it before calling the python code.<br>
It’s not pytorch’s but nvidia’s behaviour.<br>
Devices are assigned to the process before starting python therefore it doesn’t work once u are in.</p>

haoran-hash: <p>so, where should <code>os.environ['CUDA_VISIBLE_DEVICES']</code> write?<br>
above <code>import torch</code>?</p>

JuanFMontesinos: <p>it shouldn’t be inside the python script but to be set as an enviroment variable in the console such as<br>
<code>CUDA_VISIBLE_DEVICES=0,1 python your_script.py</code></p>
<p>Note that you SHOUDLN’T t set it as a permanent enviroment variable in the bashrc as it affects the whole system.</p>

haoran-hash: <p>This way I only set the GPU devices to be used for all processes, not each process.<br>
But <code>torch.cuda.set_device()</code> can set GPU device for each process.</p>

JuanFMontesinos: <p>You can manage internally (via torch commands) which gpu to use at any time.<br>
Most of the data parallel funcs allows to set that and you can set the devices manually anyway <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Just mentioning that defining cuda_visible_devices inside python won’t work no matter what u do.</p>

haoran-hash: <p>So, <code>os.environ['CUDA_VISIBLE_DEVICES]</code> and <code>torch.cuda.set_device()</code> are not conflict.<br>
Use <code>CUDA_VISIBLE_DEVICES=0,1 python your_script.py</code> to set all available GPU devices for all processes. In each process, we can also use <code>torch.cuda.set_device()</code> to specify the GPU device for this process.<br>
<strong>Is this the correct understanding?</strong> <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=10" title=":thinking:" class="emoji" alt=":thinking:"></p>

JuanFMontesinos: <p>Use <code>CUDA_VISIBLE_DEVICES=0,1 python your_script.py</code> to set all available GPU devices for all processes.<br>
I’m not aware of the intrinsecs of torch.cuda.set_device.</p>
<p>Just to mention when you pass <code>device_ids</code> this is a list which enlist the available gpus from the pytorch pov.</p>
<p>For example, if you call<br>
<code>CUDA_VISIBLE_DEVICES=5,7,9</code> there will be 3 gpus from 0 to 2.<br>
so you can pass device_ids=[0,1,2]</p>

haoran-hash: <p>I got it. Thank you.</p>