util.py

import copy
import datetime
import os
import sys
from typing import Sequence, Iterator

import pytz
import torch
import torch.distributed as dist

import globals as g
from fp16_opt import FP16_Optimizer


def toscalar(t):  # use on python scalars/pytorch scalars
    """Converts Python scalar or PyTorch tensor to Python scalar"""
    if isinstance(t, (float, int)):
        return t
    if hasattr(t, 'float'):
        t = t.float()  # half not supported on CPU
    if hasattr(t, 'item'):
        return t.item()
    else:
        assert len(t) == 0
        return t[0]


def _info(_type, value, tb):
    if hasattr(sys, 'ps1') or not sys.stderr.isatty():
        # we are in interactive mode or we don't have a tty-like
        # device, so we call the default hook
        sys.__excepthook__(type, value, tb)
    else:
        import traceback
        import pdb
        # we are NOT in interactive mode, print the exception...
        traceback.print_exception(type, value, tb)
        print()
        # ...then start the debugger in post-mortem mode.
        # pdb.pm() # deprecated
        pdb.post_mortem(tb)  # more "modern"


def pdb_on_error():
    # todo(y): doesn't work when called from other files?
    sys.excepthook = _info


def get_world_size() -> int:
    return int(os.environ.get('WORLD_SIZE', 1))


def get_global_rank() -> int:
    """Returns global rank (from env), or 0 if not set"""
    return int(os.environ.get('RANK', 0))


def one_of(l):
    assert len(l) == 2
    if l[0]:
        return l[0]
    elif l[1]:
        return l[1]
    else:
        assert f"List {l} has more than one non-zero entries"


def dist_sum_tensor(tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    return rt


def dist_mean(tensor):
    is_scalar = False
    if not isinstance(tensor, torch.Tensor):
        tensor = torch.tensor(tensor, device="cuda")
        is_scalar = True
    mean = dist_sum_tensor(tensor) / get_world_size()
    if is_scalar:
        return mean.item()
    return mean


def unwrap_model(model):
    from mem_transformer import MemTransformerLM
    i = 0
    while not isinstance(model, MemTransformerLM):
        model = model.module
        i += 1
        if i > 2:
            assert False, "Are here more than 2 wraps?"
    return model


def unwrap_fp16_optimizer(optimizer):
    if isinstance(optimizer, FP16_Optimizer):
        return optimizer.optimizer
    else:
        return optimizer


# no_op method/object that accept every signature
class NoOp:
    def __getattr__(self, *_args):
        def no_op(*_args, **_kwargs): pass

        return no_op


# Deprecated method, regular restore + DDP already broadcasts args
# def dist_restore_from_checkpoint(ddp_model, checkpoint_fn: str, force_fp16=False):
#     """Restores model wrapped in DistributedDataParallel from checkpoint file. Assumes checkpoint was saved
#     as torch.save(ddp.module) or distributed_save_checkpoint
#     """

#     if get_global_rank() == 0:
#         saved_model = torch.load(checkpoint_fn)
#         state_dict = saved_model.state_dict()
#         if force_fp16:
#             for name in state_dict:
#                 state_dict[name] = state_dict[name].half()
#         ddp_model.module.load_state_dict(state_dict)

#     pp = next(ddp_model.module.parameters())
#     print(f"{get_global_rank()}  -- Before broadcast {pp.view(-1)[0]}")
#     for p in ddp_model.module.parameters():
#         if torch.is_tensor(p):
#             dist.broadcast(p, 0)
#     print(f"{get_global_rank()}  -- After broadcast {pp.view(-1)[0]}")


def restore_from_checkpoint(model, optimizer=None, checkpoint_fn: str = '', optimizer_state_dict_fn: str = '',
                            force_fp16=False, override_lr=None):
    """Restores model wrapped in DistributedDataParallel or/and FP16_Module from checkpoint file.
    Assumes checkpoint was saved as torch.save(unwrap_model(ddp_FP16_module)).

    If optimizer_state_dict_fn is provided, also tries to restore optimizer state from state_dict saved in that file.

    Assumes optimizer is regular optimizer, not FP16Optimizer(optimizer), must wrap FP16 on top
    of restored optimizer here.
    """

    saved_model = torch.load(checkpoint_fn, map_location="cpu")
    state_dict = saved_model.state_dict()
    if force_fp16:
        for name in state_dict:
            state_dict[name] = state_dict[name].half()
    model.load_state_dict(state_dict, strict=False)

    assert 'FP16_Optimizer' not in type(optimizer).__name__, \
        f"Checkpoint restore works on PyTorch optimizers, but found {type(optimizer).__name__}, " \
        f"you must unwrap your optimizer first"
    if optimizer_state_dict_fn:
        optimizer_state_dict = torch.load(optimizer_state_dict_fn, map_location="cpu")
        # another layer of indirection added for FP16Optimizer
        if 'optimizer_state_dict' in optimizer_state_dict:
            optimizer_state_dict = optimizer_state_dict['optimizer_state_dict']
        if override_lr:
            optimizer_state_dict['param_groups'][0]['lr'] = override_lr
        optimizer.load_state_dict(optimizer_state_dict)


def dist_save_checkpoint(ddp_fp16_model, optimizer, directory: str, suffix=''):
    """Saves model/optimizer into {directory}/optimizer-{suffix}.py and {directory}/model-{suffix}.pt"""
    if get_global_rank() != 0:
        return
    with open(directory + f'/model-{suffix}.pt', 'wb') as f_1:
        torch.save(unwrap_model(ddp_fp16_model), f_1)
    with open(directory + f'/optimizer-{suffix}.pt', 'wb') as f_1:
        torch.save(unwrap_fp16_optimizer(optimizer).state_dict(), f_1)


def get_hash(o):
    import pickle
    pickle.dump(o, open('/tmp/util_hash', 'wb'))
    return hash(open('/tmp/util_hash', 'rb').read())


def cancel_shutdown():
    args = g.args
    if args.local:
        return
    if args.local_rank > 0:
        return
    os.system('shutdown -c')


def current_timestamp(timezone: str = 'America/Los_Angeles') -> str:
    """Gives timestamp formated like 2019-04-15_11-29-51.
    correct to local timezone (PDT) if running on AWS (which is UTC)"""
    pacific_tz = pytz.timezone(timezone)
    localtime = pytz.utc.localize(datetime.datetime.now(), is_dst=None).astimezone(pacific_tz)
    return localtime.strftime('%Y-%m-%d_%H-%M-%S')


def assert_close(observed, target, rtol=1e-5, atol=1e-3):
    relative = abs(target - observed) / target
    assert relative < rtol, f"rtol {rtol} exceeded at {relative}, observed={observed}, target={target}"

    absolute = abs(target - observed)
    assert absolute < rtol, f"atol {atol} exceeded at {absolute}, observed={observed}, target={target}"


def assert_args_equal(args1, args2):
    args1 = vars(args1)
    args2 = vars(args2)
    keys = set(args1.keys()).union(args2.keys())
    for key in keys:
        assert key in args1, f"{key} not found in args1"
        assert key in args2, f"{key} not found in args2"
        assert args1[key] == args2[key], f"args not equal for key={key}, {args1[key]} != {args2[key]}"


def merge_args_from_state(args, state):
    args = vars(args)
    state_args = vars(state.args)

    attr_to_merge = ['fp16', 'dynamic_loss_scale', 'static_loss_scale']
    for attr in attr_to_merge:
        assert args[attr] == state_args[attr]  # TODO(y): decide which setting has precedence when attributes conflict
        args[attr] = state_args[attr]


def flat_grad_model(model):
    flat = None
    for param in model.parameters():
        if param is None:
            continue
        if param.grad is None:
            continue
        if flat is None:
            flat = param.grad.data.reshape(-1).clone()
        else:
            flat = torch.cat((flat, param.grad.data.reshape(-1)), 0)
    return torch.Tensor() if flat is None else flat


def flat_grad_opt(optimizer):
    flat = None
    for param_group in optimizer.param_groups:
        for param in param_group['params']:
            if param is None:
                continue
            if param.grad is None:
                continue
            if flat is None:
                flat = param.grad.data.reshape(-1).clone()
            else:
                flat = torch.cat((flat, param.grad.data.reshape(-1)), 0)
    return torch.Tensor() if flat is None else flat


def flat_param(model):
    flat = None
    for param in model.parameters():
        if param is None:
            continue
        if flat is None:
            flat = param.data.reshape(-1).clone()
        else:
            flat = torch.cat((flat, param.data.reshape(-1)), 0)
    return torch.Tensor() if flat is None else flat


# Debugging
record_dict = {}


def record(tag, value):
    record_dict.setdefault(tag, {})[g.state.train_step] = value


def dump_records():
    global record_dict
    t = copy.deepcopy(record_dict)
    record_dict = {}
    return t


def get_parameters(var):
    """Walk backward from node and find all Variables in given autograd graph.
    based on https://github.com/szagoruyko/pytorchviz/blob/master/torchviz/dot.py"""

    seen = set()
    variables = set()

    def add_nodes(vv):
        if vv not in seen:
            seen.add(vv)
            if torch.is_tensor(vv):
                pass
            elif hasattr(vv, 'variable'):
                u = vv.variable
                variables.add(u)
            if hasattr(vv, 'next_functions'):
                for u in vv.next_functions:
                    if u[0] is not None:
                        add_nodes(u[0])
            if hasattr(vv, 'saved_tensors'):
                for t in vv.saved_tensors:
                    add_nodes(t)

    if isinstance(var, tuple):
        for v in var:
            add_nodes(v.grad_fn)
    else:
        add_nodes(var.grad_fn)
    return variables


class FrozenClass(object):
    __isfrozen = False

    def __setattr__(self, key, value):
        if self.__isfrozen and not hasattr(self, key):
            raise TypeError("%r is a frozen class" % self)
        object.__setattr__(self, key, value)

    def _freeze(self):
        self.__isfrozen = True


class SaveableIteratorMaker:
    """Iterator over sequences that remembers its position on unpickling"""

    def __init__(self, data: Sequence, offset: int = 0):
        self.data = data
        self.offset = offset

    def __iter__(self):
        # wrap iterator
        #        while self.offset >= len(self.data):
        #            self.offset -= len(self.data)

        for i in range(self.offset, len(self.data)):
            self.offset = i
            yield self.data[i]


def saveable_iterator(data: Sequence) -> Iterator:
    return iter(SaveableIteratorMaker(data))


def download_from_s3(url, job):
    downloaded_fn = os.path.basename(url)
    job.run(f'rm {downloaded_fn}; wget {url}')
    return downloaded_fn