agilerl 2.4.1.dev0__tar.gz → 2.4.1.dev2__tar.gz

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/PKG-INFO

@@ -1,20 +1,23 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: agilerl
-Version: 2.4.1.dev0
+Version: 2.4.1.dev2
 Summary: AgileRL is a deep reinforcement learning library focused on improving RL development through RLOps.
 License: Apache 2.0
+License-File: LICENSE
 Author: Nick Ustaran-Anderegg
 Author-email: dev@agilerl.com
-Requires-Python: >=3.10,<4.0
+Requires-Python: >=3.10,<3.13
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
+Provides-Extra: all
+Provides-Extra: llm
 Requires-Dist: SuperSuit (>=3.9.0,<4.0.0)
 Requires-Dist: accelerate (>=1.7.0,<2.0.0)
-Requires-Dist: deepspeed (>=0.17.1,<0.18.0)
+Requires-Dist: datasets (==4.4.1) ; extra == "llm" or extra == "all"
+Requires-Dist: deepspeed (>=0.17.1,<0.18.0) ; extra == "llm" or extra == "all"
 Requires-Dist: dill (>=0.3.7,<0.4.0)
 Requires-Dist: fastrand (>=1.3.0,<2.0.0)
 Requires-Dist: flatten_dict (>=0.4.2,<0.5.0)
@@ -24,11 +27,12 @@ Requires-Dist: h5py (>=3.8.0,<4.0.0)
 Requires-Dist: hydra-core (>=1.3.2,<2.0.0)
 Requires-Dist: jax[cpu] (>=0.4.31,<0.5.0)
 Requires-Dist: matplotlib (>=3.9.4,<3.10.0)
-Requires-Dist: minari (>=0.5.2,<0.6.0)
+Requires-Dist: minari[all] (==0.5.2)
 Requires-Dist: numpy (>=1.26.4,<2.0.0)
 Requires-Dist: omegaconf (>=2.3.0,<3.0.0)
+Requires-Dist: packaging (==25.0)
 Requires-Dist: pandas (>=2.2.3,<3.0.0)
-Requires-Dist: peft (>=0.15.2,<0.16.0)
+Requires-Dist: peft (>=0.18.0,<0.19.0) ; extra == "llm" or extra == "all"
 Requires-Dist: pettingzoo (>=1.23.1,<2.0.0)
 Requires-Dist: pre-commit (>=3.4.0,<4.0.0)
 Requires-Dist: pygame (>=2.6.0,<3.0.0)
@@ -39,9 +43,9 @@ Requires-Dist: tensordict (>=0.8,<0.9)
 Requires-Dist: termcolor (>=1.1.0,<2.0.0)
 Requires-Dist: torch (==2.7.1)
 Requires-Dist: tqdm (>=4.66.4,<5.0.0)
-Requires-Dist: transformers (>=4.48.1,<5.0.0)
+Requires-Dist: transformers (>=4.57.1,<5.0.0) ; extra == "llm" or extra == "all"
 Requires-Dist: ucimlrepo (>=0.0.3,<0.0.4)
-Requires-Dist: vllm (==0.10.0)
+Requires-Dist: vllm (==0.10.0) ; extra == "llm" or extra == "all"
 Requires-Dist: wandb (>=0.17.6,<0.18.0)
 Description-Content-Type: text/markdown
 
@@ -95,6 +99,16 @@ git clone https://github.com/AgileRL/AgileRL.git && cd AgileRL
 pip install -e .
 ```
 
+If you wish to install all additional dependencies please specify `[all]` or if you want to install a specific family of dependencies specify that family directly. At present, we have just one family, `[llm]`, which contains the dependencies related to our LLM RFT algorithms (datasets, deepspeed, peft, transformers, vllm).
+
+```bash
+pip install agilerl[all]
+```
+Or in development mode:
+```bash
+pip install -e ".[all]"
+```
+
 To install the ``nightly`` version of AgileRL with the latest features, use:
 
 ```bash
@@ -153,11 +167,12 @@ We are constantly updating our tutorials to showcase the latest features of Agil
   | ---------- | --------- |
   | [Bandits](https://docs.agilerl.com/en/latest/bandits/index.html) | [Neural Contextual Bandits with UCB-based Exploration (NeuralUCB)](https://docs.agilerl.com/en/latest/api/algorithms/neural_ucb.html) <br> [Neural Contextual Bandits with Thompson Sampling (NeuralTS)](https://docs.agilerl.com/en/latest/api/algorithms/neural_ts.html) |
 
-  ### LLM Reasoning Algorithms
+  ### LLM Fine-tuning Algorithms
 
   | RL         | Algorithm |
   | ---------- | --------- |
   | [On-Policy](https://docs.agilerl.com/en/latest/llm_finetuning/index.html) | [Group Relative Policy Optimization (GRPO)](https://docs.agilerl.com/en/latest/api/algorithms/grpo.html)
+  | [Off-Policy](https://docs.agilerl.com/en/latest/llm_finetuning/index.html) | [Direct Preference Optimization (DPO)](https://docs.agilerl.com/en/latest/api/algorithms/dpo.html)
 
 
 ## Train an Agent to Beat a Gym Environment

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/README.md

@@ -48,6 +48,16 @@ git clone https://github.com/AgileRL/AgileRL.git && cd AgileRL
 pip install -e .
 ```
 
+If you wish to install all additional dependencies please specify `[all]` or if you want to install a specific family of dependencies specify that family directly. At present, we have just one family, `[llm]`, which contains the dependencies related to our LLM RFT algorithms (datasets, deepspeed, peft, transformers, vllm).
+
+```bash
+pip install agilerl[all]
+```
+Or in development mode:
+```bash
+pip install -e ".[all]"
+```
+
 To install the ``nightly`` version of AgileRL with the latest features, use:
 
 ```bash
@@ -106,11 +116,12 @@ We are constantly updating our tutorials to showcase the latest features of Agil
   | ---------- | --------- |
   | [Bandits](https://docs.agilerl.com/en/latest/bandits/index.html) | [Neural Contextual Bandits with UCB-based Exploration (NeuralUCB)](https://docs.agilerl.com/en/latest/api/algorithms/neural_ucb.html) <br> [Neural Contextual Bandits with Thompson Sampling (NeuralTS)](https://docs.agilerl.com/en/latest/api/algorithms/neural_ts.html) |
 
-  ### LLM Reasoning Algorithms
+  ### LLM Fine-tuning Algorithms
 
   | RL         | Algorithm |
   | ---------- | --------- |
   | [On-Policy](https://docs.agilerl.com/en/latest/llm_finetuning/index.html) | [Group Relative Policy Optimization (GRPO)](https://docs.agilerl.com/en/latest/api/algorithms/grpo.html)
+  | [Off-Policy](https://docs.agilerl.com/en/latest/llm_finetuning/index.html) | [Direct Preference Optimization (DPO)](https://docs.agilerl.com/en/latest/api/algorithms/dpo.html)
 
 
 ## Train an Agent to Beat a Gym Environment

agilerl-2.4.1.dev2/agilerl/__init__.py

@@ -0,0 +1,18 @@
+from importlib.metadata import metadata
+from importlib.util import find_spec
+
+from packaging.requirements import Requirement
+
+
+def get_extra_dependencies(package: str, extra: str) -> list[str]:
+    requires = metadata(package).get_all("Requires-Dist") or []
+    deps = []
+    for req in requires:
+        r = Requirement(req)
+        if r.marker and r.marker.evaluate({"extra": extra}):
+            deps.append(r.name)
+    return deps
+
+
+LLM_PACKAGES = get_extra_dependencies("agilerl", "llm")
+HAS_LLM_DEPENDENCIES = all(find_spec(pkg) is not None for pkg in LLM_PACKAGES)

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/agilerl/algorithms/core/base.py

@@ -27,20 +27,14 @@ import torch
 import torch.nn.functional as F
 from accelerate import Accelerator
 from accelerate.utils import broadcast_object_list, set_seed
-from accelerate.utils.deepspeed import DeepSpeedOptimizerWrapper
-from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
 from gymnasium import spaces
-from peft import LoraConfig, PeftModel, get_peft_model, set_peft_model_state_dict
-from safetensors.torch import load_file
 from tensordict import TensorDict
 from torch._dynamo import OptimizedModule
 from torch.nn.utils import clip_grad_norm_
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import SequentialLR
-from transformers import PretrainedConfig
-from transformers.modeling_utils import PreTrainedModel
-from vllm import LLM, SamplingParams
 
+from agilerl import HAS_LLM_DEPENDENCIES
 from agilerl.algorithms.core.optimizer_wrapper import OptimizerWrapper
 from agilerl.algorithms.core.registry import (
     HyperparameterConfig,
@@ -55,7 +49,11 @@ from agilerl.protocols import (
     EvolvableAttributeDict,
     EvolvableAttributeType,
     EvolvableModule,
+    LoraConfigProtocol,
     ModuleDict,
+    PeftModelProtocol,
+    PretrainedConfigProtocol,
+    PreTrainedModelProtocol,
 )
 from agilerl.typing import (
     ActionType,
@@ -74,6 +72,7 @@ from agilerl.typing import (
 )
 from agilerl.utils.algo_utils import (
     CosineLRScheduleConfig,
+    DummyOptimizer,
     VLLMConfig,
     check_supported_space,
     chkpt_attribute_to_device,
@@ -96,11 +95,18 @@ from agilerl.utils.evolvable_networks import (
     is_image_space,
     is_vector_space,
 )
-from agilerl.utils.llm_utils import (
-    DummyOptimizer,
-    create_model_from_name_or_path,
-    gather_if_zero3,
-)
+
+if HAS_LLM_DEPENDENCIES:
+    from accelerate.utils.deepspeed import DeepSpeedOptimizerWrapper
+    from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
+    from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
+    from safetensors.torch import load_file
+    from vllm import LLM, SamplingParams
+
+    from agilerl.utils.llm_utils import (
+        create_model_from_name_or_path,
+        gather_if_zero3,
+    )
 
 __all__ = ["EvolvableAlgorithm", "RLAlgorithm", "MultiAgentRLAlgorithm"]
 
@@ -601,14 +607,16 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
         )
         optimizer = opt.optimizer if hasattr(opt, "optimizer") else None
 
-        if isinstance(opt, DeepSpeedOptimizerWrapper):
-            if isinstance(opt.optimizer, DummyOptimizer):
-                opt = getattr(
+        if isinstance(self, LLMAlgorithm):
+            if hasattr(self.actor, "optimizer"):
+                optimizer = getattr(
                     getattr(self, "actor"), "optimizer"
                 )  # If the optimizer is defined in the deepspeed config, we do this
+            else:
+                optimizer = opt.optimizer
 
             self.accelerator, self.lr_scheduler = LLMAlgorithm.update_lr(
-                opt,
+                optimizer,
                 lr=getattr(self, config.lr),
                 accelerator=self.accelerator,
                 scheduler_config=self.cosine_lr_schedule_config,
@@ -1143,6 +1151,16 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
 
         return self
 
+    def clean_up(self) -> None:
+        """
+        Clean up the algorithm by deleting the networks and optimizers.
+
+        :return: None
+        :rtype: None
+        """
+        for evo_attr in self.evolvable_attributes().values():
+            del evo_attr
+
 
 class RLAlgorithm(EvolvableAlgorithm, ABC):
     """Base object for all single-agent algorithms in the AgileRL framework.
@@ -1799,6 +1817,10 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     :type accelerator: Optional[Accelerator]
     :param name: The name of the algorithm.
     :type name: Optional[str]
+    :param model_config: The configuration for the model.
+    :type model_config: dict[str, Any] | PretrainedConfig | None
+    :param gradient_checkpointing: Whether to use gradient checkpointing.
+    :type gradient_checkpointing: bool
     """
 
     def __init__(
@@ -1813,10 +1835,10 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         seed: int,
         pad_token_id: int,
         pad_token: str,
-        lora_config: LoraConfig | None,
+        lora_config: LoraConfigProtocol | None,
         use_separate_reference_adapter: bool,
         model_name: str | None = None,
-        actor_network: PreTrainedModel | None = None,
+        actor_network: PreTrainedModelProtocol | None = None,
         micro_batch_size_per_gpu: int | None = None,
         cosine_lr_schedule_config: Optional[CosineLRScheduleConfig] = None,
         hp_config: Optional[HyperparameterConfig] = None,
@@ -1824,9 +1846,14 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         device: Union[str, torch.device] = "cpu",
         accelerator: Optional[Accelerator] = None,
         name: Optional[str] = None,
-        model_config: dict[str, Any] | PretrainedConfig | None = None,
+        model_config: dict[str, Any] | PretrainedConfigProtocol | None = None,
         gradient_checkpointing: bool = True,
     ):
+        if not HAS_LLM_DEPENDENCIES:
+            raise ImportError(
+                "LLM dependencies are not installed. Please install them using `pip install agilerl[llm]`."
+            )
+
         if model_name is None and actor_network is None:
             raise ValueError(
                 "At least one of model_name or actor_network must be provided."
@@ -1881,7 +1908,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                     )
                     lr = optim_lr
 
-        if lora_config is None and not isinstance(actor_network, PeftModel):
+        if lora_config is None and not isinstance(actor_network, PeftModelProtocol):
             warnings.warn(
                 "No LoRA config provided. AgileRL can only be used to finetune adapters at present. Using default LoRA configuration for RL finetuning."
             )
@@ -1898,15 +1925,21 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         self.use_separate_reference_adapter = use_separate_reference_adapter
         self.cosine_lr_schedule_config = cosine_lr_schedule_config
 
-        if max_grad_norm and (accelerator is not None) and accelerator.is_main_process:
-            warnings.warn(
-                "Argument 'max_grad_norm' will be overwritten by the 'gradient_clipping' value set in the deepspeed config."
-            )
-            self.max_grad_norm = None
-        else:
-            self.max_grad_norm = max_grad_norm
+        if max_grad_norm and (accelerator is not None):
+            if accelerator.is_main_process:
+                warnings.warn(
+                    "Argument 'max_grad_norm' will overwrite the equivalent value set for 'gradient_clipping' in the deepspeed config."
+                )
+            self.accelerator.state.deepspeed_plugin.deepspeed_config[
+                "gradient_clipping"
+            ] = max_grad_norm
+
+        self.max_grad_norm = max_grad_norm
         self.reduce_memory_peak = reduce_memory_peak
 
+        if self.accelerator is not None:
+            self.register_mutation_hook(self._sync_deepspeed_gradient_clipping)
+
         if self.accelerator is not None:
             self.zero_stage = self.accelerator.state.deepspeed_plugin.deepspeed_config[
                 "zero_optimization"
@@ -2041,7 +2074,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 device_map="auto"
             )
             tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
-            model = PeftModel.from_pretrained(base_model, path)
+            model = PeftModelProtocol.from_pretrained(base_model, path)
             """
         )
 
@@ -2153,6 +2186,11 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     def clean_up(self) -> None:
         """Clean up the algorithm."""
         if self.accelerator is not None:
+            # Free up GPU memory occupied by parameters
+            if hasattr(self.actor, "empty_partition_cache"):
+                self.actor.empty_partition_cache()
+            if hasattr(self.actor, "destroy"):
+                self.actor.destroy()
             (
                 self.actor,
                 self.optimizer,
@@ -2176,10 +2214,8 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         if hasattr(self, "llm"):
             del self.llm.llm_engine.model_executor
             del self.llm
-
         gc.collect()
         torch.cuda.empty_cache()
-        torch.cuda.reset_peak_memory_stats()
         torch.cuda.synchronize()
 
     def clone(self, index: Optional[int] = None, wrap: bool = True):
@@ -2214,8 +2250,8 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             input_args["wrap"] = False
             input_args["clone"] = True
 
-            actor: PeftModel = cast(
-                PeftModel,
+            actor: PeftModelProtocol = cast(
+                PeftModelProtocol,
                 (
                     self.accelerator.unwrap_model(self.actor)
                     if self.accelerator is not None
@@ -2407,12 +2443,12 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             self.reference_update_tracker += 1
 
     def _initialize_actors(
-        self, base_model: PreTrainedModel | None, add_adapters: bool = True
+        self, base_model: PreTrainedModelProtocol | None, add_adapters: bool = True
     ):
         """Initialize the actor network.
 
         :param base_model: Base model
-        :type base_model: PreTrainedModel
+        :type base_model: PreTrainedModelProtocol
         :param add_adapters: Flag to indicate if adapters should be added to the model, defaults to True
         :type add_adapters: bool, optional
         """
@@ -2422,7 +2458,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 self.pretrained_model_name_or_path
             )
 
-        if isinstance(base_model, PeftModel) and add_adapters:
+        if isinstance(base_model, PeftModelProtocol) and add_adapters:
             # Handles backwards compatibility with user providing a peft model as the actor network
             if self.lora_config is None:
                 adapter_name = list(base_model.peft_config.keys())
@@ -2432,7 +2468,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             if "default" in list(base_model.peft_config.keys()):
                 base_model.peft_config.pop("default")
 
-        self.actor: PeftModel = (
+        self.actor: PeftModelProtocol = (
             get_peft_model(base_model, self.lora_config, adapter_name="actor")
             if add_adapters
             else base_model
@@ -2581,7 +2617,6 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     def _move_model_to_vllm(self) -> None:
         """Move the deepspeed model to vllm."""
 
-        # TODO: Add support for ZeRO Stage 3
         if self.accelerator is not None:
             self.accelerator.wait_for_everyone()
         model_ref = self.accelerator.unwrap_model(self.actor)
@@ -2949,3 +2984,28 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
 
         if self.accelerator is not None:
             self.accelerator.wait_for_everyone()
+
+    def _sync_deepspeed_gradient_clipping(self) -> None:
+        """Synchronizes max_grad_norm with DeepSpeed gradient_clipping config.
+        Registered as a mutation hook to ensure consistency after mutations.
+        """
+        if self.accelerator is None:
+            return
+
+        if (
+            "gradient_clipping"
+            not in self.accelerator.state.deepspeed_plugin.deepspeed_config
+        ):
+            return
+
+        ds_config = self.accelerator.state.deepspeed_plugin.deepspeed_config
+        if ds_config["gradient_clipping"] != self.max_grad_norm:
+            self.accelerator.state.deepspeed_plugin.deepspeed_config[
+                "gradient_clipping"
+            ] = self.max_grad_norm
+
+        if hasattr(self.actor, "optimizer"):
+            if hasattr(self.actor.optimizer, "grad_clip"):
+                self.actor.optimizer.grad_clip = self.max_grad_norm
+            if hasattr(self.actor.optimizer, "clip_grad"):
+                self.actor.optimizer.clip_grad = self.max_grad_norm

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/agilerl/algorithms/core/optimizer_wrapper.py

@@ -2,19 +2,27 @@ import inspect
 from typing import Any, Optional, Union
 
 import torch.nn as nn
-from peft import PeftModel
 from torch.optim import Optimizer
 
+from agilerl import HAS_LLM_DEPENDENCIES
 from agilerl.modules import EvolvableModule, ModuleDict
 from agilerl.protocols import EvolvableAlgorithm
 from agilerl.typing import OptimizerType, StateDict
-from agilerl.utils.llm_utils import DummyOptimizer
+from agilerl.utils.algo_utils import DummyOptimizer
+
+if HAS_LLM_DEPENDENCIES:
+    from peft import PeftModel
+
+    PeftModelType = PeftModel
+else:
+    PeftModelType = "PeftModel"
+
 
 ModuleList = list[EvolvableModule]
 _Optimizer = Union[
     type[OptimizerType], dict[str, type[OptimizerType]], type[DummyOptimizer]
 ]
-_Module = Union[EvolvableModule, ModuleDict, ModuleList, PeftModel]
+_Module = Union[EvolvableModule, ModuleDict, ModuleList, PeftModelType]
 
 
 def init_from_multiple(

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/agilerl/algorithms/core/registry.py

@@ -9,7 +9,7 @@ from torch.optim import Optimizer
 
 from agilerl.protocols import EvolvableAlgorithm
 from agilerl.typing import NetworkType
-from agilerl.utils.llm_utils import DummyOptimizer
+from agilerl.utils.algo_utils import DummyOptimizer
 
 
 @dataclass

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/agilerl/algorithms/dpo.py

@@ -5,11 +5,10 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 from accelerate import Accelerator
-from peft import LoraConfig
-from transformers import PreTrainedModel
 
 from agilerl.algorithms.core.base import LLMAlgorithm
 from agilerl.algorithms.core.registry import HyperparameterConfig, NetworkGroup
+from agilerl.protocols import LoraConfigProtocol, PreTrainedModelProtocol
 from agilerl.typing import ExperiencesType, LLMObsType
 from agilerl.utils.algo_utils import get_experiences_samples
 from agilerl.utils.llm_utils import PreferenceGym
@@ -25,7 +24,7 @@ class DPO(LLMAlgorithm):
     :param model_name: Model name
     :type model_name: str, optional
     :param actor_network: HuggingFace LLM
-    :type actor_network: PreTrainedModel
+    :type actor_network: PreTrainedModelProtocol
     :param model_config: Model configuration, to be used when creating the model from a name or path
     :param hp_config: RL hyperparameter mutation configuration, defaults to None, whereby algorithm mutations are disabled.
     :type hp_config: HyperparameterConfig, optional
@@ -50,7 +49,7 @@ class DPO(LLMAlgorithm):
     :param device: Device for accelerated computing, 'cpu' or 'cuda', defaults to 'cpu'
     :type device: str, optional
     :param lora_config: Config for LoRA, defaults to None
-    :type lora_config: LoraConfig, optional
+    :type lora_config: LoraConfigProtocol, optional
     :param accelerator: Accelerator for distributed computing, defaults to None
     :type accelerator: accelerate.Accelerator(), optional
     :param wrap: Wrap models for distributed training upon creation, defaults to True
@@ -70,7 +69,7 @@ class DPO(LLMAlgorithm):
         pad_token_id: int,
         pad_token: str,
         model_name: str | None = None,
-        actor_network: PreTrainedModel | None = None,
+        actor_network: PreTrainedModelProtocol | None = None,
         model_config: dict[str, Any] | None = None,
         hp_config: HyperparameterConfig | None = None,
         index: int = 0,
@@ -83,7 +82,7 @@ class DPO(LLMAlgorithm):
         micro_batch_size_per_gpu: int | None = None,
         reduce_memory_peak: bool = False,
         device: str = "cpu",
-        lora_config: LoraConfig | None = None,
+        lora_config: LoraConfigProtocol | None = None,
         accelerator: Accelerator | None = None,
         wrap: bool = True,
         clone: bool = False,

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/agilerl/algorithms/grpo.py

@@ -1,17 +1,18 @@
 import gc
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import numpy as np
 import torch
 from accelerate import Accelerator
-from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
-from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
-from peft import LoraConfig, PeftModel
-from transformers import GenerationConfig
-from transformers.modeling_utils import PreTrainedModel
 
+from agilerl import HAS_LLM_DEPENDENCIES
 from agilerl.algorithms.core import LLMAlgorithm
 from agilerl.algorithms.core.registry import HyperparameterConfig, NetworkGroup
+from agilerl.protocols import (
+    LoraConfigProtocol,
+    PeftModelProtocol,
+    PreTrainedModelProtocol,
+)
 from agilerl.typing import ExperiencesType, LLMObsType
 from agilerl.utils.algo_utils import (
     CosineLRScheduleConfig,
@@ -23,10 +24,8 @@ from agilerl.utils.llm_utils import (
     ReasoningGym,
 )
 
-DeepSpeedOptimizerType = Union[
-    DeepSpeedZeroOptimizer,  # ZeRO Stage 1 & 2 optimizer
-    DeepSpeedZeroOptimizer_Stage3,  # ZeRO Stage 3 optimizer
-]
+if HAS_LLM_DEPENDENCIES:
+    from transformers import GenerationConfig
 
 
 class GRPO(LLMAlgorithm):
@@ -39,7 +38,7 @@ class GRPO(LLMAlgorithm):
     :param model_name: Model name
     :type model_name: str, optional
     :param actor_network: HuggingFace LLM
-    :type actor_network: PreTrainedModel
+    :type actor_network: PreTrainedModelProtocol
     :param model_config: Model configuration, to be used when creating the model from a name or path
     :type model_config: dict[str, Any], optional
     :param hp_config: RL hyperparameter mutation configuration, defaults to None, whereby algorithm mutations are disabled.
@@ -77,7 +76,7 @@ class GRPO(LLMAlgorithm):
     :param max_model_len: Maximum context window length, defaults to None
     :type max_model_len: int, optional
     :param lora_config: Config for LoRA, defaults to None
-    :type lora_config: LoraConfig, optional
+    :type lora_config: LoraConfigProtocol, optional
     :param cosine_lr_schedule_config: Config for cosine lr scheduling, defaults to None
     :type cosine_lr_schedule_config: CosineLRScheduleConfig, optional
     :param accelerator: Accelerator for distributed computing, defaults to None
@@ -105,7 +104,7 @@ class GRPO(LLMAlgorithm):
         pad_token_id: int,
         pad_token: str,
         model_name: str | None = None,
-        actor_network: PreTrainedModel | None = None,
+        actor_network: PreTrainedModelProtocol | None = None,
         model_config: dict[str, Any] | None = None,
         hp_config: Optional[HyperparameterConfig] = None,
         index: int = 0,
@@ -127,7 +126,7 @@ class GRPO(LLMAlgorithm):
         max_output_tokens: int | None = 1024,
         min_output_tokens: Optional[int] = None,
         max_model_len: Optional[int] = None,
-        lora_config: Optional[LoraConfig] = None,
+        lora_config: Optional[LoraConfigProtocol] = None,
         cosine_lr_schedule_config: Optional[CosineLRScheduleConfig] = None,
         accelerator: Optional[Accelerator] = None,
         device: str = "cpu",
@@ -188,8 +187,8 @@ class GRPO(LLMAlgorithm):
         ), "Policy update epochs must be greater than or equal to one."
         if actor_network is not None:
             assert isinstance(
-                actor_network, (PeftModel, PreTrainedModel)
-            ), "Actor network must be a PeftModel or PreTrainedModel"
+                actor_network, (PeftModelProtocol, PreTrainedModelProtocol)
+            ), "Actor network must be a PeftModelProtocol or PreTrainedModelProtocol"
 
         self.clip_coef = clip_coef
         self.update_epochs = update_epochs

{agilerl-2.4.1.dev0 → agilerl-2.4.1.dev2}/agilerl/algorithms/ilql.py

@@ -1223,6 +1223,20 @@ class ILQL(nn.Module):
         self.fitness = checkpoint["fitness"]
         self.steps = checkpoint["steps"]
 
+    def clean_up(self) -> None:
+        """Clean up the networks"""
+        del self.model
+        del self.actor
+        del self.actor_target
+        del self.v
+        del self.q
+        del self.target_q
+        del self.pi
+        del self.optimizer
+        if self.double_q:
+            del self.q2
+            del self.target_q2
+
 
 class ILQL_Policy:
     def __init__(self, iql_model: ILQL, kind: str, **generation_kwargs) -> None: