Skip to content

vllm.model_executor.offloader.base

Base classes for model parameter offloading.

Classes:

  • BaseOffloader –

    Base class for model parameter offloading strategies.

  • NoopOffloader –

    No-op offloader that returns modules as-is without any offloading.

Functions:

BaseOffloader

Bases: ABC

Base class for model parameter offloading strategies.

Offloaders control how model parameters are stored and loaded during inference. Different strategies trade memory for compute/transfer time.

Methods:

Attributes:

Source code in vllm/model_executor/offloader/base.py 
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 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
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
class BaseOffloader(ABC):
    """Base class for model parameter offloading strategies.

    Offloaders control how model parameters are stored and loaded during
    inference. Different strategies trade memory for compute/transfer time.
    """

    supports_tower_offload: bool = False
    """Whether `wrap_modules` also accepts modules routed by
    `SupportsMultiModal._mark_tower_model`, outside the `make_layers` call.

    Offloaders whose `wrap_modules` may only be called on the decoder layer
    stack (e.g. `PrefetchOffloader`, which schedules prefetches over a
    circular layer stack) must keep this `False`.
    """

    @abstractmethod
    def wrap_modules(
        self,
        modules_generator: Generator[nn.Module, None, None],
        prefix: str = "",
    ) -> list[nn.Module]:
        """Wrap modules with offloading logic.

        Args:
            modules_generator: Generator yielding modules to potentially offload.
            prefix: Name prefix prepended to parameter names before matching
                them against the offloading parameter set. Used when the
                modules are not the full model, so that name segments stay
                fully qualified (e.g. `visual` for a tower module).

        Returns:
            List of modules, potentially with offloading hooks installed.
        """
        pass

    def post_init(self):
        """Called after model construction completes.

        Offloaders can use this to:
        - Finalize parameter storage
        - Start initial prefetching
        - Allocate shared resources
        """
        return

    def sync_prev_onload(self) -> None:  # noqa: B027
        """Sync previous onload operations. Override in subclasses."""
        pass

    def join_after_forward(self) -> None:  # noqa: B027
        """Join streams after forward. Override in subclasses."""
        pass

    def _wait_for_layer(self, layer_idx: int) -> None:  # noqa: B027
        """Wait for layer prefetch. Override in subclasses."""
        pass

    def _start_prefetch(self, layer_idx: int) -> None:  # noqa: B027
        """Start layer prefetch. Override in subclasses."""
        pass

supports_tower_offload = False class-attribute instance-attribute

Whether wrap_modules also accepts modules routed by SupportsMultiModal._mark_tower_model, outside the make_layers call.

Offloaders whose wrap_modules may only be called on the decoder layer stack (e.g. PrefetchOffloader, which schedules prefetches over a circular layer stack) must keep this False.

_start_prefetch(layer_idx)

Start layer prefetch. Override in subclasses.

Source code in vllm/model_executor/offloader/base.py 
105
106
107
def _start_prefetch(self, layer_idx: int) -> None:  # noqa: B027
    """Start layer prefetch. Override in subclasses."""
    pass

_wait_for_layer(layer_idx)

Wait for layer prefetch. Override in subclasses.

Source code in vllm/model_executor/offloader/base.py 
101
102
103
def _wait_for_layer(self, layer_idx: int) -> None:  # noqa: B027
    """Wait for layer prefetch. Override in subclasses."""
    pass

join_after_forward()

Join streams after forward. Override in subclasses.

Source code in vllm/model_executor/offloader/base.py 
97
98
99
def join_after_forward(self) -> None:  # noqa: B027
    """Join streams after forward. Override in subclasses."""
    pass

post_init()

Called after model construction completes.

Offloaders can use this to: - Finalize parameter storage - Start initial prefetching - Allocate shared resources

Source code in vllm/model_executor/offloader/base.py 
83
84
85
86
87
88
89
90
91
def post_init(self):
    """Called after model construction completes.

    Offloaders can use this to:
    - Finalize parameter storage
    - Start initial prefetching
    - Allocate shared resources
    """
    return

sync_prev_onload()

Sync previous onload operations. Override in subclasses.

Source code in vllm/model_executor/offloader/base.py 
93
94
95
def sync_prev_onload(self) -> None:  # noqa: B027
    """Sync previous onload operations. Override in subclasses."""
    pass

wrap_modules(modules_generator, prefix='') abstractmethod

Wrap modules with offloading logic.

Parameters:

  • modules_generator

    (Generator[Module, None, None]) –

    Generator yielding modules to potentially offload.

  • prefix

    (str, default: '' ) –

    Name prefix prepended to parameter names before matching them against the offloading parameter set. Used when the modules are not the full model, so that name segments stay fully qualified (e.g. visual for a tower module).

Returns:

  • list[Module] –

    List of modules, potentially with offloading hooks installed.

Source code in vllm/model_executor/offloader/base.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
@abstractmethod
def wrap_modules(
    self,
    modules_generator: Generator[nn.Module, None, None],
    prefix: str = "",
) -> list[nn.Module]:
    """Wrap modules with offloading logic.

    Args:
        modules_generator: Generator yielding modules to potentially offload.
        prefix: Name prefix prepended to parameter names before matching
            them against the offloading parameter set. Used when the
            modules are not the full model, so that name segments stay
            fully qualified (e.g. `visual` for a tower module).

    Returns:
        List of modules, potentially with offloading hooks installed.
    """
    pass

NoopOffloader

Bases: BaseOffloader

No-op offloader that returns modules as-is without any offloading.

Methods:

Source code in vllm/model_executor/offloader/base.py 
110
111
112
113
114
115
116
117
118
119
class NoopOffloader(BaseOffloader):
    """No-op offloader that returns modules as-is without any offloading."""

    def wrap_modules(
        self,
        modules_generator: Generator[nn.Module, None, None],
        prefix: str = "",
    ) -> list[nn.Module]:
        """Return modules unchanged."""
        return list(modules_generator)

wrap_modules(modules_generator, prefix='')

Return modules unchanged.

Source code in vllm/model_executor/offloader/base.py 
113
114
115
116
117
118
119
def wrap_modules(
    self,
    modules_generator: Generator[nn.Module, None, None],
    prefix: str = "",
) -> list[nn.Module]:
    """Return modules unchanged."""
    return list(modules_generator)

create_offloader(offload_config)

Create an offloader based on the offload configuration.

Uses the explicit offload_backend selector. When set to "auto", selects prefetch if offload_group_size > 0, UVA if cpu_offload_gb > 0, otherwise noop.

Source code in vllm/model_executor/offloader/base.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
    """Create an offloader based on the offload configuration.

    Uses the explicit ``offload_backend`` selector.  When set to ``"auto"``,
    selects prefetch if ``offload_group_size > 0``, UVA if
    ``cpu_offload_gb > 0``, otherwise noop.
    """
    from vllm.model_executor.offloader.prefetch import PrefetchOffloader
    from vllm.model_executor.offloader.uva import UVAOffloader

    backend = offload_config.offload_backend
    uva = offload_config.uva
    prefetch = offload_config.prefetch

    if backend == "auto":
        if prefetch.offload_group_size > 0:
            backend = "prefetch"
        elif uva.cpu_offload_gb > 0:
            backend = "uva"
        else:
            return NoopOffloader()

    if backend == "prefetch":
        return PrefetchOffloader(
            group_size=prefetch.offload_group_size,
            num_in_group=prefetch.offload_num_in_group,
            prefetch_step=prefetch.offload_prefetch_step,
            offload_params=prefetch.offload_params,
            mode="cpu",
        )
    elif backend == "uva":
        return UVAOffloader(
            cpu_offload_max_bytes=int(uva.cpu_offload_gb * 1024**3),
            cpu_offload_params=uva.cpu_offload_params,
        )
    else:
        return NoopOffloader()

get_offloader()

Get the global offloader instance.

Source code in vllm/model_executor/offloader/base.py 
126
127
128
def get_offloader() -> BaseOffloader:
    """Get the global offloader instance."""
    return _instance

set_offloader(instance)

Set the global offloader instance.

Source code in vllm/model_executor/offloader/base.py 
131
132
133
134
135
136
137
138
def set_offloader(instance: BaseOffloader) -> None:
    """Set the global offloader instance."""
    global _instance
    _instance = instance
    if isinstance(instance, NoopOffloader):
        logger.debug_once("Offloader set to NoopOffloader (no offloading).")
    else:
        logger.info_once("Offloader set to %s", type(instance).__name__)

should_pin_memory()

Check if pinned memory should be used for weight offloading.

Combines the platform capability check with the user override env var. On unified-memory systems (e.g. GH200) pinned memory eats into GPU memory, so users can disable it via VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY.

Source code in vllm/model_executor/offloader/base.py 
23
24
25
26
27
28
29
30
31
32
def should_pin_memory() -> bool:
    """Check if pinned memory should be used for weight offloading.

    Combines the platform capability check with the user override env var.
    On unified-memory systems (e.g. GH200) pinned memory eats into GPU
    memory, so users can disable it via VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY.
    """
    return (
        is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
    )