meta

meta

MetaMixin

MetaMixin(*args, dispatch: bool = False, meta_buffers: bool = True, rename: Optional[Dict[str, str]] = None, **kwargs)

Bases: LoadableMixin

Mixin that adds lazy loading with meta tensors and deferred dispatch.

When dispatch=False (the default), the model is initialized with meta tensors (no memory allocation) using :meth:_load_meta. Real weights are loaded later when :meth:dispatch is called, or automatically on the first :meth:trace / :meth:generate call.

When dispatch=True or a pre-loaded torch.nn.Module is passed, the model is loaded immediately via :meth:_load.

PARAMETER	DESCRIPTION
*args	Forwarded to :meth:_load_meta or :meth:_load. DEFAULT: ()
dispatch	If True, load real weights immediately. Defaults to False (lazy / meta-tensor initialization). TYPE: bool DEFAULT: False
meta_buffers	If True, buffers are also created on the meta device. Defaults to True. TYPE: bool DEFAULT: True
rename	Module path aliases. TYPE: Optional[Dict[str, str]] DEFAULT: None
**kwargs	Forwarded to :meth:_load_meta or :meth:_load. DEFAULT: {}

ATTRIBUTE	DESCRIPTION
dispatched	Whether real weights have been loaded. TYPE: bool
args	Saved positional arguments for deferred :meth:dispatch.
kwargs	Saved keyword arguments for deferred :meth:dispatch.

dispatched instance-attribute

dispatched = False

args instance-attribute

args = args

kwargs instance-attribute

kwargs = kwargs

dispatch

dispatch() -> None

Load real weights into the model, replacing meta tensors.

Calls :meth:_load with the saved constructor arguments and updates the Envoy tree to reflect the newly loaded module. After this call, :attr:dispatched is True.

interleave

interleave(fn: Callable, *args, **kwargs)

Run fn with interleaved interventions, auto-dispatching if needed.

If the model has not been dispatched yet and this is not a scan, :meth:dispatch is called first and fn is rebound to the newly loaded module.

__setstate__

__setstate__(state)