File size: 2,205 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | # Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is licensed under a Creative Commons
# Attribution-NonCommercial-ShareAlike 4.0 International License.
# You should have received a copy of the license along with this
# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
import os
import torch
from datetime import timedelta
from accelerate import Accelerator
from accelerate import DistributedDataParallelKwargs
try:
from accelerate import InitProcessGroupKwargs
except Exception:
InitProcessGroupKwargs = None
#----------------------------------------------------------------------------
ACCELERATOR = None
def init():
global ACCELERATOR
timeout_minutes = int(os.environ.get("REOAC_DIST_TIMEOUT_MIN", "60"))
if timeout_minutes < 1:
timeout_minutes = 60
ddp_kwargs = DistributedDataParallelKwargs(
find_unused_parameters=False,
broadcast_buffers=False,
)
kwargs_handlers = [ddp_kwargs]
if InitProcessGroupKwargs is not None:
kwargs_handlers.append(
InitProcessGroupKwargs(timeout=timedelta(minutes=timeout_minutes))
)
ACCELERATOR = Accelerator(kwargs_handlers=kwargs_handlers)
#----------------------------------------------------------------------------
def get_accelerator():
return ACCELERATOR
def get_rank():
return ACCELERATOR.process_index
#----------------------------------------------------------------------------
def get_local_rank():
return ACCELERATOR.local_process_index
#----------------------------------------------------------------------------
def get_world_size():
return ACCELERATOR.num_processes
#----------------------------------------------------------------------------
def update_progress(cur, total):
_ = cur, total
#----------------------------------------------------------------------------
def print0(*args, **kwargs):
if get_rank() == 0:
print(*args, **kwargs)
#----------------------------------------------------------------------------
if __name__ == "__main__":
init()
a = torch.zeros(3, device="cuda") + get_rank()
aa = ACCELERATOR.gather(a)
print(aa)
|