scrapegoat/Neural-Audio-Codec

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7"""Torch distributed utilities."""8 9import typing as tp10 11import torch12 13 14def rank():15    if torch.distributed.is_initialized():16        return torch.distributed.get_rank()17    else:18        return 019 20 21def world_size():22    if torch.distributed.is_initialized():23        return torch.distributed.get_world_size()24    else:25        return 126 27 28def is_distributed():29    return world_size() > 130 31 32def all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM):33    if is_distributed():34        return torch.distributed.all_reduce(tensor, op)35 36 37def _is_complex_or_float(tensor):38    return torch.is_floating_point(tensor) or torch.is_complex(tensor)39 40 41def _check_number_of_params(params: tp.List[torch.Tensor]):42    # utility function to check that the number of params in all workers is the same,43    # and thus avoid a deadlock with distributed all reduce.44    if not is_distributed() or not params:45        return46    #print('params[0].device ', params[0].device)47    tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long)48    all_reduce(tensor)49    if tensor.item() != len(params) * world_size():50        # If not all the workers have the same number, for at least one of them,51        # this inequality will be verified.52        raise RuntimeError(f"Mismatch in number of params: ours is {len(params)}, "53                           "at least one worker has a different one.")54 55 56def broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0):57    """Broadcast the tensors from the given parameters to all workers.58    This can be used to ensure that all workers have the same model to start with.59    """60    if not is_distributed():61        return62    tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)]63    _check_number_of_params(tensors)64    handles = []65    for tensor in tensors:66        handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True)67        handles.append(handle)68    for handle in handles:69        handle.wait()70 71 72def sync_buffer(buffers, average=True):73    """74    Sync grad for buffers. If average is False, broadcast instead of averaging.75    """76    if not is_distributed():77        return78    handles = []79    for buffer in buffers:80        if torch.is_floating_point(buffer.data):81            if average:82                handle = torch.distributed.all_reduce(83                    buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True)84            else:85                handle = torch.distributed.broadcast(86                    buffer.data, src=0, async_op=True)87            handles.append((buffer, handle))88    for buffer, handle in handles:89        handle.wait()90        if average:91            buffer.data /= world_size92 93 94def sync_grad(params):95    """96    Simpler alternative to DistributedDataParallel, that doesn't rely97    on any black magic. For simple models it can also be as fast.98    Just call this on your model parameters after the call to backward!99    """100    if not is_distributed():101        return102    handles = []103    for p in params:104        if p.grad is not None:105            handle = torch.distributed.all_reduce(106                p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True)107            handles.append((p, handle))108    for p, handle in handles:109        handle.wait()110        p.grad.data /= world_size()111 112 113def average_metrics(metrics: tp.Dict[str, float], count=1.):114    """Average a dictionary of metrics across all workers, using the optional115    `count` as unormalized weight.116    """117    if not is_distributed():118        return metrics119    keys, values = zip(*metrics.items())120    device = 'cuda' if torch.cuda.is_available() else 'cpu'121    tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32)122    tensor *= count123    all_reduce(tensor)124    averaged = (tensor[:-1] / tensor[-1]).cpu().tolist()125    return dict(zip(keys, averaged))126