This currently assumes that weights for a given layer are loaded in order because the layerwise reloading logic allocates buffers per layer:
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.
# 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))
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.
🚀 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:
However, the current
fastsafetensorsiterator 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:
I see two options:
Alternatives
No response
Additional context
No response
Before submitting a new issue...