ShalomKing/infinitetalk

1from contextlib import contextmanager2 3import torch4 5@contextmanager6def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False):7    old_register_parameter = torch.nn.Module.register_parameter8    if include_buffers:9        old_register_buffer = torch.nn.Module.register_buffer10 11    def register_empty_parameter(module, name, param):12        old_register_parameter(module, name, param)13        if param is not None:14            param_cls = type(module._parameters[name])15            kwargs = module._parameters[name].__dict__16            kwargs["requires_grad"] = param.requires_grad17            module._parameters[name] = param_cls(18                module._parameters[name].to(device), **kwargs19            )20 21    def register_empty_buffer(module, name, buffer, persistent=True):22        old_register_buffer(module, name, buffer, persistent=persistent)23        if buffer is not None:24            module._buffers[name] = module._buffers[name].to(device)25 26    def patch_tensor_constructor(fn):27        def wrapper(*args, **kwargs):28            kwargs["device"] = device29            return fn(*args, **kwargs)30 31        return wrapper32 33    if include_buffers:34        tensor_constructors_to_patch = {35            torch_function_name: getattr(torch, torch_function_name)36            for torch_function_name in ["empty", "zeros", "ones", "full"]37        }38    else:39        tensor_constructors_to_patch = {}40 41    try:42        torch.nn.Module.register_parameter = register_empty_parameter43        if include_buffers:44            torch.nn.Module.register_buffer = register_empty_buffer45        for torch_function_name in tensor_constructors_to_patch.keys():46            setattr(47                torch,48                torch_function_name,49                patch_tensor_constructor(getattr(torch, torch_function_name)),50            )51        yield52    finally:53        torch.nn.Module.register_parameter = old_register_parameter54        if include_buffers:55            torch.nn.Module.register_buffer = old_register_buffer56        for (57            torch_function_name,58            old_torch_function,59        ) in tensor_constructors_to_patch.items():60            setattr(torch, torch_function_name, old_torch_function)