diff --git a/docs/content/docs/sft/overview.mdx b/docs/content/docs/sft/overview.mdx index 861bbc866a..d244263a27 100644 --- a/docs/content/docs/sft/overview.mdx +++ b/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 diff --git a/examples/train/sft/run_sft_gemma4_e2b_qlora.sh b/examples/train/sft/run_sft_gemma4_e2b_qlora.sh new file mode 100755 index 0000000000..cea7dd0752 --- /dev/null +++ b/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="" \ + "$@" diff --git a/pyproject.toml b/pyproject.toml index 881d6c7e71..f3954d6231 100644 --- a/pyproject.toml +++ b/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'", diff --git a/skyrl/backends/skyrl_train/distributed/fsdp_strategy.py b/skyrl/backends/skyrl_train/distributed/fsdp_strategy.py index 4220aaafcb..95adee986b 100644 --- a/skyrl/backends/skyrl_train/distributed/fsdp_strategy.py +++ b/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 diff --git a/skyrl/backends/skyrl_train/distributed/fsdp_utils.py b/skyrl/backends/skyrl_train/distributed/fsdp_utils.py index e23d36fab0..fec7f42a5b 100644 --- a/skyrl/backends/skyrl_train/distributed/fsdp_utils.py +++ b/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) diff --git a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py index ca9f5acd64..4c68e5c2ba 100644 --- a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py +++ b/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,6 +423,7 @@ 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, @@ -424,15 +431,20 @@ def init_model(self, model_path): 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, diff --git a/skyrl/backends/skyrl_train/workers/model_wrapper.py b/skyrl/backends/skyrl_train/workers/model_wrapper.py index 8235cddb73..00d5fb98a9 100644 --- a/skyrl/backends/skyrl_train/workers/model_wrapper.py +++ b/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, diff --git a/skyrl/train/config/__init__.py b/skyrl/train/config/__init__.py index fd161d069f..ea4e0b3453 100644 --- a/skyrl/train/config/__init__.py +++ b/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", diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 57c7521910..6fc3cea149 100644 --- a/skyrl/train/config/config.py +++ b/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 ): diff --git a/skyrl/train/config/sft_config.py b/skyrl/train/config/sft_config.py index e39912738e..73c6366c4a 100644 --- a/skyrl/train/config/sft_config.py +++ b/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 diff --git a/tests/backends/skyrl_train/distributed/test_fsdp_utils.py b/tests/backends/skyrl_train/distributed/test_fsdp_utils.py new file mode 100644 index 0000000000..2ddb8c0fc9 --- /dev/null +++ b/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] diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 5afc726c43..bd19461b1c 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -154,6 +154,13 @@ def test_trainer_config_rejects_invalid_vocab_entropy_chunking(field_name, value TrainerConfig(**{field_name: value}) +def test_ref_4bit_rejects_megatron(): + with pytest.raises(ValueError, match="trainer.ref.model.bitsandbytes_4bit"): + SkyRLTrainConfig.from_cli_overrides( + ["trainer.strategy=megatron", "trainer.ref.model.bitsandbytes_4bit.enabled=true"] + ) + + def test_cli_overrides_plus_prefix_rejected(): with pytest.raises(ValueError, match="The '\\+' prefix"): SkyRLTrainConfig.from_cli_overrides(["+new_field=value"]) diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 3237e508cd..62fab7ce50 100644 --- a/tests/train/test_sft_config.py +++ b/tests/train/test_sft_config.py @@ -147,6 +147,53 @@ def test_remove_microbatch_padding_propagates(self): skyrl_cfg_on = build_skyrl_config_for_sft(cfg_on) assert skyrl_cfg_on.trainer.remove_microbatch_padding is True + def test_flash_attention_propagates(self): + cfg = _sft_cfg_from_overrides(["flash_attn=false"]) + skyrl_cfg = build_skyrl_config_for_sft(cfg) + assert skyrl_cfg.trainer.flash_attn is False + + +class TestQLoRAConfig: + def test_4bit_config_bridges_to_policy_model(self): + cfg = _sft_cfg_from_overrides( + [ + "strategy=fsdp", + "model.path=test/my-model", + "model.lora.rank=8", + "model.bitsandbytes_4bit.enabled=true", + "model.bitsandbytes_4bit.quant_type=fp4", + "model.bitsandbytes_4bit.use_double_quant=false", + ] + ) + + skyrl_cfg = build_skyrl_config_for_sft(cfg) + + quantization = skyrl_cfg.trainer.policy.model.bitsandbytes_4bit + assert quantization.enabled is True + assert quantization.quant_type == "fp4" + assert quantization.use_double_quant is False + + def test_4bit_requires_lora(self): + cfg = _sft_cfg_from_overrides( + ["strategy=fsdp", "model.path=test/my-model", "model.bitsandbytes_4bit.enabled=true"] + ) + + with pytest.raises(ValueError, match="requires model.lora.rank > 0"): + validate_sft_cfg(cfg) + + def test_4bit_rejects_megatron(self): + cfg = _sft_cfg_from_overrides( + [ + "strategy=megatron", + "model.path=test/my-model", + "model.lora.rank=8", + "model.bitsandbytes_4bit.enabled=true", + ] + ) + + with pytest.raises(ValueError, match="only supported with strategy='fsdp'"): + validate_sft_cfg(cfg) + class TestMegatronConfigOverrides: """Megatron parallelism config overrides propagate correctly.""" @@ -207,6 +254,13 @@ def test_lora_target_modules_propagate(self): skyrl_cfg = build_skyrl_config_for_sft(cfg) assert skyrl_cfg.trainer.policy.model.lora.target_modules == "all-linear" + def test_lora_target_module_list_propagates(self): + cfg = _sft_cfg_from_overrides( + ["model.path=test/my-model", "model.lora.rank=16", "model.lora.target_modules=[q_a_proj]"] + ) + skyrl_cfg = build_skyrl_config_for_sft(cfg) + assert skyrl_cfg.trainer.policy.model.lora.target_modules == ["q_a_proj"] + def test_lora_disabled_by_default(self): cfg = _sft_cfg_from_overrides([]) skyrl_cfg = build_skyrl_config_for_sft(cfg) diff --git a/uv.lock b/uv.lock index ba929b031a..5d5a40ac23 100644 --- a/uv.lock +++ b/uv.lock @@ -554,6 +554,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/5c/fb93d3092640a24dfb7bd7727a24016d7c01774ca013e60efd3f683c8002/backrefs-7.0-py314-none-any.whl", hash = "sha256:a6448b28180e3ca01134c9cf09dcebafad8531072e09903c5451748a05f24bc9", size = 412349, upload-time = "2026-04-28T16:28:02.412Z" }, ] +[[package]] +name = "bitsandbytes" +version = "0.50.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'extra-5-skyrl-fsdp') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra != 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra != 'extra-5-skyrl-fsdp' and extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version >= '3.13' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra != 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-jax') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra != 'extra-5-skyrl-fsdp' and extra != 'extra-5-skyrl-jax' and extra != 'extra-5-skyrl-megatron') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version < '3.13' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra != 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and extra != 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version >= '3.13' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, + { name = "packaging", marker = "sys_platform == 'linux' or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/6c/b3c2a6b05e0fb06c6258b675436b2b090192450d7f283826e6323471e270/bitsandbytes-0.50.0-py3-none-macosx_14_0_arm64.whl", hash = "sha256:c6f482df4dc18c7c150246577025de40dd0935031ea10a1e4122599a525b39bb", size = 123217, upload-time = "2026-07-24T19:48:51.724Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3a/fe0f0c7186436038319f15156386e844abc2ae1e54b8cda3f53b81a96cb9/bitsandbytes-0.50.0-py3-none-manylinux_2_24_aarch64.whl", hash = "sha256:824d931d8e77d7db09bb26d940268d36920e0757e2a2f49cce8766e309e13048", size = 23776757, upload-time = "2026-07-25T01:34:17.305Z" }, + { url = "https://files.pythonhosted.org/packages/22/08/9501f4fc830448a6862bd5313df94a7dd1ae678f4f81087b96569d4a6f8b/bitsandbytes-0.50.0-py3-none-manylinux_2_24_x86_64.whl", hash = "sha256:173d137610468bec9cddbaa2e049254e97792657ab984e3e737bec1772c1668c", size = 40860117, upload-time = "2026-07-25T01:34:21.049Z" }, + { url = "https://files.pythonhosted.org/packages/78/40/7b2253f94ec8c6235606aa393f315d941d02c44f141b7bccae14202b3fea/bitsandbytes-0.50.0-py3-none-win_amd64.whl", hash = "sha256:7f9ebf51d76a00341711cdd557af97555274b5aad3444975ff9927866a18bc11", size = 37961095, upload-time = "2026-07-25T01:34:24.809Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e9/3cbda2ff7574753a493fd2bc27c1dc64cb51d97534ad2136f64e68d34344/bitsandbytes-0.50.0-py3-none-win_arm64.whl", hash = "sha256:9c6f63f7f35c4425661256daca56664d2e4f05a5d230a9a148cb14df207bcca2", size = 128508, upload-time = "2026-07-25T01:34:27.322Z" }, +] + [[package]] name = "blake3" version = "1.0.9" @@ -8644,6 +8662,9 @@ miniswe = [ { name = "vllm-router", version = "0.1.14.post1", source = { url = "https://github.com/SumanthRH/router/releases/download/0.1.14.post1/vllm_router-0.1.14.post1-cp38-abi3-manylinux_2_35_x86_64.whl" }, marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-5-skyrl-miniswe') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-gpu' and extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu')" }, { name = "wandb" }, ] +qlora = [ + { name = "bitsandbytes", marker = "sys_platform == 'linux' or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, +] ray = [ { name = "ray", extra = ["default"] }, ] @@ -8700,6 +8721,7 @@ requires-dist = [ { name = "aiosqlite", marker = "extra == 'tinker'" }, { name = "alembic", marker = "extra == 'dev'" }, { name = "asyncpg", marker = "extra == 'tinker'" }, + { name = "bitsandbytes", marker = "sys_platform == 'linux' and extra == 'qlora'", specifier = ">=0.47.0" }, { name = "causal-conv1d", marker = "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'fsdp'", url = "https://github.com/erictang000/causal-conv1d/releases/download/v1.6.1.post4-torch2.11/causal_conv1d-1.6.1-cp312-cp312-linux_x86_64.whl" }, { name = "causal-conv1d", marker = "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'megatron'", url = "https://github.com/erictang000/causal-conv1d/releases/download/v1.6.1.post4-torch2.11/causal_conv1d-1.6.1-cp312-cp312-linux_x86_64.whl" }, { name = "causal-conv1d", marker = "(python_full_version != '3.12.*' and sys_platform == 'linux' and extra == 'fsdp') or (platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'fsdp')" }, @@ -8819,7 +8841,7 @@ requires-dist = [ { name = "xxhash", specifier = ">=3.0.0" }, { name = "zstandard", specifier = ">=0.23.0" }, ] -provides-extras = ["gpu", "tpu", "tinker", "ray", "aws", "gcp", "azure", "jax", "skyrl-train", "fsdp", "megatron", "miniswe", "harbor", "dev"] +provides-extras = ["gpu", "tpu", "tinker", "ray", "aws", "gcp", "azure", "jax", "skyrl-train", "fsdp", "qlora", "megatron", "miniswe", "harbor", "dev"] [[package]] name = "skyrl-gym"