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

[bvolpato]

Summary

Add opt-in bitsandbytes 4-bit loading for FSDP policy and reference models 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 policy and reference models 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.
• Accept explicit LoRA target and exclusion module lists as documented.
• 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 fsdp --extra dev pytest -q tests/train tests/backends/skyrl_train --ignore=tests/backends/skyrl_train/gpu -m 'not vllm' (1402 passed, 5 skipped, 5 deselected after rebasing onto current main)
• pre-commit run --all-files --config .pre-commit-config.yaml (Ruff, Black, secret scan passed)
• Real Gemma 4 E2B reference-worker probe loaded 4-bit modules and completed a finite forward on a 16 GB RTX 5070 Ti.
• Combined with [chore][deps] Support GLM-5.2 configs #1974, a fresh 10-step Gemma 4 E2B SkyRL QLoRA SFT run reduced loss monotonically from 17.6720 to 0.0127; every gradient was finite.
• Combined with [chore][deps] Support GLM-5.2 configs #1974, a fresh 10-step GLM-5.2 LoRA SkyRL SFT run using an explicit q_a_proj target reduced loss monotonically from 12.4953 to 11.3287; every gradient was finite.
• Combined with [chore][deps] Support GLM-5.2 configs #1974, one-step inference-optimization/GLM-5.2-0.8B-A0.8B GRPO QLoRA on a 16 GB RTX 5070 Ti completed 4-bit policy load, vLLM rollout, reward, advantages, FSDP backward, optimizer, and post-update LoRA sync. Absolute advantage was 0.7071, gradient norm was 0.0449, and process exited 0. Rollout used a test-only dense MLA fallback because native sparse MLA rejects this tiny checkpoint's 192-wide compressed KV head on SM120.
• 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 PR. New behavior is disabled by default.

• Combined with [chore][deps] Support GLM-5.2 configs #1974, a 12-step exact-model GRPO QLoRA run completed on one 16 GB RTX 5070 Ti. Every step had nonzero GRPO advantages and gradients, then synchronized updated adapters back to vLLM. Mean toy reward increased from 0.6938 over steps 1-4 to 0.7147 over steps 9-12, with expected sampling noise. Dense MLA required prefix caching disabled because vLLM v0.23 compressed-cache gather kernels do not accept this model’s 192-wide head.

[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.

[gemini-code-assist]

Code Review

This pull request introduces support for FSDP QLoRA (4-bit base-weight quantization with LoRA adapters) using bitsandbytes, including configuration options, validation rules, and documentation/examples (specifically for Gemma 4 E2B). It updates model loading and FSDP sharding logic to handle 4-bit quantized weights, and fixes embedding wrapping behavior in FSDP. The reviewer noted that while FSDPPolicyWorkerBase supports 4-bit quantization, FSDPRefWorkerBase in fsdp_worker.py does not yet pass these parameters to HFModelWrapper, which could cause the reference model to load in full precision during RL training and lead to high memory usage or OOM.