Abstract optimizer

Author: EduardDurech

docs/examples/config.rst

@@ -614,12 +614,17 @@ Optim
 .. code:: yaml
 
    optim:
+     optimizer: AdamW
+     optimizer_impl: torch.optim
      lr: 1e-5
      weight_decay: 0.01
      warmup_steps_ratio: 0.1
      clip_grad: 1.0
      lr_scheduler: cosine
+     override_optimizer_config: null
 
+- ``optimizer``: Optimizer class name (e.g., ``"AdamW"``, ``"AdamW8bit"``, ``"_AdamW"``). The class name as it appears in the module.
+- ``optimizer_impl``: Module path to import optimizer from (e.g., ``"torch.optim"``, ``"torchao.optim"``, ``"bitsandbytes.optim"``).
 - ``optim.lr``: Learning rate for the optimizer.
 - ``optim.weight_decay``: Weight decay for the optimizer.
 - ``optim.warmup_steps_ratio``: Ratio of warmup steps to total training steps.
@@ -629,6 +634,8 @@ Optim
   - ``cosine``: Cosine learning rate scheduler with warmup (default).
   - ``wsd``: Warmup-Stable-Decay scheduler that provides a stable learning rate phase between warmup and decay phases.
 
+- ``override_optimizer_config``: Dictionary of additional optimizer-specific keyword arguments. For example, to use ``torchao.optim``'s ``_AdamW`` with BF16 stochastic rounding: ``{"bf16_stochastic_round": true}``
+
 Model
 ~~~~~~~~~~~~
 

recipe/prime/prime_fsdp_workers.py

@@ -40,6 +40,7 @@
 )
 from verl.utils.import_utils import import_external_libs
 from verl.utils.profiler import log_gpu_memory_usage
+from verl.workers.config.optimizer import build_optimizer
 from verl.workers.fsdp_workers import create_device_mesh, get_sharding_strategy
 from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager
 
@@ -87,7 +88,6 @@ def __init__(self, config):
 
     def _build_reward_ref_model_optimizer(self, config):
         # the following line is necessary
-        from torch import optim
         from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
         from torch.distributed.fsdp import MixedPrecision
 
@@ -219,12 +219,7 @@ def _build_reward_ref_model_optimizer(self, config):
             cpu_offload=None,
         )
 
-        reward_optimizer = optim.AdamW(
-            reward_module.parameters(),
-            lr=config.model.optim.lr,
-            betas=config.model.optim.get("betas", (0.9, 0.999)),
-            weight_decay=config.model.optim.get("weight_decay", 1e-2),
-        )
+        reward_optimizer = build_optimizer(reward_module.parameters(), config.model.optim)
 
         total_steps = config.model.optim.get("total_training_steps", 0)
         num_warmup_steps = int(config.model.optim.get("lr_warmup_steps", -1))

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -7,6 +7,8 @@ actor_rollout_ref:
   actor:
     optim:
       _target_: verl.workers.config.FSDPOptimizerConfig
+      optimizer: AdamW
+      optimizer_impl: torch.optim
       lr: 1.0e-06
       lr_warmup_steps_ratio: 0.0
       total_training_steps: -1
@@ -19,6 +21,7 @@ actor_rollout_ref:
       min_lr_ratio: 0.0
       num_cycles: 0.5
       warmup_style: constant
+      override_optimizer_config: null
     fsdp_config:
       _target_: verl.workers.config.FSDPEngineConfig
       wrap_policy:
@@ -302,6 +305,8 @@ data:
 critic:
   optim:
     _target_: verl.workers.config.FSDPOptimizerConfig
+    optimizer: AdamW
+    optimizer_impl: torch.optim
     lr: 1.0e-05
     lr_warmup_steps_ratio: 0.0
     total_training_steps: -1
@@ -314,6 +319,7 @@ critic:
     min_lr_ratio: 0.0
     num_cycles: 0.5
     warmup_style: constant
+    override_optimizer_config: null
   model:
     fsdp_config:
       _target_: verl.workers.config.FSDPEngineConfig

verl/trainer/config/optim/fsdp.yaml

@@ -1,6 +1,13 @@
 # Target class for this configuration
 _target_: verl.workers.config.FSDPOptimizerConfig
 
+# Optimizer class name (e.g., "AdamW", "AdamW8bit", "_AdamW", "Adam")
+optimizer: AdamW
+
+# Module path to import optimizer
+# Examples: "torch.optim", "torchao.optim", "bitsandbytes.optim"
+optimizer_impl: torch.optim
+
 # Learning rate
 lr: 1e-3
 
@@ -31,3 +38,11 @@ num_cycles: 0.5
 # LR warmup style: "constant" or "cosine"
 warmup_style: constant
 
+# Additional optimizer-specific keyword arguments
+# Example for torchao with bf16 stochastic rounding:
+# optimizer_impl: torchao.optim
+# optimizer: _AdamW
+# override_optimizer_config:
+#   bf16_stochastic_round: true
+override_optimizer_config: null
+

verl/trainer/fsdp_sft_trainer.py

@@ -34,7 +34,7 @@
 from omegaconf import DictConfig, OmegaConf
 from peft import LoraConfig, TaskType, get_peft_model
 from tensordict import TensorDict
-from torch import nn, optim
+from torch import nn
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.fsdp import CPUOffload, MixedPrecision, ShardingStrategy
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
@@ -73,6 +73,7 @@
     get_ulysses_sequence_parallel_world_size,
     ulysses_pad_and_slice_inputs,
 )
+from verl.workers.config.optimizer import build_optimizer
 from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager
 
 logger = logging.getLogger(__file__)
@@ -317,12 +318,7 @@ def _build_model_optimizer(self):
 
         log_gpu_memory_usage("After FSDP wrapping", logger=logger)
 
-        self.optimizer = optim.AdamW(
-            self.fsdp_model.parameters(),
-            lr=self.config.optim.lr,
-            betas=self.config.optim.betas,
-            weight_decay=self.config.optim.weight_decay,
-        )
+        self.optimizer = build_optimizer(self.fsdp_model.parameters(), self.config.optim)
 
         log_gpu_memory_usage("After initialize optimizer", logger=logger)
 

verl/workers/config/optimizer.py

@@ -19,7 +19,7 @@
 
 from verl.base_config import BaseConfig
 
-__all__ = ["OptimizerConfig", "FSDPOptimizerConfig", "McoreOptimizerConfig"]
+__all__ = ["OptimizerConfig", "FSDPOptimizerConfig", "McoreOptimizerConfig", "build_optimizer"]
 
 
 @dataclass
@@ -58,15 +58,22 @@ class FSDPOptimizerConfig(OptimizerConfig):
     """FSDP optimizer configuration extending base OptimizerConfig.
 
     Args:
+        optimizer (str): Optimizer class name (e.g., "AdamW", "AdamW8bit", "_AdamW").
+        optimizer_impl (str): Module path to import optimizer from (e.g., "torch.optim", "torchao.optim",
+            "bitsandbytes.optim").
         lr (float): Learning rate.
         min_lr_ratio (Optional[float]): Minimum LR ratio for cosine schedule.
         warmup_style (str): LR warmup style: "constant" or "cosine".
         num_cycles (float): Number of cosine cycles in LR schedule.
+        override_optimizer_config (Optional[dict]): Additional optimizer-specific keyword arguments.
     """
 
+    optimizer: str = "AdamW"
+    optimizer_impl: str = "torch.optim"
     min_lr_ratio: Optional[float] = None
     warmup_style: str = "constant"
     num_cycles: float = 0.5
+    override_optimizer_config: Optional[dict] = None
 
     def __post_init__(self):
         assert self.warmup_style in ["constant", "cosine"]
@@ -101,3 +108,59 @@ class McoreOptimizerConfig(OptimizerConfig):
     lr_wsd_decay_steps: Optional[int] = None
     use_checkpoint_opt_param_scheduler: bool = False
     override_optimizer_config: Optional[dict] = None
+
+
+def build_optimizer(parameters, config: FSDPOptimizerConfig):
+    """Build an optimizer based on the configuration.
+
+    Dynamically imports and instantiates an optimizer class from the specified module.
+
+    Args:
+        parameters: Model parameters to optimize
+        config: FSDPOptimizerConfig with optimizer settings
+
+    Returns:
+        Optimizer instance
+
+    Examples:
+        # PyTorch AdamW
+        config.optimizer_impl = "torch.optim"
+        config.optimizer = "AdamW"
+
+        # TorchAO AdamW with bf16 stochastic rounding
+        config.optimizer_impl = "torchao.optim"
+        config.optimizer = "_AdamW"
+        config.override_optimizer_config = {"bf16_stochastic_round": True}
+
+        # BitsAndBytes AdamW 8bit
+        config.optimizer_impl = "bitsandbytes.optim"
+        config.optimizer = "AdamW8bit"
+    """
+    import importlib
+
+    optimizer_args = {
+        "lr": config.lr,
+        "weight_decay": config.weight_decay,
+    }
+
+    optimizer_name_lower = config.optimizer.lower()
+    if "adam" in optimizer_name_lower or "ademamix" in optimizer_name_lower:
+        optimizer_args["betas"] = config.betas
+
+    if config.override_optimizer_config is not None:
+        optimizer_args.update(config.override_optimizer_config)
+
+    try:
+        module = importlib.import_module(config.optimizer_impl)
+        optimizer_cls = getattr(module, config.optimizer)
+    except ImportError as e:
+        raise ImportError(
+            f"Failed to import module '{config.optimizer_impl}'. Make sure the package is installed. Error: {e}"
+        ) from e
+    except AttributeError as e:
+        raise AttributeError(
+            f"Optimizer '{config.optimizer}' not found in module '{config.optimizer_impl}'. "
+            f"Available optimizers: {dir(module)}"
+        ) from e
+
+    return optimizer_cls(parameters, **optimizer_args)

verl/workers/engine/fsdp/transformer_impl.py

@@ -354,14 +354,10 @@ def _build_fsdp_module(self, module):
         return module
 
     def _build_optimizer(self, module):
-        from torch import optim
+        from verl.workers.config.optimizer import build_optimizer
+
+        optimizer = build_optimizer(module.parameters(), self.optimizer_config)
 
-        optimizer = optim.AdamW(
-            module.parameters(),
-            lr=self.optimizer_config.lr,
-            betas=self.optimizer_config.betas,
-            weight_decay=self.optimizer_config.weight_decay,
-        )
         return optimizer
 
     def _build_lr_scheduler(self, optimizer):

verl/workers/fsdp_workers.py

@@ -86,6 +86,7 @@
 from verl.utils.profiler.performance import reduce_timing, topk_reduce_ratio_min_max
 from verl.utils.py_functional import convert_to_regular_types
 from verl.workers.config import FSDPCriticConfig, FSDPEngineConfig, HFModelConfig, RolloutConfig
+from verl.workers.config.optimizer import build_optimizer
 from verl.workers.rollout import get_rollout_class
 from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager
 
@@ -279,7 +280,6 @@ def _build_model_optimizer(
         role="actor",
         enable_activation_offload=False,
     ):
-        from torch import optim
         from torch.distributed.fsdp import CPUOffload, MixedPrecision
         from transformers import (
             AutoConfig,
@@ -520,12 +520,7 @@ def _build_model_optimizer(
         if role == "actor" and optim_config is not None:
             from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup
 
-            actor_optimizer = optim.AdamW(
-                actor_module_fsdp.parameters(),
-                lr=optim_config.lr,
-                betas=optim_config.get("betas", (0.9, 0.999)),
-                weight_decay=optim_config.get("weight_decay", 1e-2),
-            )
+            actor_optimizer = build_optimizer(actor_module_fsdp.parameters(), optim_config)
 
             total_steps = optim_config.get("total_training_steps", 0)
             num_warmup_steps = int(optim_config.get("lr_warmup_steps", -1))
@@ -866,7 +861,7 @@ def update_actor(self, data: DataProto):
             metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / (1024**3)
 
             lr = self.actor_lr_scheduler.get_last_lr()[0]
-            metrics["actor/lr"] = lr
+            metrics["actor/lr"] = lr.item() if torch.is_tensor(lr) else lr
             self.actor_lr_scheduler.step()
 
             # TODO: here, we should return all metrics
@@ -1187,7 +1182,6 @@ def __init__(self, config: FSDPCriticConfig):
 
     def _build_critic_model_optimizer(self, config):
         # the following line is necessary
-        from torch import optim
         from torch.distributed.fsdp import MixedPrecision
 
         from verl.utils.model import load_valuehead_model, print_model_size
@@ -1368,12 +1362,7 @@ def _build_critic_model_optimizer(self, config):
 
         log_gpu_memory_usage("After critic FSDP", logger=None)
 
-        critic_optimizer = optim.AdamW(
-            critic_module.parameters(),
-            lr=config.optim.lr,
-            betas=config.optim.get("betas", (0.9, 0.999)),
-            weight_decay=config.optim.get("weight_decay", 1e-2),
-        )
+        critic_optimizer = build_optimizer(critic_module.parameters(), config.optim)
 
         total_steps = config.optim.get("total_training_steps", 0)
         num_warmup_steps = int(config.optim.get("lr_warmup_steps", -1))