Skip to content

vllm.model_executor.model_loader.reload

Layerwise weight reloading utilities for vLLM.

This module provides functionality to reload model weights layer-by-layer, which is useful for weight updates without full model reconstruction.

Limitations: 1. Composition with CPU offloading has not been implemented 2. Tied parameters will only reflect processing from one of the parent layers (for example, only processing from embed_tokens will have an effect) 3. This design assumes that the number of weights loaded from disk is the same as the number of weights created at model init time. This is not true for quant methods which (1) pad weights or (2) load qkv weights into the same parameter. Both of these cases are non-issues for today's quant methods, but future quantizations may cause reloading to fail

Modules:

  • bake –

    The load plan's bake: run load_weights once over TraceTensors, on

  • direct –

    Direct weight reload (VLLM_RELOAD_MODE=direct).

  • geometry –

    Integer geometry of the load plan: one traced copy_ as a box of the

  • inplace –
  • layerwise –
  • lifecycle –

    Begin and complete a checkpoint-format weight update in the reload mode

  • meta –
  • plan –

    Load plan: which part of each checkpoint tensor this rank needs, and where

  • sanitize –
  • torchao_decorator –
  • trace –

    Tracing a model's weight loaders on zero-storage checkpoint tensors.

  • utils –

Classes:

Functions:

LoadPlan dataclass

Methods:

  • get_boxes –

    Parts of the checkpoint tensor this rank needs; None: nothing.

  • get_dest –

    (part of a box, live tensor) pairs: copying each part, reshaped to

Source code in vllm/model_executor/model_loader/reload/plan.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
@dataclass(frozen=True)
class LoadPlan:
    entries: dict[str, Entry]  # per declared name, in metadata order
    dests: dict[str, Dest]  # live tensors built at plan_loads; not data

    def get_boxes(self, name: str) -> tuple[Box, ...] | None:
        """Parts of the checkpoint tensor this rank needs; ``None``: nothing.

        For engines that send or read only those slices (sharded RDT, NCCL
        M2N follow-ups); an engine that receives full tensors calls
        ``load_planned`` instead."""
        return self.entries[name].boxes

    def get_dest(self, name: str) -> list[tuple[Box, torch.Tensor]] | None:
        """(part of a box, live tensor) pairs: copying each part, reshaped to
        its tensor's shape, into that tensor is exactly what the loader would
        have written. ``None``: fallback, send the full tensor through
        ``load_weights``. Write only between ``start_reload`` and
        ``finish_reload``; never read.

        ``load_planned`` copies full tensors through it; an engine that sends
        only ``get_boxes`` copies its slices the same way."""
        if self.entries[name].writes is None:
            return None
        return list(self.dests[name])

get_boxes(name)

Parts of the checkpoint tensor this rank needs; None: nothing.

For engines that send or read only those slices (sharded RDT, NCCL M2N follow-ups); an engine that receives full tensors calls load_planned instead.

Source code in vllm/model_executor/model_loader/reload/plan.py 
71
72
73
74
75
76
77
def get_boxes(self, name: str) -> tuple[Box, ...] | None:
    """Parts of the checkpoint tensor this rank needs; ``None``: nothing.

    For engines that send or read only those slices (sharded RDT, NCCL
    M2N follow-ups); an engine that receives full tensors calls
    ``load_planned`` instead."""
    return self.entries[name].boxes

get_dest(name)

(part of a box, live tensor) pairs: copying each part, reshaped to its tensor's shape, into that tensor is exactly what the loader would have written. None: fallback, send the full tensor through load_weights. Write only between start_reload and finish_reload; never read.

load_planned copies full tensors through it; an engine that sends only get_boxes copies its slices the same way.

Source code in vllm/model_executor/model_loader/reload/plan.py 
79
80
81
82
83
84
85
86
87
88
89
90
def get_dest(self, name: str) -> list[tuple[Box, torch.Tensor]] | None:
    """(part of a box, live tensor) pairs: copying each part, reshaped to
    its tensor's shape, into that tensor is exactly what the loader would
    have written. ``None``: fallback, send the full tensor through
    ``load_weights``. Write only between ``start_reload`` and
    ``finish_reload``; never read.

    ``load_planned`` copies full tensors through it; an engine that sends
    only ``get_boxes`` copies its slices the same way."""
    if self.entries[name].writes is None:
        return None
    return list(self.dests[name])

check_plannable(model)

Refuse a configuration no load plan supports: a reload mode other than direct, and EPLB. Reads only configuration, so an engine can call it before start_reload changes the model. Direct reload's own refusals run in begin_direct_reload.

Raises:

Source code in vllm/model_executor/model_loader/reload/plan.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def check_plannable(model: torch.nn.Module) -> None:
    """Refuse a configuration no load plan supports: a reload mode other than
    direct, and EPLB. Reads only configuration, so an engine can call it
    before ``start_reload`` changes the model. Direct reload's own refusals
    run in ``begin_direct_reload``.

    Raises:
        RuntimeError: the configuration cannot be planned.

    """
    if not direct_reload_enabled():
        raise RuntimeError("load plan: requires VLLM_RELOAD_MODE=direct")
    for name, module in model.named_modules():
        moe_config = getattr(module, "moe_config", None)
        parallel = getattr(moe_config, "moe_parallel_config", None)
        if getattr(parallel, "enable_eplb", False):
            raise RuntimeError(f"load plan: EPLB is not supported ({name})")

finalize_layerwise_processing(model, model_config)

Apply processing to any layers which were not layerwise processed during loading. This includes attention layers and layers which have weight elements which are not loaded (due to padding).

This function should be applied after initialize_layerwise_reload is applied unwrap the layerwise weight loaders.

Parameters:

  • model

    (Module) –

    model to finalize processing for

  • model_config

    (ModelConfig) –

    config needed for applying processing to attention layers

Source code in vllm/model_executor/model_loader/reload/layerwise.py 
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelConfig):
    """Apply processing to any layers which were not layerwise processed during loading.
    This includes attention layers and layers which have weight elements which are not
    loaded (due to padding).

    This function should be applied after `initialize_layerwise_reload` is applied
    unwrap the layerwise weight loaders.

    Args:
        model: model to finalize processing for
        model_config: config needed for applying processing to attention layers

    """
    if hasattr(model, "_original_do_torchao_reload"):
        model._do_torchao_reload = model._original_do_torchao_reload

    deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = []

    for layer in model.modules():
        info = get_layerwise_info(layer)
        if not info.can_load():
            info.reset()
            continue

        # Deferred attention-like layers are processed after all other layers
        if is_deferred_attention_layer(layer):
            deferred_attn.append((layer, info))
            continue

        # No weights were loaded
        if info.load_numel <= 0:
            # first load: checkpoint did not contain weights for this layer
            if info.kernel_tensors is None:
                _layerwise_process(layer, info)
                continue

            # reloading: place kernel tensors back as a fallback. Always place, even
            # when nothing is loadable (load_numel_total == 0), so parameter-alias
            # buffers on such layers are restored rather than left deleted.
            if info.load_numel_total > 0:  # type: ignore[operator]
                logger.warning("%s: Failed to load weights", layer.__class__.__name__)
            _place_kernel_tensors(layer, info)

        # Process non-attention layers which did not load all elements. This can happen
        # if the created weight has extra padding elements which are not loaded
        # Having too many of these delayed layers can lead to excess memory usage
        # see Limitations(4)
        elif info.defer_processing or info.load_numel < info.load_numel_total:  # type: ignore[operator]
            logger.debug("%s: Delayed processing", layer.__class__.__name__)
            _layerwise_process(layer, info)

        info.reset()

    # Process attention layers after all other layers are done
    for layer, info in deferred_attn:
        _finalize_attention_layer(layer, info, model_config)
        info.reset()

    LOADING_LAYERS.clear()

finish_reload(model, model_config, device)

Complete the weight update start_reload began.

Source code in vllm/model_executor/model_loader/reload/lifecycle.py 
24
25
26
27
28
29
30
31
def finish_reload(
    model: torch.nn.Module, model_config: ModelConfig, device: torch.device
) -> None:
    """Complete the weight update `start_reload` began."""
    if direct_reload_enabled():
        finish_direct_reload(model, model_config, device)
    else:
        finalize_layerwise_reload(model, model_config)

initialize_layerwise_reload(model, *, defer_processing=False)

Set up layerwise weight loading with deferred processing.

Must be called after record_metadata_for_reloading. This function: 1. Saves current kernel tensors for later copying 2. Restores layer parameters/buffers from metadata (on meta device) 3. Wraps weight loaders to defer processing until all weights are loaded

When all weights for a layer are loaded, the wrapped loaders will: 1. Materialize the layer onto the target device 2. Load all cached weights 3. Run quantization processing if applicable 4. Copy processed values back to original tensor storage

With defer_processing=True, only finalize_layerwise_reload triggers processing.

Source code in vllm/model_executor/model_loader/reload/layerwise.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
@torch.no_grad()
def initialize_layerwise_reload(
    model: torch.nn.Module, *, defer_processing: bool = False
):
    """Set up layerwise weight loading with deferred processing.

    Must be called after `record_metadata_for_reloading`. This function:
    1. Saves current kernel tensors for later copying
    2. Restores layer parameters/buffers from metadata (on meta device)
    3. Wraps weight loaders to defer processing until all weights are loaded

    When all weights for a layer are loaded, the wrapped loaders will:
    1. Materialize the layer onto the target device
    2. Load all cached weights
    3. Run quantization processing if applicable
    4. Copy processed values back to original tensor storage

    With defer_processing=True, only finalize_layerwise_reload triggers processing.
    """
    # disable torchao reloading to avoid infinite recursion
    model._original_do_torchao_reload = getattr(model, "_do_torchao_reload", False)
    model._do_torchao_reload = False

    for layer in model.modules():
        info = get_layerwise_info(layer)

        # Skip if the layer has already been initialized
        if info.can_load():
            continue

        # Save current tensors for later copying
        info.kernel_tensors = get_layer_params_buffers(layer)
        # snapshot now: restore_layer_on_meta drops alias buffers from the live set
        info.kernel_non_persistent_buffers = set(layer._non_persistent_buffers_set)

        # Restore layer parameters/buffers onto meta device
        restore_layer_on_meta(layer, info)

        # Wrap weight loaders to buffer loading
        initialize_online_processing(layer)
        info.defer_processing = defer_processing

load_planned(model, plan, weights)

Load full checkpoint tensors as plan says: copy each direct name's boxes into its destinations, skip names this rank needs nothing of, and pass fallbacks to model.load_weights. Call between start_reload and finish_reload.

Source code in vllm/model_executor/model_loader/reload/plan.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def load_planned(
    model: torch.nn.Module,
    plan: LoadPlan,
    weights: Iterable[tuple[str, torch.Tensor]],
) -> None:
    """Load full checkpoint tensors as ``plan`` says: copy each direct name's
    boxes into its destinations, skip names this rank needs nothing of, and
    pass fallbacks to ``model.load_weights``. Call between ``start_reload``
    and ``finish_reload``."""
    fallbacks = []
    for name, full in weights:
        dest = plan.get_dest(name)
        if dest is None:
            fallbacks.append((name, full))
        for box, tensor in dest or ():
            part = full[tuple(slice(start, stop) for start, stop in box)]
            tensor.copy_(part.reshape(tensor.shape))
    if fallbacks:
        model.load_weights(fallbacks)

plan_loads(model, metadata)

Trace model.load_weights over the declared checkpoint tensors once and return where each lands on this rank.

Call between start_reload and finish_reload: the plan traces the real tensors in the layout begin_direct_reload applied. Incoming tensors are data-free traces, so nothing is loaded. Loaders' constant writes land, as on the load_weights path; so do writes into non-checkpoint GPU tensors fed by a read of a checkpoint tensor (a mirror), which are not restored: that name falls back and the update rewrites them. Module, parameter and buffer state a loader sets stays, as it would after load_weights.

Raises:

Source code in vllm/model_executor/model_loader/reload/plan.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def plan_loads(model: torch.nn.Module, metadata: "Iterable[ParamMeta]") -> LoadPlan:
    """Trace ``model.load_weights`` over the declared checkpoint tensors once
    and return where each lands on this rank.

    Call between ``start_reload`` and ``finish_reload``: the plan traces
    the real tensors in the layout ``begin_direct_reload`` applied. Incoming
    tensors are data-free traces, so nothing is loaded. Loaders' constant
    writes land, as on the ``load_weights`` path; so do writes into
    non-checkpoint GPU tensors fed by a read of a checkpoint tensor (a
    mirror), which are not restored: that name falls back and the update
    rewrites them. Module, parameter and buffer state a loader sets stays,
    as it would after ``load_weights``.

    Raises:
        ValueError: a name is declared twice.
        RuntimeError: the model cannot be planned.

    """
    check_plannable(model)
    if not direct_reload_active(model):
        raise RuntimeError(
            "load plan: call plan_loads between start_reload and finish_reload; "
            "the plan takes its layout from what the direct reload restored"
        )
    metas: dict[str, ParamMeta] = {}
    for meta in metadata:
        if meta.name in metas:
            raise ValueError(f"load plan: {meta.name!r} is declared twice")
        metas[meta.name] = meta
    # TODO(#54477 F2): refuse when a name of the required set (cold-start
    # checkpoint names minus `is_unused_checkpoint_weight` and
    # VLLM_DIRECT_RELOAD_EXCLUDE_REGEX) is missing from `metas`.

    t0 = time.perf_counter()
    targets = Targets(model)
    traced = trace_loads(model, metas, targets)
    if traced.lonely:
        lonely = traced.lonely
        raise RuntimeError(
            f"load plan: {len(lonely)} fallback names write nothing themselves "
            f"(e.g. {lonely[:3]}); their loaders combine several checkpoint "
            "tensors. Move the combine into process_weights_after_loading."
        )
    for name, reason in _overlaps(traced, targets).items():
        traced.reasons.setdefault(name, reason)
    entries = {name: _entry(meta, traced, targets) for name, meta in metas.items()}
    dests = {
        name: tuple((w.src, _dest(targets, w, e.dtype)) for w in e.writes)
        for name, e in entries.items()
        if e.writes is not None
    }
    plan = LoadPlan(entries=entries, dests=dests)
    _report(plan, time.perf_counter() - t0)
    return plan

record_metadata_for_reloading(model)

Record layer metadata needed for later reloading.

Stores parameter and buffer metadata as meta tensors for restoration. Must be called before initialize_layerwise_reload.

Source code in vllm/model_executor/model_loader/reload/layerwise.py 
74
75
76
77
78
79
80
81
82
83
def record_metadata_for_reloading(model: torch.nn.Module):
    """Record layer metadata needed for later reloading.

    Stores parameter and buffer metadata as meta tensors for restoration.
    Must be called before `initialize_layerwise_reload`.
    """
    for layer in model.modules():
        info = get_layerwise_info(layer)
        info.restore_metadata = capture_layer_to_meta(layer)
        info.restore_device = torch.get_default_device()

start_reload(model)

Begin a weight update in the configured mode.

Source code in vllm/model_executor/model_loader/reload/lifecycle.py 
16
17
18
19
20
21
def start_reload(model: torch.nn.Module) -> None:
    """Begin a weight update in the configured mode."""
    if direct_reload_enabled():
        begin_direct_reload(model)
    else:
        initialize_layerwise_reload(model)

support_quantized_model_reload_from_hp_weights(original_load_weights)

Decorator for load_weights method for AutoWeightsLoader.load_weights to support reloading high precision (bfloat16/float16/float32) weight for an already quantized model, this involves restoring the weights to a high precision weights and then online quantize the weights.

Only applies to torchao quantized models. Assumes that all model weights are loaded within a single weights iterator (cannot perform batched updates)

Source code in vllm/model_executor/model_loader/reload/torchao_decorator.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def support_quantized_model_reload_from_hp_weights(
    original_load_weights: Callable[Concatenate["AutoWeightsLoader", _P], _R],
) -> Callable[Concatenate["AutoWeightsLoader", _P], _R]:
    """Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
    reloading high precision (bfloat16/float16/float32) weight for an already quantized
    model, this involves restoring the weights to a high precision weights and
    then online quantize the weights.

    Only applies to torchao quantized models. Assumes that all model weights are
    loaded within a single weights iterator (cannot perform batched updates)
    """

    @wraps(original_load_weights)
    def patched_model_load_weights(
        self: "AutoWeightsLoader",
        *args: _P.args,
        **kwargs: _P.kwargs,
    ) -> _R:
        model = self.module

        if not getattr(model, "_do_torchao_reload", False):
            return original_load_weights(self, *args, **kwargs)

        initialize_layerwise_reload(model)
        loaded_weights = original_load_weights(self, *args, **kwargs)
        finalize_layerwise_processing(model, model._model_config)

        return loaded_weights

    return patched_model_load_weights