Logical GPU in PyTorch

amlgeek: <p>Hi!</p>
<p>Is it possible to divide a single physical GPU into several logical like in TF?</p>
<p>I have a node with 4 GPUs that I want to subdivide in 8 logical GPUs in order to train 8 models in parallel.</p>
<p>Is this possible?</p>
<p>4 x NVIDIA GTX1080 Ti with 11GB into 8 x logical GPUs with 5GB each.</p>
<p>Thanks!</p>

ptrblck: <p>I don’t know how TF “divides” the GPU, but you can use <a href="https://pytorch.org/docs/stable/generated/torch.cuda.set_per_process_memory_fraction.html"><code>torch.cuda.set_per_process_memory_fraction</code></a> to share the device memory between different processes.</p>

amlgeek: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> !</p>
<p>I will read a little more about this method!</p>
<p>amlgeek</p>

vfdev-5: <p>Here is a TF code that creates 2 logical GPUs from 1 physical GPU:</p>
<pre><code class="lang-python">gpus = tf.config.list_physical_devices("GPU")
if gpus:
  # Create 2 virtual GPUs with 1GB memory each
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024),
         tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.list_logical_devices("GPU")
    print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
</code></pre>
<p>Source: <a href="https://www.tensorflow.org/recommenders/examples/diststrat_retrieval#set_up_two_virtual_gpus" class="inline-onebox" rel="noopener nofollow ugc">Recommending movies: retrieval with distribution strategy  |  TensorFlow Recommenders</a></p>
<p><strong>EDIT:</strong> However, TF says itself that NCCL is not supported:</p>
<pre><code class="lang-auto">WARNING:tensorflow:NCCL is not supported when using virtual GPUs, fallingback to reduction to one device
</code></pre>
<p>Using “gloo” backend we can create N processes using the same GPU and perform collective ops:</p>
<pre><code class="lang-python"># Run as 
# torchrun --nproc_per_node=4 test_virtual_gpus.py

import os
import time
import torch
import torch.distributed as dist
import torch.distributed.distributed_c10d as c10d

def pprint(msg):
    dist.get_rank()
    # We add sleep to avoid printing clutter
    time.sleep(0.5 * rank)
    print(rank, msg)

if __name__ == "__main__":
    
    # See https://pytorch.org/docs/stable/distributed.html#which-backend-to-use
    dist.init_process_group("gloo")
    
    rank = dist.get_rank()
    ws = dist.get_world_size()
    pprint(f"Hello from process {rank} / {os.environ['LOCAL_RANK']} among {ws} others")
        
    t = torch.tensor(rank, device="cuda")
    pprint(t)
    
    dist.all_reduce(t)
    pprint(f"All reduced output: {t}")

    dist.destroy_process_group()

</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
0 Hello from process 0 / 0 among 4 others
1 Hello from process 1 / 1 among 4 others
2 Hello from process 2 / 2 among 4 others
0 tensor(0, device='cuda:0')
3 Hello from process 3 / 3 among 4 others
1 tensor(1, device='cuda:0')
2 tensor(2, device='cuda:0')
3 tensor(3, device='cuda:0')
0 All reduced output: 6
1 All reduced output: 6
2 All reduced output: 6
3 All reduced output: 6
</code></pre>
<p>PS: Keeping in mind that all NCCL-related wrappers and ops wont work in such configuration.</p>