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:

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
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.
    """

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

        Args:
            modules_generator: Generator yielding modules to potentially offload.

        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

_start_prefetch(layer_idx)

Start layer prefetch. Override in subclasses.

Source code in vllm/model_executor/offloader/base.py 
91
92
93
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 
87
88
89
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 
83
84
85
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 
69
70
71
72
73
74
75
76
77
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 
79
80
81
def sync_prev_onload(self) -> None:  # noqa: B027
    """Sync previous onload operations. Override in subclasses."""
    pass

wrap_modules(modules_generator) abstractmethod

Wrap modules with offloading logic.

Parameters:

  • modules_generator

    (Generator[Module, None, None]) –

    Generator yielding modules to potentially offload.

Returns:

  • list[Module] –

    List of modules, potentially with offloading hooks installed.

Source code in vllm/model_executor/offloader/base.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
@abstractmethod
def wrap_modules(
    self,
    modules_generator: Generator[nn.Module, None, None],
) -> list[nn.Module]:
    """Wrap modules with offloading logic.

    Args:
        modules_generator: Generator yielding modules to potentially offload.

    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 
 96
 97
 98
 99
100
101
102
103
104
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],
    ) -> list[nn.Module]:
        """Return modules unchanged."""
        return list(modules_generator)

wrap_modules(modules_generator)

Return modules unchanged.

Source code in vllm/model_executor/offloader/base.py 
 99
100
101
102
103
104
def wrap_modules(
    self,
    modules_generator: Generator[nn.Module, None, None],
) -> 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 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
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 
111
112
113
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 
116
117
118
119
120
121
122
123
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
    )