Skip to content

[train] Support LoRA-only weight syncing for Megatron backend #1336

@SumanthRH

Description

@SumanthRH

Summary

Currently, for weight syncing with LoRA, SkyRL's Megatron backend will fuse the weights and sync the merged weights with vLLM replicas. SkyRL should support syncing just the LoRA weights. and enabling LoRA serving on the vLLM side.

This was previously attempted in #885 but it also requires vLLM changes.

  1. Change Megatron's weight extractor to only transfer trainable, LoRA weights (this should be doable with export_adapter_weights (link) in Megatron-Bridge
  2. Fix vLLM's weight loading in LRUCacheWorkerLoRAManager upstream to allow LoRA weight loading with in-memory LoRA weights
  3. Upgrade SkyRL to use the vLLM release version with the above changes after 2. lands

Reference PRs:
verl: verl-project/verl#4632
Nemo-RL: NVIDIA-NeMo/RL#1797 (three PRs linked and some relevant benchmarking for LoRA only weight update code path there)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions