matorus/replit-coder

1from contextlib import contextmanager2import torch3import torch.nn as nn4 5@contextmanager6def init_empty_weights(include_buffers: bool=False):7    """Meta initialization context manager.8 9    A context manager under which models are initialized with all parameters10    on the meta device, therefore creating an empty model. Useful when just11    initializing the model would blow the available RAM.12 13    Args:14        include_buffers (`bool`, *optional*, defaults to `False`): Whether or15            not to also put all buffers on the meta device while initializing.16 17    Example:18    ```python19    import torch.nn as nn20 21    # Initialize a model with 100 billions parameters in no time and without using any RAM.22    with init_empty_weights():23        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])24    ```25 26    <Tip warning={true}>27 28    Any model created under this context manager has no weights. As such you can't do something like29    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].30 31    </Tip>32    """33    with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:34        yield f35 36@contextmanager37def init_on_device(device: torch.device, include_buffers: bool=False):38    """Device initialization context manager.39 40    A context manager under which models are initialized with all parameters41    on the specified device.42 43    Args:44        device (`torch.device`): Device to initialize all parameters on.45        include_buffers (`bool`, *optional*, defaults to `False`): Whether or46            not to also put all buffers on the meta device while initializing.47 48    Example:49    ```python50    import torch.nn as nn51 52    with init_on_device(device=torch.device("cuda")):53        tst = nn.Liner(100, 100)  # on `cuda` device54    ```55    """56    old_register_parameter = nn.Module.register_parameter57    if include_buffers:58        old_register_buffer = nn.Module.register_buffer59 60    def register_empty_parameter(module, name, param):61        old_register_parameter(module, name, param)62        if param is not None:63            param_cls = type(module._parameters[name])64            kwargs = module._parameters[name].__dict__65            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)66 67    def register_empty_buffer(module, name, buffer):68        old_register_buffer(module, name, buffer)69        if buffer is not None:70            module._buffers[name] = module._buffers[name].to(device)71    if include_buffers:72        tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}73    else:74        tensor_constructors_to_patch = {}75 76    def patch_tensor_constructor(fn):77 78        def wrapper(*args, **kwargs):79            kwargs['device'] = device80            return fn(*args, **kwargs)81        return wrapper82    try:83        nn.Module.register_parameter = register_empty_parameter84        if include_buffers:85            nn.Module.register_buffer = register_empty_buffer86        for torch_function_name in tensor_constructors_to_patch.keys():87            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))88        yield89    finally:90        nn.Module.register_parameter = old_register_parameter91        if include_buffers:92            nn.Module.register_buffer = old_register_buffer93        for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():94            setattr(torch, torch_function_name, old_torch_function)