Add PiSSA initialization for Megatron LoRA adapters

[kalectory]

Summary

Following Meng et al., this adds PiSSA initialization for Megatron LoRA adapters.

• Initialize dense and per-expert grouped-MoE adapters from principal singular components across TP, EP, and ETP layouts, while preserving the original model as a residual base plus adapter.
• Add an offline producer that exports the residual Hugging Face model and a standard step-zero adapter/optimizer checkpoint for normal training resume.
• Add matched Qwen3-0.6B and Qwen3-30B-A3B producer and training examples based on the existing Megatron LoRA recipes.

Results

Initialization numerical error

PiSSA initialization preserves the base model at approximately the same log-probability error scale as standard LoRA initialization.

Qwen3-0.6B training

The included Qwen3-0.6B PiSSA recipe was trained against the corresponding SkyRL Megatron LoRA example.

Follow-ups

• Evaluate shared-expert MoE PiSSA as a separate algorithm variant.
• Support LoRA/PiSSA targeting for Qwen3.5/3.6 GDN projections.
• Add an FSDP PiSSA artifact producer.
• Evaluate FP32 adapter storage to reduce initialization error.
• Add an optional rank-zero CPU SVD fallback if supported large models exhibit CUDA OOM or cuSOLVER failures.

Testing

uv run --extra dev --extra megatron pytest tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_pissa_init_worker.py -k 'not preserves_base_model_forward' -q
uvx --from ruff==0.11.9 ruff check skyrl/backends/skyrl_train/workers/megatron/pissa_init_worker.py skyrl/train/entrypoints/pissa_init.py tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_pissa_init_worker.py

The focused suite passes 20 tests, with the two multi-GPU identity cases deselected. A real two-process collective smoke test also passes. SkyRL's pinned Black check, launcher shell syntax checks, pre-commit workflow, and SkyRL Gym CI pass.

The test suite includes dense TP=2/PP=2 and grouped-MoE TP=4/EP=8/ETP=1 identity cases. The upstream GPU workflow could not start for this fork PR because the Anyscale credential is unavailable to fork workflows; it failed before allocating hardware or running tests.

[gemini-code-assist]

Code Review

This pull request introduces support for offline PiSSA (Principal Singular values and Singular vectors Adaptation) initialization with Megatron, including launcher scripts, documentation, and tests for both dense and grouped-expert MoE models. The feedback focuses on enhancing the robustness and safety of the implementation: running the memory-intensive SVD decomposition on the CPU instead of the GPU to prevent OOM errors on large models, avoiding the discouraged use of .data in PyTorch in favor of .detach(), and passing an integer instead of a float for the training steps parameter.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 47b983c. Configure here.}