meta¶
Building a model on the meta device, dispatching real weights only on demand.
Loading a large model's weights is slow and memory-hungry, but most of what
nnsight needs to plan a trace — the module tree and the shapes that flow through
it — is fixed by the model's structure, not its weights. So a model is built
lazily: every module is constructed on the meta device, which records a
tensor's shape and dtype but allocates no storage. The envoy tree, scan,
and remote-key resolution all work against this weightless skeleton.
MetaDevice is the mechanism — a torch function mode that forces every
tensor created (or moved) within it onto meta, however it was created.
Meta is the policy — it does the meta build up front and dispatches
(loads real weights and swaps them into the same envoy tree) the first time the
model actually has to run.
.. code-block:: python
model = TransformersModel("openai-community/gpt2") # meta build, no weights
model.scan("Hello") # inspect shapes, still no weights
with model.trace("Hello"): # dispatches real weights on first run
...
The meta build reconstructs structure only, so a subclass's Meta._load_meta
drops weight- and placement-related kwargs (device_map, max_memory, ...):
they mean nothing on storageless meta tensors and take effect only at dispatch.
Bases: TorchFunctionMode
Force every tensor created within the context onto the meta device.
Setting a default device only covers factory calls that omit device=;
this also rewrites explicit device= arguments, so tensors land on meta
no matter how they are created.
real opens a nested window where this forcing is suspended, for the
parts of a lazy build that need genuine tensors (e.g. a component whose
constructor moves a buffer with .to(...), which a meta tensor can't do).
classmethod
¶
Suspend meta-forcing within this block so tensors are created for real.
Bases: Loadable
A model wrapper that builds on meta and dispatches real weights on demand.
Extends Loadable with a two-phase
build: _load_meta constructs the weightless skeleton up front, and
dispatch (triggered on the first interleave that isn't a scan)
loads real weights via _load and swaps them into the existing envoy tree.
Passing a ready module, or dispatch=True, skips the meta phase and loads
eagerly.
Load real weights via _load and swap them into the envoy tree.
Idempotent — a no-op once dispatched. Triggered automatically on the first
real (non-scan) interleave; call it directly to load eagerly.
scan(*args: Any, fn: Any = None, backend: Any = None, **kwargs: Any) -> ScanningTracer
Inspect activation shapes without loading real weights or computing.
Just like trace, but the forward
runs under a fake-tensor mode so only tensor metadata (shapes, dtypes)
propagates — the model is never dispatched, so this works on a
meta-initialized model with no weights in memory:
>>> model = TransformersModel("openai-community/gpt2") # not dispatched
>>> with model.scan("Hello World"):
... hidden = model.transformer.h[0].output[0].shape
>>> print(hidden) # torch.Size([...])
The values read inside the block are fake tensors valid only within it;
read their .shape / .dtype here rather than saving the tensor out.
| PARAMETER | DESCRIPTION |
|---|---|
*args
|
Inputs to scan, forwarded to the model's forward pass.
TYPE:
|
fn
|
Method to run, resolved against the module. Defaults to
TYPE:
|
backend
|
Backend that runs the captured block. Defaults to in place.
TYPE:
|
**kwargs
|
Keyword inputs forwarded to the forward pass.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
ScanningTracer
|
A |