Summary
Recently, we merged #2062 , which changes SkyRL's offload/backload semantics with LoRA + colocate_all=True as follows: LoRA parameters of the trainer stay resident on the GPU throughout training with colocate_all=True.
There are two follow-ups to #2062:
- We should document the new way of offload/backload with LoRA + colocated training temporarily
- LoRA weights stay in GPU all the time, even after weight sync is done and we're generating with the vLLM engines. We could have OOMs with high vLLM GPU util ratios that users didn't encounter previously.
- Revisit offload/backload semantics for LoRA: SkyRL previously backloaded all parameters, including frozen base weights before a weight sync:
|
def _prepare_for_weight_sync(self) -> None: |
|
"""Load policy weights and apply the configured optimizer offload policy.""" |
|
if not self.colocate_all: |
|
return |
|
offload_optimizer = self.cfg.trainer.policy.optimizer_config.offload_after_step |
|
self._ensure_on_gpu( |
|
"policy", |
|
need_optimizer=False, |
|
need_model=True, |
|
) |
We should actually book-keep trainable vs non-trainable weights differently and only backload trainable weights for weight sync with colocate_all=True. After weight sync, all weights should be offloaded. This will help avoid the OOM issue in 1. as well.
Summary
Recently, we merged #2062 , which changes SkyRL's offload/backload semantics with LoRA +
colocate_all=Trueas follows: LoRA parameters of the trainer stay resident on the GPU throughout training withcolocate_all=True.There are two follow-ups to #2062:
SkyRL/skyrl/backends/skyrl_train/workers/worker_dispatch.py
Lines 681 to 690 in 7207e4d
We should actually book-keep trainable vs non-trainable weights differently and only backload trainable weights for weight sync with
colocate_all=True. After weight sync, all weights should be offloaded. This will help avoid the OOM issue in 1. as well.