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

docs/content/docs/sft/overview.mdx

@@ -32,6 +32,21 @@ bash examples/train/sft/run_sft_fsdp.sh
 Trains `Qwen/Qwen2.5-0.5B-Instruct` on 1 GPU with FSDP (max length 512, batch
 size 4, 10 steps).
 
+### QLoRA with Gemma 4 E2B
+
+```bash
+bash examples/train/sft/run_sft_gemma4_e2b_qlora.sh
+```
+
+This example loads `google/gemma-4-E2B-it` with bitsandbytes NF4 base weights
+and trains LoRA adapters with BF16 computation. Accept the Gemma license on
+Hugging Face before running it. The `qlora` extra installs bitsandbytes.
+
+QLoRA is available only with FSDP and requires `model.lora.rank > 0`. Configure
+it through `model.bitsandbytes_4bit`; `quant_type` accepts `nf4` or `fp4`, and
+`use_double_quant` controls nested quantization. Gemma 4 uses SDPA because its
+attention layout is incompatible with this FlashAttention 2 path.
+
 ### Megatron (multi-GPU with TP/PP)
 
 ```bash

examples/train/sft/run_sft_gemma4_e2b_qlora.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+set -x
+
+# Text-only SFT of Gemma 4 E2B with a 4-bit base model and LoRA adapters.
+# Accept the Gemma license on Hugging Face before running the Google checkpoint.
+
+uv run --isolated --extra fsdp --extra qlora \
+    python -m skyrl.train.main_sft \
+    strategy=fsdp \
+    model.path=google/gemma-4-E2B-it \
+    model.bitsandbytes_4bit.enabled=true \
+    model.lora.rank=8 \
+    model.lora.alpha=16 \
+    model.lora.target_modules=all-linear \
+    train_datasets="['yahma/alpaca-cleaned']" \
+    train_dataset_splits="['train[:100]']" \
+    messages_key=messages \
+    max_length=128 \
+    num_steps=10 \
+    batch_size=1 \
+    micro_train_batch_size_per_gpu=1 \
+    remove_microbatch_padding=false \
+    flash_attn=false \
+    seed=42 \
+    optimizer_config.lr=2e-5 \
+    optimizer_config.weight_decay=0.0 \
+    optimizer_config.max_grad_norm=1.0 \
+    optimizer_config.num_warmup_steps=0 \
+    optimizer_config.scheduler=constant_with_warmup \
+    placement.num_nodes=1 \
+    placement.num_gpus_per_node=1 \
+    fsdp_config.cpu_offload=false \
+    fsdp_config.reshard_after_forward=true \
+    fsdp_config.wrap_policy.transformer_layer_cls_to_wrap="['Gemma4Model']" \
+    logger=console \
+    project_name=skyrl_sft_qlora \
+    run_name=gemma4_e2b_qlora \
+    ckpt_path="" \
+    ckpt_interval=0 \
+    resume_from="" \
+    "$@"

pyproject.toml

@@ -127,6 +127,10 @@ fsdp = [
     "torchvision; sys_platform == 'linux'",
 ]
 
+qlora = [
+    "bitsandbytes>=0.47.0; sys_platform == 'linux'",
+]
+
 megatron = [
     "skyrl[skyrl-train]",
     "transformer-engine[pytorch]==2.11.0; sys_platform == 'linux'",

skyrl/backends/skyrl_train/distributed/fsdp_strategy.py

@@ -92,6 +92,7 @@ def __init__(
 
         # LoRA related configs
         self.is_lora = self.model_config.lora.rank > 0 if self.model_config is not None else False
+        self.is_4bit = self.model_config.bitsandbytes_4bit.enabled if self.model_config is not None else False
 
         self.time_steps = defaultdict(int)
 
@@ -221,6 +222,16 @@ def _fsdp_init_model(self, model, is_train=True, is_wrapped=False):
             "reshard_after_forward": self.fsdp_config.reshard_after_forward,
         }
         module = model.model if is_wrapped else model
+
+        if self.is_4bit and self.world_size == 1:
+            return module
+
+        # Params4bit cannot be moved through the meta device. Each rank already
+        # loaded the same quantized checkpoint, so shard it in place instead.
+        if self.is_4bit:
+            apply_fsdp2(module, fsdp_kwargs, self.fsdp_config)
+            return module
+
         full_state = module.state_dict()
 
         # Move the entire module to meta before apply_fsdp2 so the sharded

skyrl/backends/skyrl_train/distributed/fsdp_utils.py

@@ -224,11 +224,20 @@ def apply_fsdp2(model, fsdp_kwargs, config: Union[FSDPConfig, DictConfig]):
     assert len(fsdp_transformer_layer_cls_to_wrap) > 0 and fsdp_transformer_layer_cls_to_wrap[0] is not None
 
     modules = []
+    wrapped_module_names = []
+    embeddings = []
+    tie_word_embeddings = getattr(getattr(model, "config", None), "tie_word_embeddings", False)
     for name, module in model.named_modules():
-        if module.__class__.__name__ in fsdp_transformer_layer_cls_to_wrap or (
-            isinstance(module, nn.Embedding) and not model.config.tie_word_embeddings
-        ):
+        if module.__class__.__name__ in fsdp_transformer_layer_cls_to_wrap:
             modules.append(module)
+            wrapped_module_names.append(name)
+        elif isinstance(module, nn.Embedding) and not tie_word_embeddings:
+            embeddings.append((name, module))
+
+    for name, module in embeddings:
+        if any(name.startswith(f"{parent_name}.") for parent_name in wrapped_module_names if parent_name):
+            continue
+        modules.append(module)
 
     for idx, module in enumerate(modules):
         fully_shard(module, **fsdp_kwargs)

skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py

@@ -147,18 +147,24 @@ def init_model(self, model_path, num_training_steps: int = None):
         self.strategy = strategy
 
         self._is_lora = self.cfg.policy.model.lora.rank > 0
+        quantization = self.cfg.policy.model.bitsandbytes_4bit
+        if quantization.enabled and not self._is_lora:
+            raise ValueError("4-bit policy training requires LoRA (policy.model.lora.rank > 0).")
 
         model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
         is_multimodal = hasattr(model_config, "vision_config") and model_config.vision_config is not None
         self._is_multimodal_lm_only = self.cfg.policy.language_model_only and is_multimodal
-        use_meta = should_use_meta_init(
+        use_meta = not quantization.enabled and should_use_meta_init(
             use_meta_tensor=not model_config.tie_word_embeddings, mesh=self.strategy.device_mesh
         )
 
         wrapped_model = HFModelWrapper(
             model_path,
             use_flash_attention_2=self.cfg.flash_attn,
-            bf16=self.cfg.policy.inference_only_init,
+            bf16=self.cfg.policy.inference_only_init or quantization.enabled,
+            load_in_4bit=quantization.enabled,
+            bnb_4bit_quant_type=quantization.quant_type,
+            bnb_4bit_use_double_quant=quantization.use_double_quant,
             lora_rank=self.cfg.policy.model.lora.rank,
             lora_alpha=self.cfg.policy.model.lora.alpha,
             lora_dropout=self.cfg.policy.model.lora.dropout,
@@ -417,22 +423,28 @@ def init_model(self, model_path):
         assert self.cfg.strategy == "fsdp"
         strategy = FSDPStrategy(
             fsdp_config=self.cfg.ref.fsdp_config,
+            model_config=self.cfg.ref.model,
             fsdp_strategy=self.cfg.strategy,
             seed=self.cfg.seed,
             micro_train_batch_size_per_gpu=self.cfg.micro_train_batch_size_per_gpu,
         )
         strategy.setup_distributed()
         self.strategy = strategy
 
+        quantization = self.cfg.ref.model.bitsandbytes_4bit
+
         model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-        use_meta = should_use_meta_init(
+        use_meta = not quantization.enabled and should_use_meta_init(
             use_meta_tensor=not model_config.tie_word_embeddings, mesh=self.strategy.device_mesh
         )
 
         wrapped_model = HFModelWrapper(
             model_path,
             use_flash_attention_2=self.cfg.flash_attn,
-            bf16=self.cfg.bf16,
+            bf16=self.cfg.bf16 or quantization.enabled,
+            load_in_4bit=quantization.enabled,
+            bnb_4bit_quant_type=quantization.quant_type,
+            bnb_4bit_use_double_quant=quantization.use_double_quant,
             sequence_parallel_size=self.cfg.ref.sequence_parallel_size,
             remove_microbatch_padding=self.cfg.remove_microbatch_padding,
             model_config_kwargs=self.cfg.ref.model_config_kwargs,

skyrl/backends/skyrl_train/workers/model_wrapper.py

@@ -3,6 +3,7 @@
 # https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/actor.py
 # https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/model.py
 
+import re
 from typing import Optional, Union
 
 import numpy as np
@@ -33,6 +34,24 @@
 )
 
 
+def _language_model_only_lora_exclusions(model: nn.Module, exclude_modules):
+    prefixes = [
+        name for name, _ in model.named_modules() if name == "language_model" or name.endswith(".language_model")
+    ]
+    if not prefixes:
+        return exclude_modules
+
+    prefix = min(prefixes, key=len)
+    outside_language_model = rf"^(?!{re.escape(prefix)}(?:\.|$)).*$"
+    if not exclude_modules:
+        return outside_language_model
+    if isinstance(exclude_modules, str):
+        return rf"(?:{outside_language_model})|(?:{exclude_modules})"
+
+    excluded_names = "|".join(rf"(?:.*\.)?{re.escape(name)}" for name in exclude_modules)
+    return rf"(?:{outside_language_model})|(?:{excluded_names})"
+
+
 class HFModelWrapper(nn.Module):
     """
     Base class for wrapped HF models in reinforcement learning.
@@ -44,6 +63,8 @@ class HFModelWrapper(nn.Module):
         use_flash_attention_2 (bool, optional): Whether to utilize Flash Attention 2.0 for improved performance. Defaults to False.
         bf16 (bool, optional): Enable bfloat16 precision for model computations. Defaults to True.
         load_in_4bit (bool, optional): Load the model in 4-bit precision. Defaults to False.
+        bnb_4bit_quant_type (str, optional): bitsandbytes 4-bit data type. Defaults to "nf4".
+        bnb_4bit_use_double_quant (bool, optional): Enable nested quantization. Defaults to True.
         lora_rank (int, optional): Rank for LoRA adaptation. Defaults to 0.
         lora_alpha (int, optional): Alpha parameter for LoRA. Defaults to 16.
         lora_dropout (float, optional): Dropout rate for LoRA layers. Defaults to 0.
@@ -62,6 +83,8 @@ def __init__(
         use_flash_attention_2=False,
         bf16=True,
         load_in_4bit=False,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
         # TODO(shu): combine all LoRA specific configs into one place?
         lora_rank=0,
         lora_alpha=16,
@@ -97,9 +120,10 @@ def __init__(
                 assert bf16, "we only support bnb_4bit_compute_dtype = bf16"
                 nf4_config = BitsAndBytesConfig(
                     load_in_4bit=True,
-                    bnb_4bit_quant_type="nf4",
-                    bnb_4bit_use_double_quant=True,
+                    bnb_4bit_quant_type=bnb_4bit_quant_type,
+                    bnb_4bit_use_double_quant=bnb_4bit_use_double_quant,
                     bnb_4bit_compute_dtype=torch.bfloat16,
+                    bnb_4bit_quant_storage=torch.bfloat16,
                 )
             else:
                 nf4_config = None
@@ -186,6 +210,8 @@ def __init__(
             if lora_rank > 0:
                 # https://github.com/huggingface/peft/issues/137
                 self.model.enable_input_require_grads()
+                if language_model_only:
+                    exclude_modules = _language_model_only_lora_exclusions(self.model, exclude_modules)
                 lora_config = LoraConfig(
                     task_type=TaskType.CAUSAL_LM,
                     r=lora_rank,

skyrl/train/config/__init__.py

@@ -1,6 +1,7 @@
 from skyrl.train.config.config import (
     AlgorithmConfig,
     BaseConfig,
+    BitsAndBytes4BitConfig,
     ChatTemplateConfig,
     CISPOConfig,
     ClipCovConfig,
@@ -81,6 +82,7 @@
     "DynamicSamplingConfig",
     "OffPolicyCorrectionConfig",
     "BaseConfig",
+    "BitsAndBytes4BitConfig",
     "MixedPrecisionConfig",
     "MegatronDDPConfig",
     "MegatronLoraConfig",

skyrl/train/config/config.py

@@ -107,11 +107,11 @@ class SkyRLLoraConfig(BaseConfig):
     lora_sync_path: str = "/tmp/skyrl_lora_sync"
     """Directory where LoRA adapter weights are saved and synchronized between the training and inference processes.
     Must be accessible to all workers in distributed setups."""
-    target_modules: str = "all-linear"
+    target_modules: Union[str, List[str]] = "all-linear"
     """Modules to apply LoRA to.
     ``"all-linear"`` targets every linear layer for FSDP/PEFT, and is remapped to a fixed module list
     on Megatron. A list of specific module names can be given instead."""
-    exclude_modules: Optional[str] = None
+    exclude_modules: Optional[Union[str, List[str]]] = None
     """Modules to exclude from LoRA."""
     init_method: str = "kaiming"
     """For FSDP, corresponds to ``init_lora_weights`` in PEFT.
@@ -128,6 +128,18 @@ class SkyRLLoraConfig(BaseConfig):
     >= ``max_loras`` if explicitly set."""
 
 
+@dataclass
+class BitsAndBytes4BitConfig(BaseConfig):
+    """bitsandbytes 4-bit base-weight quantization for FSDP QLoRA."""
+
+    enabled: bool = False
+    """Load base weights in 4-bit. Policy training requires LoRA when enabled."""
+    quant_type: Literal["fp4", "nf4"] = "nf4"
+    """4-bit data type. NF4 is recommended for normally distributed pretrained weights."""
+    use_double_quant: bool = True
+    """Quantize first-stage quantization constants to save more memory."""
+
+
 @dataclass
 class FakeInt4QatConfig(BaseConfig):
     """Fake-INT4 quantization-aware training for MoE experts (Megatron only).
@@ -167,6 +179,7 @@ class ModelConfig(BaseConfig):
     path: Optional[str] = None
     """HuggingFace model path (or local directory) for this model."""
     lora: SkyRLLoraConfig = field(default_factory=SkyRLLoraConfig)
+    bitsandbytes_4bit: BitsAndBytes4BitConfig = field(default_factory=BitsAndBytes4BitConfig)
     fake_int4_qat: FakeInt4QatConfig = field(default_factory=FakeInt4QatConfig)
 
     def __post_init__(self) -> None:
@@ -1520,6 +1533,21 @@ def __post_init__(self):
                 "sync preserves the inference engine's INT4 base weights."
             )
 
+        if self.policy.model.bitsandbytes_4bit.enabled:
+            if self.strategy != "fsdp":
+                raise ValueError(
+                    "`trainer.policy.model.bitsandbytes_4bit.enabled=True` is only supported with "
+                    "`trainer.strategy=fsdp`."
+                )
+            if self.policy.model.lora.rank <= 0:
+                raise ValueError(
+                    "`trainer.policy.model.bitsandbytes_4bit.enabled=True` requires "
+                    "`trainer.policy.model.lora.rank > 0` for QLoRA."
+                )
+
+        if self.ref.model.bitsandbytes_4bit.enabled and self.strategy != "fsdp":
+            raise ValueError("`trainer.ref.model.bitsandbytes_4bit` is only supported with FSDP.")
+
         if self.logprobs_chunk_size is not None and (
             not isinstance(self.logprobs_chunk_size, int) or self.logprobs_chunk_size <= 0
         ):

skyrl/train/config/sft_config.py

@@ -140,6 +140,8 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig":
     model_config_kwargs: dict = field(default_factory=dict)
     """Pass-through kwargs for the HuggingFace model config (FSDP backends).
     For Megatron, use ``megatron_config.transformer_config_kwargs`` instead."""
+    flash_attn: bool = True
+    """Use FlashAttention 2 for FSDP models. Disable for models with unsupported head dimensions."""
     use_torch_compile: bool = False
     """Apply torch.compile to logits calculation."""
     record_memory: bool = False
@@ -558,6 +560,11 @@ def validate_sft_cfg(cfg: SFTConfig) -> None:
             raise ValueError(f"num_epochs must be > 0, got {cfg.num_epochs}")
     if not cfg.model.path:
         raise ValueError("model.path must be set")
+    if cfg.model.bitsandbytes_4bit.enabled:
+        if cfg.strategy != "fsdp":
+            raise ValueError("model.bitsandbytes_4bit.enabled=True is only supported with strategy='fsdp'.")
+        if cfg.model.lora.rank <= 0:
+            raise ValueError("model.bitsandbytes_4bit.enabled=True requires model.lora.rank > 0 for QLoRA.")
     if cfg.dummy_run_full_ctx and cfg.dummy_run_max_steps <= 0:
         raise ValueError(f"dummy_run_max_steps must be > 0, got {cfg.dummy_run_max_steps}")
     if cfg.max_training_steps is not None and cfg.max_training_steps <= 0:
@@ -656,6 +663,7 @@ def build_skyrl_config_for_sft(sft_cfg: SFTConfig) -> SkyRLTrainConfig:
     cfg.trainer.policy.sequence_parallel_size = sft_cfg.sequence_parallel_size
     cfg.trainer.policy.model_config_kwargs = sft_cfg.model_config_kwargs
     cfg.trainer.policy.use_torch_compile = sft_cfg.use_torch_compile
+    cfg.trainer.flash_attn = sft_cfg.flash_attn
     cfg.trainer.policy.record_memory = sft_cfg.record_memory
     cfg.trainer.policy.torch_profiler_config = sft_cfg.torch_profiler_config
 

tests/backends/skyrl_train/distributed/test_fsdp_utils.py

@@ -0,0 +1,29 @@
+from types import SimpleNamespace
+
+import torch.nn as nn
+
+from skyrl.backends.skyrl_train.distributed import fsdp_utils
+
+
+class WrappedBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = nn.Embedding(8, 4)
+
+
+class DummyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(tie_word_embeddings=False)
+        self.block = WrappedBlock()
+
+
+def test_apply_fsdp2_does_not_wrap_embedding_inside_selected_parent(monkeypatch):
+    model = DummyModel()
+    wrapped = []
+    monkeypatch.setattr(fsdp_utils, "fully_shard", lambda module, **kwargs: wrapped.append(module))
+    config = SimpleNamespace(wrap_policy={"transformer_layer_cls_to_wrap": ["WrappedBlock"]})
+
+    fsdp_utils.apply_fsdp2(model, {}, config)
+
+    assert wrapped == [model.block, model]