ml.utils.large_models

Utility functions for dealing with large models.

ml.utils.large_models.init_empty_weights(include_buffers: bool = False) → Iterator[None]

Avoid instantiating weights when initializing a model.

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.

Parameters:

include_buffers – Whether or not to also put all buffers on the meta device while initializing.

Yields:

An empty context manager

ml.utils.large_models.meta_to_empty_func(device: device | str, dtype: dtype | None = None) → Callable[[Tensor], Tensor]