from contextlib import contextmanager from typing import Any, Callable, Optional import torch import torch.nn as nn @contextmanager def init_empty_weights(include_buffers: bool = False): """Meta initialization context manager. A context manager under which models are initialized with all parameters on the meta device, therefore creating an empty model. Useful when just initializing the model would blow the available RAM. Args: include_buffers (`bool`, *optional*, defaults to `False`): Whether or not to also put all buffers on the meta device while initializing. Example: ```python import torch.nn as nn # Initialize a model with 100 billions parameters in no time and without using any RAM. with init_empty_weights(): tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)]) ``` Any model created under this context manager has no weights. As such you can't do something like `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`]. """ with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f: yield f @contextmanager def init_on_device(device: torch.device, include_buffers: bool = False): """Device initialization context manager. A context manager under which models are initialized with all parameters on the specified device. Args: device (`torch.device`): Device to initialize all parameters on. include_buffers (`bool`, *optional*, defaults to `False`): Whether or not to also put all buffers on the meta device while initializing. Example: ```python import torch.nn as nn with init_on_device(device=torch.device("cuda")): tst = nn.Liner(100, 100) # on `cuda` device ``` """ old_register_parameter = nn.Module.register_parameter if include_buffers: old_register_buffer = nn.Module.register_buffer def register_empty_parameter( self: torch.nn.Module, name: str, param: Optional[torch.nn.Parameter] ): old_register_parameter(self, name, param) if param is not None: parameter = self._parameters[name] assert parameter is not None param_cls = type(parameter) kwargs = parameter.__dict__ self._parameters[name] = param_cls(parameter.to(device), **kwargs) def register_empty_buffer( self: torch.nn.Module, name: str, tensor: Optional[torch.Tensor], persistent: bool = True, ): old_register_buffer(self, name, tensor, persistent=persistent) if tensor is not None: named_buffer = self._buffers[name] assert named_buffer is not None self._buffers[name] = named_buffer.to(device) if include_buffers: tensor_constructors_to_patch = { torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ["empty", "zeros", "ones", "full"] } else: tensor_constructors_to_patch = {} def patch_tensor_constructor(fn: Callable): def wrapper(*args: Any, **kwargs: Any): kwargs["device"] = device return fn(*args, **kwargs) return wrapper try: nn.Module.register_parameter = register_empty_parameter if include_buffers: nn.Module.register_buffer = register_empty_buffer for torch_function_name in tensor_constructors_to_patch.keys(): setattr( torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)), ) yield finally: nn.Module.register_parameter = old_register_parameter if include_buffers: nn.Module.register_buffer = old_register_buffer for ( torch_function_name, old_torch_function, ) in tensor_constructors_to_patch.items(): setattr(torch, torch_function_name, old_torch_function)