+++++++++++++++++++++++++++++++++++++++++++++++++++++++
                      PAGE 3/5

+++++++++++++++++++++++++++++++++++++++++++++++++++++++
To do that, we need to determine what are the parts of code that create randomness. Once determined, we put them in locks to synchronize the access to them. Such parts need to be the smallest possible so we do not lock the entire code and make it sequential. Remember locks are blocking. So, you have to use them wisely over the smallest region of  instructions.
Code 3: Synchronize threads using locks to reset the PRNG status (thread-safe) like nothing happened.
import threading
import random

import numpy as np

import torch
import torch.nn as nn
from torch.nn import DataParallel


def set_seed(seed):
    """
    For seed to some modules.
    :param seed: int. The seed.
    :return:
    """
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()
        self.thread_lock = threading.Lock()

    def forward(self, x):
        if x.is_cuda:
            device = x.get_device()
            print("DEVICE: {}".format(x.get_device()))
            print("x.size() = {}".format(x.size()))
            print("x:\n {}".format(x))
        else:
            device = torch.device("cpu")

        prngs0 = torch.random.get_rng_state().type(torch.float).numpy()
        # computing sum(abs(diff)**2) is a better way to have a summary of the PRNGs instead of computing only the sum.
        # Two different PRNGs may have the same sum (may happen if we generate few random numbers).
        print("DEVICE {} PRNG STATUS BEFORE: \n {}".format(torch.cuda.current_device(),
                                                           np.abs(np.diff(prngs0)**2).sum()))

        self.thread_lock.acquire()
        set_seed(1)  # re-seed
        delta = torch.rand(x.size()).to(device)  # the danger zone.
        set_seed(1)  # re-seed like nothing happened!!!
        self.thread_lock.release()

        prngs1 = torch.random.get_rng_state().type(torch.float).numpy()
        print("DEVICE {} PRNG STATUS AFTER: \n {}".format(torch.cuda.current_device(),
                                                          np.abs(np.diff(prngs1)**2).sum()))

        print("DEVICE {} DIFF PRNG STATUS ABS(BEFORE - AFTER).SUM(): \n {}".format(
            torch.cuda.current_device(),   np.abs(prngs0 - prngs1).sum()))
        print("Delta:\n {}".format(delta))
        x = x + delta
        return x


if __name__ == "__main__":
    set_seed(0)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = Model()
    model = DataParallel(model)
    model.to(device)

    x = torch.rand(8, 3)
    x = x.to(device)
    print("X in:\n {}".format(x))
    print("X out:\n {}".format(model(x)))


run:
X in:
 tensor([[0.4963, 0.7682, 0.0885],
        [0.1320, 0.3074, 0.6341],
        [0.4901, 0.8964, 0.4556],
        [0.6323, 0.3489, 0.4017],
        [0.0223, 0.1689, 0.2939],
        [0.5185, 0.6977, 0.8000],
        [0.1610, 0.2823, 0.6816],
        [0.9152, 0.3971, 0.8742]], device='cuda:0')
DEVICE: 0
x.size() = torch.Size([4, 3])
DEVICE: 1
x.size() = torch.Size([4, 3])
x:
 tensor([[0.0223, 0.1689, 0.2939],
        [0.5185, 0.6977, 0.8000],
        [0.1610, 0.2823, 0.6816],
        [0.9152, 0.3971, 0.8742]], device='cuda:1')
DEVICE 1 PRNG STATUS BEFORE: 
 46790328.0
DEVICE 1 PRNG STATUS AFTER: 
 47205040.0
DEVICE 1 DIFF PRNG STATUS ABS(BEFORE - AFTER).SUM(): 
 213383.0
Delta:
 tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999],
        [0.3971, 0.7544, 0.5695],
        [0.4388, 0.6387, 0.5247]], device='cuda:1')
x:
 tensor([[0.4963, 0.7682, 0.0885],
        [0.1320, 0.3074, 0.6341],
        [0.4901, 0.8964, 0.4556],
        [0.6323, 0.3489, 0.4017]], device='cuda:0')
DEVICE 0 PRNG STATUS BEFORE: 
 47205040.0
DEVICE 0 PRNG STATUS AFTER: 
 47205040.0
DEVICE 0 DIFF PRNG STATUS ABS(BEFORE - AFTER).SUM(): 
 0.0
Delta:
 tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999],
        [0.3971, 0.7544, 0.5695],
        [0.4388, 0.6387, 0.5247]], device='cuda:0')
X out:
 tensor([[1.2539, 1.0475, 0.4915],
        [0.8667, 0.3367, 1.4339],
        [0.8872, 1.6508, 1.0251],
        [1.0711, 0.9876, 0.9264],
        [0.7800, 0.4482, 0.6970],
        [1.2532, 0.7269, 1.5999],
        [0.5582, 1.0366, 1.2511],
        [1.3540, 1.0358, 1.3988]], device='cuda:0')

The above code is deterministic. Over each GPU, the PRNG generates the same random numbers.
Now, you may say that this is a good thing, and we can produce reproducibility whatever the number of GPUs. Not so fast!!! we will arrive to an issue (see code 5). In short, random number generation is sequential.
Ok, now we arrive to produce the same randomness over each GPU which is fine. But, you may not want that. You may want to have different BUT deterministic randomness on each GPU. To do that, we need to seed each thread differently. The code below provides a way of doing that.