File size: 4,320 Bytes
ec0a9aa | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import logging
import os
import random
import numpy as np
import torch
import torch.distributed as dist
from omegaconf import DictConfig, ListConfig, OmegaConf
def requires_grad(model: torch.nn.Module, flag: bool = True) -> None:
"""
Set requires_grad flag for all parameters in a model.
"""
for p in model.parameters():
p.requires_grad = flag
def set_seed(seed):
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
def str_to_dtype(x: str):
if x == "fp32":
return torch.float32
elif x == "fp16":
return torch.float16
elif x == "bf16":
return torch.bfloat16
else:
raise RuntimeError(f"Only fp32, fp16 and bf16 are supported, but got {x}")
def merge_args(args1, args2):
"""
Merge two argparse Namespace objects.
"""
if args2 is None:
return args1
for k in args2._content.keys():
if k in args1.__dict__:
v = getattr(args2, k)
if isinstance(v, ListConfig) or isinstance(v, DictConfig):
v = OmegaConf.to_object(v)
setattr(args1, k, v)
else:
raise RuntimeError(f"Unknown argument {k}")
return args1
def all_exists(paths):
return all(os.path.exists(path) for path in paths)
def get_logger():
return logging.getLogger(__name__)
def create_logger(logging_dir=None):
"""
Create a logger that writes to a log file and stdout.
"""
if dist.get_rank() == 0:
additional_args = dict()
if logging_dir is not None:
additional_args["handlers"] = [
logging.StreamHandler(),
logging.FileHandler(f"{logging_dir}/log.txt"),
]
logging.basicConfig(
level=logging.INFO,
format="[\033[34m%(asctime)s\033[0m] %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
**additional_args,
)
logger = logging.getLogger(__name__)
else: # dummy logger (does nothing)
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())
return logger
def all_to_all(input_: torch.Tensor, gather_dim: int, scatter_dim: int) -> torch.Tensor:
assert gather_dim != scatter_dim
assert 0 <= gather_dim < input_.ndim
assert 0 <= scatter_dim < input_.ndim
world_size = dist.get_world_size()
assert input_.size(scatter_dim) % world_size == 0
if world_size == 1:
return input_
inputs = [x.contiguous() for x in input_.chunk(world_size, dim=scatter_dim)]
outputs = [torch.empty_like(x) for x in inputs]
dist.all_to_all(outputs, inputs)
return torch.cat(outputs, dim=gather_dim)
def sp_split(input_: torch.Tensor) -> torch.Tensor:
size = dist.get_world_size()
rank = dist.get_rank()
if size == 1:
return input_
assert input_.size(1) % size == 0
return input_.chunk(size, dim=1)[rank].contiguous()
def sp_gather(input_: torch.Tensor) -> torch.Tensor:
size = dist.get_world_size()
rank = dist.get_rank()
if size == 1:
return input_
output = [torch.empty_like(input_) for _ in range(size)]
dist.all_gather(output, input_)
return torch.cat(output, dim=1)
def _setup_dist_env_from_slurm():
import subprocess
from time import sleep
while not os.environ.get("MASTER_ADDR", ""):
try:
os.environ["MASTER_ADDR"] = subprocess.check_output(
"sinfo -Nh -n %s | head -n 1 | awk '{print $1}'" %
os.environ['SLURM_NODELIST'],
shell=True,
).decode().strip()
except:
pass
sleep(1)
os.environ["MASTER_PORT"] = str(18183)
os.environ["RANK"] = os.environ["SLURM_PROCID"]
os.environ["WORLD_SIZE"] = os.environ["SLURM_NPROCS"]
os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"]
os.environ["LOCAL_WORLD_SIZE"] = os.environ["SLURM_NTASKS_PER_NODE"]
def init_process_groups():
if any([
x not in os.environ
for x in ["RANK", "WORLD_SIZE", "MASTER_PORT", "MASTER_ADDR"]
]):
_setup_dist_env_from_slurm()
dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
|