medical
AleksanderObuchowski's picture
Add files using upload-large-folder tool
5ceacbc verified
import logging
import os
import pickle
import requests
import tenacity
import time
import shutil
import torch
import torch.distributed as dist
from PIL import Image
from torchvision.utils import make_grid
from fvcore.nn import FlopCountAnalysis
from fvcore.nn import flop_count_table
from fvcore.nn import flop_count_str
logger = logging.getLogger(__name__)
NORM_MODULES = [
torch.nn.BatchNorm1d,
torch.nn.BatchNorm2d,
torch.nn.BatchNorm3d,
torch.nn.SyncBatchNorm,
# NaiveSyncBatchNorm inherits from BatchNorm2d
torch.nn.GroupNorm,
torch.nn.InstanceNorm1d,
torch.nn.InstanceNorm2d,
torch.nn.InstanceNorm3d,
torch.nn.LayerNorm,
torch.nn.LocalResponseNorm,
]
def register_norm_module(cls):
NORM_MODULES.append(cls)
return cls
def is_main_process():
rank = 0
if 'OMPI_COMM_WORLD_SIZE' in os.environ:
rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
return rank == 0
@torch.no_grad()
def analysis_model(model, dump_input, verbose=False):
model.eval()
flops = FlopCountAnalysis(model, dump_input)
total = flops.total()
model.train()
params_total = sum(p.numel() for p in model.parameters())
params_learned = sum(
p.numel() for p in model.parameters() if p.requires_grad
)
logger.info(f"flop count table:\n {flop_count_table(flops)}")
if verbose:
logger.info(f"flop count str:\n {flop_count_str(flops)}")
logger.info(f" Total flops: {total / 1000 / 1000:.3f}M,")
logger.info(f" Total params: {params_total / 1000 / 1000:.3f}M,")
logger.info(f" Learned params: {params_learned / 1000 / 1000:.3f}M")
return total, flop_count_table(flops), flop_count_str(flops)
def gather_tensors(tensor):
"""
Performs all_gather operation on the provided tensors.
*** Warning ***: torch.distributed.all_gather has no gradient.
"""
tensors_gather = [
torch.ones_like(tensor)
for _ in range(int(os.environ['WORLD_SIZE']))
]
dist.all_gather(tensors_gather, tensor, async_op=False)
# need to do this to restore propagation of the gradients
tensors_gather[int(os.environ['RANK'])] = tensor
output = torch.cat(tensors_gather, dim=0)
return output
def is_valid_url(url):
try:
from urllib import parse
return parse.urlparse(str(url)).scheme != ''
except Exception:
return False
@tenacity.retry(stop=tenacity.stop_after_attempt(3))
def download_file(url, filepath):
logger.info(f'Downloading from {url} to {filepath.absolute()}.')
with requests.get(url, stream=True, allow_redirects=True, timeout=60) as r:
if r.status_code > 200:
raise RuntimeError(f'Failed in downloading from {url}, status code {r.status_code}.')
with open(filepath, 'wb') as f:
shutil.copyfileobj(r.raw, f, length=4194304)
class DistributionGridFactory:
"""
DistributionGrid Factory for helping create, cache and share the DistributionGrid based on the usage.
The DistributionGrid con be shared cross modules only the when this 3 conditions:
1. expert parallel group size
2. expert parallel replica group size,
are the same.
"""
distribution_grid_cache = {}
@classmethod
def get_distribution_grid(cls,
expert_parallel_group_size,
expert_parallel_replica_group_size,
ddp_type):
"""
Get the DistributionGrid by the conditions.
Args:
expert_parallel_group_size: expert parallel group size
expert_parallel_replica_group_size: expert parallel replica group size
ddp_type: distributed data parallel type. "DDP" of the recipe, only allow ddp_type is "MAINZ", "OSS" or "ShardedDDP".
Returns: new created DistributionGrid or shared DistributionGrid.
Notes: Currently get_distribution_grid only support "DDP" is "MAINZ", "OSS" or "ShardedDDP".
"""
# TODO: Support cases that "DDP" is "FSDP".
# For "FSDP", we use the DG of self.opt['fsdp_expert_grid'] which is initialize in DistributedTrainer directly.
ddp_type = ddp_type.upper()
assert ddp_type in ["MAINZ", "OSS", "SHARDEDDDP"], f'DistributionGrid Factory only support "DDP" is "MAINZ",' \
f' "OSS" or "ShardedDDP".' \
f' But currently "DDP" is {ddp_type}'
cached_distributed_grid = cls.distribution_grid_cache.get(
(expert_parallel_group_size, expert_parallel_replica_group_size), None)
if cached_distributed_grid is not None:
return cached_distributed_grid
else:
from ort_moe.grids import DistributionGrid
distributed_grid = DistributionGrid(expert_parallel_group_size=expert_parallel_group_size,
expert_parallel_replica_group_size=expert_parallel_replica_group_size)
cls.distribution_grid_cache[expert_parallel_group_size,
expert_parallel_replica_group_size] = distributed_grid
return distributed_grid
def get_world_size():
if not dist.is_available():
return 1
if not dist.is_initialized():
return 1
return dist.get_world_size()
def get_rank():
if not dist.is_available():
return 0
if not dist.is_initialized():
return 0
return dist.get_rank()
def synchronize():
"""
Helper function to synchronize (barrier) among all processes when
using distributed training
"""
if not dist.is_available():
return
if not dist.is_initialized():
return
world_size = dist.get_world_size()
rank = dist.get_rank()
if world_size == 1:
return
def _send_and_wait(r):
if rank == r:
tensor = torch.tensor(0, device="cuda")
else:
tensor = torch.tensor(1, device="cuda")
dist.broadcast(tensor, r)
while tensor.item() == 1:
time.sleep(1)
_send_and_wait(0)
# now sync on the main process
_send_and_wait(1)
def all_gather(data):
"""
Run all_gather on arbitrary picklable data (not necessarily tensors)
Args:
data: any picklable object
Returns:
list[data]: list of data gathered from each rank
"""
world_size = get_world_size()
if world_size == 1:
return [data]
# serialized to a Tensor
buffer = pickle.dumps(data)
storage = torch.ByteStorage.from_buffer(buffer)
tensor = torch.ByteTensor(storage).to("cuda")
# obtain Tensor size of each rank
local_size = torch.LongTensor([tensor.numel()]).to("cuda")
size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
dist.all_gather(size_list, local_size)
size_list = [int(size.item()) for size in size_list]
max_size = max(size_list)
# receiving Tensor from all ranks
# we pad the tensor because torch all_gather does not support
# gathering tensors of different shapes
tensor_list = []
for _ in size_list:
tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda"))
if local_size != max_size:
padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda")
tensor = torch.cat((tensor, padding), dim=0)
dist.all_gather(tensor_list, tensor)
data_list = []
for size, tensor in zip(size_list, tensor_list):
buffer = tensor.cpu().numpy().tobytes()[:size]
data_list.append(pickle.loads(buffer))
return data_list
def all_gather_cpu(data):
"""
Run all_gather on arbitrary picklable data (not necessarily tensors).
Args:
data: any picklable object
group: a torch process group. By default, will use a group which
contains all ranks on gloo backend.
Returns:
list[data]: list of data gathered from each rank
"""
def _get_global_gloo_group():
"""
Return a process group based on gloo backend, containing all the ranks
The result is cached.
"""
if dist.get_backend() == "nccl":
return dist.new_group(backend="gloo")
else:
return dist.group.WORLD
if get_world_size() == 1:
return [data]
group = _get_global_gloo_group() # use CPU group by default, to reduce GPU RAM usage.
world_size = dist.get_world_size(group)
if world_size == 1:
return [data]
output = [None for _ in range(world_size)]
dist.all_gather_object(output, data, group=group)
return output
def reduce_dict(input_dict, average=True):
"""
Args:
input_dict (dict): all the values will be reduced
average (bool): whether to do average or sum
Reduce the values in the dictionary from all processes so that process with rank
0 has the averaged results. Returns a dict with the same fields as
input_dict, after reduction.
"""
world_size = get_world_size()
if world_size < 2:
return input_dict
with torch.no_grad():
names = []
values = []
# sort the keys so that they are consistent across processes
for k in sorted(input_dict.keys()):
names.append(k)
values.append(input_dict[k])
values = torch.stack(values, dim=0)
dist.reduce(values, dst=0)
if dist.get_rank() == 0 and average:
# only main process gets accumulated, so only divide by
# world_size in this case
values /= world_size
reduced_dict = {k: v for k, v in zip(names, values)}
return reduced_dict
def broadcast_data(data):
if not torch.distributed.is_initialized():
return data
rank = dist.get_rank()
if rank == 0:
data_tensor = torch.tensor(data + [0], device="cuda")
else:
data_tensor = torch.tensor(data + [1], device="cuda")
torch.distributed.broadcast(data_tensor, 0)
while data_tensor.cpu().numpy()[-1] == 1:
time.sleep(1)
return data_tensor.cpu().numpy().tolist()[:-1]
def reduce_sum(tensor):
if get_world_size() <= 1:
return tensor
tensor = tensor.clone()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
return tensor
def save_result(result, filename):
output_folder = os.path.dirname(filename)
basename = os.path.splitext(os.path.basename(filename))[0]
os.makedirs(output_folder, exist_ok=True)
if isinstance(result, torch.Tensor) and result.ndim in [3,4]:
if result.ndim==3 and result.size(0) not in [1,3]:
result = make_grid(result.unsqueeze(1))
elif result.ndim==4:
result = make_grid(result)
else:
result = make_grid([result])
im = Image.fromarray(result.clamp_(0, 255).permute(1, 2, 0).to(torch.uint8).numpy())
im.save(os.path.join(output_folder, '{}.png'.format(basename)))
else:
torch.save(result, os.path.join(output_folder, '{}.pth'.format(basename)))