Skip to content

[Feature]: Support fast weight updates from disk #48644

Description

@SumanthRH

🚀 The feature, motivation and pitch

vLLM currently uses layerwise reloading while updating weights during RL training without recompilation:
https://docs.vllm.ai/en/latest/training/layerwise/#what-is-layerwise-reloading

This currently assumes that weights for a given layer are loaded in order because the layerwise reloading logic allocates buffers per layer:

https://docs.vllm.ai/en/latest/training/layerwise/#troubleshooting-excessive-memory-usage

There can be excessive memory usage when weights are received out of order.

In some use-cases like disaggregated RL, it is desirable to update weights from disk. In this case, it would be good to leverage GPU Direct Storage to directly load weights into GPU memory.

Users would like to do something like:

# in gpu_worker.py
from vllm.model_executor.model_loader.weight_utils import fastsafetensors_weights_iterator

self.model.load_weights(fastsafetensors_weights_iterator(files, use_tqdm_on_load=False))

However, the current fastsafetensors iterator will lead to large temporary memory allocation due to out of order loading.

For Qwen 3.5 35B, I saw temporary memory grow to > 40 GB due to the out of order loading issue, with warnings like:

WARNING 07-14 08:53:56 [layerwise.py:199] Allocating 42305.8 MB of device memory to buffers to load ['FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE'] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.
WARNING 07-14 08:53:56 [layerwise.py:199] Allocating 42307.9 MB of device memory to buffers to load ['FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE'] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.
WARNING 07-14 08:53:56 [layerwise.py:199] Allocating 42310.0 MB of device memory to buffers to load ['FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE'] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.
WARNING 07-14 08:53:56 [layerwise.py:199] Allocating 42312.1 MB of device memory to buffers to load ['FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE', 'FusedMoE'] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.

I see two options:

  1. Provide a different safetensors iterator that uses GDS but is compatible with layerwise reloading, bounding the temporary memory usage in some way
  2. Provide a separate path for updating weights from disk that bypasses the current layerwise reloading and can still use the existing fastsafetensors iterator

Alternatives

No response

Additional context

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions