Commit 7c3542d
Add PiSSA initialization for Megatron LoRA adapters
PiSSA (https://arxiv.org/abs/2404.02948) seeds LoRA A/B from the principal
singular components of the base weight and freezes the residual W_res = W - AB
in place of the full W, so the adapter starts on the principal subspace instead
of at zero.
- pissa_init.py: pure, CPU-testable pissa_decompose() + TP-aware apply_pissa_init()
that gathers each base weight across the TP group, SVD-decomposes, and writes
the residual base + A/B back in megatron-bridge's sharded layout.
- SkyRLLoraConfig.pissa flag; applied in the megatron worker's lora pre-wrap hook
(after HF weights load, before register_pristine), so no megatron-bridge edits.
- CPU unit tests for the decomposition math (reconstruction, principal
correctness, alpha/rank scale invariance).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>1 parent db73ad6 commit 7c3542d
4 files changed
Lines changed: 534 additions & 85 deletions
File tree
- skyrl
- backends/skyrl_train/workers/megatron
- train/config
- tests/backends/skyrl_train/workers
0 commit comments