[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 single-GPU example.

Changes

• Add model.bitsandbytes_4bit with NF4/FP4 and double-quantization controls.
• Load policy and reference models with BF16 compute and quantization storage.
• Keep bitsandbytes Params4bit off the meta-device path and avoid redundant single-rank FSDP wrapping.
• Reject FSDP-native CPU offload for unwrapped single-rank 4-bit models and direct users to manual phase offload.
• Preserve packed 4-bit payloads during multi-rank FSDP2 all-gathers by disabling FSDP parameter casting while retaining BF16 compute and FP32 gradient reduction.
• Prevent embeddings inside selected parent modules from being wrapped twice.
• Restrict language_model_only LoRA adapters to the language model.
• Support explicit LoRA target and exclusion lists.
• Expose SFT FlashAttention selection, validate its padding invariant, and add a Gemma-safe SDPA example.
• Add the qlora extra, configuration tests, wrap-policy coverage, docs, and example script.

New behavior is disabled by default.

Repro

After accepting the Gemma license on Hugging Face:

bash examples/train/sft/run_sft_gemma4_e2b_qlora.sh

Validation

• Current focused config, SFT, and FSDP tests (212 passed)
• Two-rank FSDP2 QLoRA byte-preservation probe on one RTX 5070 Ti: 14 quantized modules, 23 BF16 NaN storage values, zero packed-byte mismatches, finite loss and gradients, and successful optimizer step
• Full pre-commit and uv lock --check
• Fresh 4-bit Gemma 4 E2B load on RTX 5070 Ti: 525 quantized parameters, finite BF16 logits, shape (1, 2, 262144)
• Prior SkyRL Train CPU suite (1402 passed, 5 skipped, 5 deselected)
• Real 4-bit reference-worker load and finite forward
• Gemma 4 E2B 100-step QLoRA run: held-out loss 4.2274 to 1.1558; finite gradients throughout
• Exported 38 MB PEFT adapter reloaded over the same 4-bit base with finite logits
• Gemma 4 E4B 100-step run: held-out loss 2.7891 to 1.1067
• Combined with [chore][deps] Support GLM-5.2 configs #1974, exact-model GLM-5.2 GRPO QLoRA completed rollout, backward, optimizer updates, and adapter sync with finite gradients

Single-rank and colocated RL paths were validated on a 16 GB RTX 5070 Ti. Multi-rank FSDP2 was validated with two Gloo ranks sharing that GPU; separate-GPU NCCL remains untested locally.

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

[cursor]

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

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