Alchan/mpt-7b-chat

1from contextlib import contextmanager2from typing import Any, Callable, Optional3import torch4import torch.nn as nn5 6@contextmanager7def init_empty_weights(include_buffers: bool=False):8    """Meta initialization context manager.9 10    A context manager under which models are initialized with all parameters11    on the meta device, therefore creating an empty model. Useful when just12    initializing the model would blow the available RAM.13 14    Args:15        include_buffers (`bool`, *optional*, defaults to `False`): Whether or16            not to also put all buffers on the meta device while initializing.17 18    Example:19    ```python20    import torch.nn as nn21 22    # Initialize a model with 100 billions parameters in no time and without using any RAM.23    with init_empty_weights():24        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])25    ```26 27    <Tip warning={true}>28 29    Any model created under this context manager has no weights. As such you can't do something like30    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].31 32    </Tip>33    """34    with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:35        yield f36 37@contextmanager38def init_on_device(device: torch.device, include_buffers: bool=False):39    """Device initialization context manager.40 41    A context manager under which models are initialized with all parameters42    on the specified device.43 44    Args:45        device (`torch.device`): Device to initialize all parameters on.46        include_buffers (`bool`, *optional*, defaults to `False`): Whether or47            not to also put all buffers on the meta device while initializing.48 49    Example:50    ```python51    import torch.nn as nn52 53    with init_on_device(device=torch.device("cuda")):54        tst = nn.Liner(100, 100)  # on `cuda` device55    ```56    """57    old_register_parameter = nn.Module.register_parameter58    if include_buffers:59        old_register_buffer = nn.Module.register_buffer60 61    def register_empty_parameter(self: torch.nn.Module, name: str, param: Optional[torch.nn.Parameter]):62        old_register_parameter(self, name, param)63        if param is not None:64            parameter = self._parameters[name]65            assert parameter is not None66            param_cls = type(parameter)67            kwargs = parameter.__dict__68            self._parameters[name] = param_cls(parameter.to(device), **kwargs)69 70    def register_empty_buffer(self: torch.nn.Module, name: str, tensor: Optional[torch.Tensor], persistent: bool=True):71        old_register_buffer(self, name, tensor, persistent=persistent)72        if tensor is not None:73            named_buffer = self._buffers[name]74            assert named_buffer is not None75            self._buffers[name] = named_buffer.to(device)76    if include_buffers:77        tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}78    else:79        tensor_constructors_to_patch = {}80 81    def patch_tensor_constructor(fn: Callable):82 83        def wrapper(*args: Any, **kwargs: Any):84            kwargs['device'] = device85            return fn(*args, **kwargs)86        return wrapper87    try:88        nn.Module.register_parameter = register_empty_parameter89        if include_buffers:90            nn.Module.register_buffer = register_empty_buffer91        for torch_function_name in tensor_constructors_to_patch.keys():92            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))93        yield94    finally:95        nn.Module.register_parameter = old_register_parameter96        if include_buffers:97            nn.Module.register_buffer = old_register_buffer98        for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():99            setattr(torch, torch_function_name, old_torch_function)