pytorch中的nn.Module.register_buffer及DDP模块的一些小知识

作用：给一个nn的Module添加一个buffer，而且不会被认为是一个model的parameter。BN的running_mean就是一个buffer。这个buffer会被para一起，存储到模型里面(除非显示的使用persistent=False)。是否persistant的区别就是会不会出现在module的state_dict。

例子：

self.register_buffer('running_mean', torch.zeros(num_features))

如果要采用register_buffer的话，要注意DDP的工作原理，因为DDP模块每一次forward都会把root process的buffer，广播到其他所有的process。

首先介绍一些关于分布式编程的知识：

与一个机器处理相比，分布式系统就是要把数据分散到不同的机器上去进行处理。随之而来的就是单个节点之间要进行信息的交换。在常见的nccl的backend之下，broadcast，all_reduce, reduce, all_gather, gather(注意这个不可以)。reduce_scattter。前面加不加all的含义就是，所有的device会不会都得到最后的结果。例如all_reduce可以被拿来进行求梯度平均的操作。

下面要介绍一下什么是hook。

这个问题一直围绕我很久了，现在终于要解决这个问题。

hook提供了一种机制，程序提供了hook接口之后。用户可以编写一个hook函数，并且将这个hook函数挂到hook接口上。

def register_hook(self, hook):
        r"""Registers a backward hook.
        The hook will be called every time a gradient with respect to the
        Tensor is computed. The hook should have the following signature::
            hook(grad) -> Tensor or None
        The hook should not modify its argument, but it can optionally return
        a new gradient which will be used in place of :attr:`grad`.
        This function returns a handle with a method ``handle.remove()``
        that removes the hook from the module.
        Example::
            >>> v = torch.tensor([0., 0., 0.], requires_grad=True)
            >>> h = v.register_hook(lambda grad: grad * 2)  # double the gradient
            >>> v.backward(torch.tensor([1., 2., 3.]))
            >>> v.grad
             2
             4
             6
            [torch.FloatTensor of size (3,)]
            >>> h.remove()  # removes the hook

下面介绍一些，pytorch中的一些分布式的一些函数。

torch.distributed.all_reduce，将所有机器上的数据进行reduce，结果就是他们所有tensor的结果都相同。

这里还涉及到一些Op，对all_reduce来讲，其实就是SUM，PRODUCT，MIN，MAX，BAND，BOR，BXOR。

# We have 2 process groups, 2 ranks, data type myst be the same.
tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
# tensor([1, 2]) # Rank 0
# tensor([3, 4]) # Rank 1
dist.all_reduce(tensor, op=ReduceOp.SUM)
# tensor([4, 6]) # Rank 0
# tensor([4, 6]) # Rank 1

torch.distributed.reduce，这个和上面那一个函数其实是很像的。但是多了一个argument，dst。表示的是destination rank(目标进程)。

torch.distributed.all_gather，将这个进程组的tensor全部搞到一起，组成一个list。

tensor_list = [torch.zero(2, dtype=torch.int64) for _ in range(2)]
# result [tensor([0, 0]), tensor([0, 0])] # Rank 0 and 1
tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
# tensor([1, 2]) # Rank 0
# tensor([3, 4]) # Rank 1
dist.all_gather(tensor_list, tensor)
# [tensor([1, 2]), tensor([3, 4])] # Rank 0
# [tensor([1, 2]), tensor([3, 4])] # Rank 1

在多机，每一个机器多卡的情况下，有一些概念和practice的代码还是值得多研究研究的。

一图胜千言:

N是nodes个数，机器个数。G是每一个机器上，gpu的个数。W 是所有node上跑的所有进程总数(也叫World Size)。L是指每台机器上跑的process数量。a local rank in [0, L-1] and a global rank in [0, W-1]。

在DDP启动之外，每一个进程都需要知道它的global rank和local rank。只要有了这个信息。所有的process就可以创建一个ProcessGroup，这之后所有的process都能参与到collective communication上面来了(比如AllReduce)。

其中一个方法是使用launch.py。那么它的左右有哪些呢？首先，它会通过环境变量的方式，将world size，global rank， master address，master port，通过命令行参数的方式将local_rank传递给每一个instance。要使用它，你就需要满足俩条件

(a). 它必须为每一个worker提供entry-point。比如不应该使用torch.multiprocessing.spawn来launch子进程。（在下文中就是spmd这个函数）。

(b). 它必须使用环境变量来初始化每一个进程组。

接下来是一个示范。

--local_rank，被launch.py传递的一个命令行参数

--local_world_size 命令行参数，常常是1或者是每一个node上的gpu个数。这个是由我们给定的。

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument("--local_world_size", type=int, default=1)
    args = parser.parse_args()
    spmd_main(args.local_world_size, args.local_rank)

下面是spmd的一个实例。这个process_group就是用backend来初始化一下。其他的东西都是使用了launch.py设置的一些环境变量。

def spmd_main(local_world_size, local_rank):
    # These are the parameters used to initialize the process group
    env_dict = {
        key: os.environ[key]
        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")
    }
    print(f"[{os.getpid()}] Initializing process group with: {env_dict}")
    dist.init_process_group(backend="nccl")
    print(
        f"[{os.getpid()}] world_size = {dist.get_world_size()}, "
        + f"rank = {dist.get_rank()}, backend={dist.get_backend()}"
    )

    demo_basic(local_world_size, local_rank)

    # Tear down the process group
    dist.destroy_process_group()

def demo_basic(local_world_size, local_rank):

    # setup devices for this process. For local_world_size = 2, num_gpus = 8,
    # rank 1 uses GPUs [0, 1, 2, 3] and
    # rank 2 uses GPUs [4, 5, 6, 7].
    n = torch.cuda.device_count() // local_world_size
    device_ids = list(range(local_rank * n, (local_rank + 1) * n))

    print(
        f"[{os.getpid()}] rank = {dist.get_rank()}, "
        + f"world_size = {dist.get_world_size()}, n = {n}, device_ids = {device_ids}"
    )

    model = ToyModel().cuda(device_ids[0])
    ddp_model = DDP(model, device_ids)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_ids[0])
    loss_fn(outputs, labels).backward()
    optimizer.step()

介绍到这里基本就结束了。

为了简单的使用单机，多卡，这里贴出一些脚本。

首先，你应该启动n个process，确保每一个process都能独立的work on a single gpu。你可以使用CUDA_VISIBLE_DEVICES，或者使用torch.cuda.set_device(i)。i 从0-N-1。在每一个进程之中，你应该使用如下的code来建立这个module。

torch.distributed.init_process_group(
	backend='nccl', world_size=N, init_method='...'
)
model = DistributedDataParallel(model, device_ids=[i], output_device=i)

现在给出在inference的时候，应该使用的data sampler。它的作用是将inference的输入，拆成连续的段。有利于我们之后的操作。

class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):
    """
    Distributed Sampler that subsamples indicies sequentially,
    making it easier to collate all results at the end.
    Even though we only use this sampler for eval and predict (no training),
    which means that the model params won't have to be synced (i.e. will not hang
    for synchronization even if varied number of forward passes), we still add extra
    samples to the sampler to make it evenly divisible (like in `DistributedSampler`)
    to make it easy to `gather` or `reduce` resulting tensors at the end of the loop.
    """
 
    def __init__(self, dataset, batch_size, rank=None, num_replicas=None):
        if num_replicas is None:
            if not torch.distributed.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = torch.distributed.get_world_size()
        if rank is None:
            if not torch.distributed.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = torch.distributed.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.batch_size = batch_size
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
        self.total_size = self.num_samples * self.num_replicas
 
    def __iter__(self):
        indices = list(range(len(self.dataset)))
        # add extra samples to make it evenly divisible
        indices += [indices[-1]] * (self.total_size - len(indices))
        # subsample
        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
        return iter(indices)
 
    def __len__(self):
        return self.num_samples

现在直接给出将所有卡的数据直接all_gather到一起的代码。使用的前提是，我们必须使用相同的

def distrbited_concat(tensor, num_examples):
    output_tensors = [tensor.clone() for _ in range(num_examples)]
    torch.distributed.all_gather(output_tensors, tensor)
    concat = torch.cat(output_tensors, dim = 0)
    return concat

现在要解决的问题是进程之间的同步问题。下面将详细的讨论一下进程之间的同步问题，以免出现不是很显然的bug。在torch.distributed这个package中，提供了一个barrier的接口，其作用就是等所有的进程都跑到了barrier之后，他们才会接着往下面进行。

(a). 只在某个进程中执行，不需要同步。
```
if torch.get_rank() == 0:
    something_to_implement
```
(b). 简单同步，那么只需要一行我们就可以完成。
```
torch.distributed.barrier()
```

reference

https://blog.csdn.net/weixin_34237703/article/details/112493632

https://blog.csdn.net/weixin_39959298/article/details/111141826#t8

https://github.com/pytorch/examples/blob/master/distributed/ddp/README.md

https://blog.csdn.net/zandaoguang/article/details/111714014

原文地址：https://www.cnblogs.com/JohnRan/p/15098245.html

框架中的DDP和buffer

reference