[feat][fsdp] Add QLoRA and Gemma 4 training support

[bvolpato]

Summary

Add opt-in bitsandbytes 4-bit loading for FSDP policy training so SkyRL can train LoRA adapters against a quantized base model in SFT and RL. Include a runnable Gemma 4 E2B example for small single-GPU experiments.

Changes

• Add model.bitsandbytes_4bit config with NF4/FP4 and double-quantization controls.
• Load 4-bit policies with BF16 compute and BF16 quantization storage for FSDP compatibility.
• Avoid moving bitsandbytes Params4bit through the meta device.
• Skip redundant FSDP wrapping for single-rank 4-bit runs to avoid all-gather memory duplication.
• Keep embeddings inside explicitly selected parent wrap boundaries instead of wrapping them twice.
• Restrict language_model_only LoRA adapters to the language model so text-only vLLM sync excludes audio and vision towers.
• Expose SFT flash_attn selection and use SDPA plus a Gemma-safe FSDP boundary in the example.
• Add qlora dependency extra, config tests, wrap-policy regression coverage, and SFT docs.

Repro

After accepting the Gemma license on Hugging Face:

bash examples/train/sft/run_sft_gemma4_e2b_qlora.sh

Extra config arguments can reduce this to a one-step smoke:

bash examples/train/sft/run_sft_gemma4_e2b_qlora.sh \
  'train_dataset_splits=["train[:1]"]' \
  max_length=64 \
  num_steps=1 \
  disable_cache=true \
  num_workers=0

Testing

• uv run --isolated --extra dev --extra skyrl-train pytest tests/train/test_sft_config.py tests/backends/skyrl_train/distributed/test_fsdp_utils.py -q (62 passed)
• uv run --isolated --extra dev pre-commit run --all-files (Ruff, Black, gitleaks passed)
• uv run --isolated --extra fsdp --extra qlora --extra dev pytest -q tests/backends/skyrl_train/models/test_models.py (4 passed)
• Gemma 4 E2B meta-init checks covered all-linear and explicit LoRA targets: 276 and 50 language-model adapters respectively, with zero audio or vision tower adapters.
• One-step colocated Gemma 4 E2B GRPO QLoRA smoke on a 16 GB RTX 5070 Ti used a deterministic positive-reward validation environment and completed vLLM 4-bit rollout, policy logprobs, optimizer update, and initial/post-update LoRA sync. Metrics were reward 1.0, advantage 1.0, policy loss -1.0, and finite grad norm 6.78125; process exited 0.
• 100-step sequential QLoRA run using the compatible public unsloth/gemma-4-E2B-it mirror on a 16 GB RTX 5070 Ti: 200 distinct training examples, plus a disjoint 31-example held-out set. Held-out loss fell from 4.2274 before training to 1.1885 at step 50 and 1.1558 at step 100; gradients stayed finite and all optimizer steps completed.
• HF export produced a 38 MB PEFT adapter. Reloading it over the same 4-bit base produced finite logits with shape (1, 4, 262144).
• Gemma 4 E4B QLoRA also fits a 16 GB RTX 5070 Ti. A 20-step one-example overfit reduced loss from 2.7434 to 0.8364 with finite gradients.
• A separate E4B run trained on 200 distinct examples and evaluated 31 disjoint examples: held-out loss fell from 2.7891 before training to 1.1250 at step 50 and 1.1067 at step 100.

Did this cause any problems?

Rollback by reverting this commit. New behavior is disabled by default.

[gemini-code-assist]

Code Review

This pull request introduces support for FSDP QLoRA by integrating 4-bit base-weight quantization using bitsandbytes. It adds the necessary configuration options, updates model loading and wrapping logic, handles in-place sharding for 4-bit parameters, and ensures embeddings within wrapped parent modules are not redundantly wrapped. Feedback on the changes includes addressing a regex bug in LoRA exclusions when exclude_modules is empty, optimizing module traversal in apply_fsdp2 to a single pass, and replacing configuration assertions with explicit ValueError exceptions to prevent issues if Python is run with optimization flags.