Skip to content

Add dense Megatron LoRA-XS support - #2

Draft
kalectory wants to merge 7 commits into
neil/pissafrom
neil/lora-xs
Draft

Add dense Megatron LoRA-XS support#2
kalectory wants to merge 7 commits into
neil/pissafrom
neil/lora-xs

Conversation

@kalectory

@kalectory kalectory commented Aug 19, 2026

Copy link
Copy Markdown
Owner
  • Adds dense Megatron LoRA-XS with frozen SVD factors and a trainable rank-by-rank core, building on the PiSSA distributed initialization path.
  • Keeps algorithm behavior in lora_xs.py; the shared worker only selects the PEFT class and invokes its optional export context.
  • Adds offline step-zero initialization, Qwen3-0.6B GSM8K example scripts, and focused coverage for factorization, TP sharding, gradient isolation, fused-forward handling, and export restoration. Stacked on Add PiSSA initialization for Megatron LoRA adapters NovaSky-AI/SkyRL#2069.

Testing

NVTE_FLASH_ATTN=0 uv run --isolated --extra dev --extra megatron pytest tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_lora_xs.py -q

9 passed; verifies the focused LoRA-XS initialization and export behavior.

uv run --isolated --with pre-commit pre-commit run --files skyrl/backends/skyrl_train/workers/megatron/lora_xs.py skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_lora_xs.py

Formatting, lint, and secret checks passed.

On 8x H200 with Qwen3-0.6B/GSM8K, offline initialization plus one 16-trajectory GRPO batch completed policy training, saved global_step_1, and synchronized the updated weights to vLLM.

@kalectory
kalectory force-pushed the neil/lora-xs branch 2 times, most recently from a71bf49 to 884ed23 Compare August 19, 2026 22:10
Signed-off-by: Neil Kale <263453039+kalectory@users.noreply.github.com>
Signed-off-by: Neil Kale <263453039+kalectory@users.noreply.github.com>
Signed-off-by: Neil Kale <263453039+kalectory@users.noreply.github.com>
Signed-off-by: Neil Kale <263453039+kalectory@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant