agilerl 2.4.1__tar.gz → 2.4.1.dev0__tar.gz

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

@@ -1,23 +1,20 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.3
 Name: agilerl
-Version: 2.4.1
+Version: 2.4.1.dev0
 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,<3.13
+Requires-Python: >=3.10,<4.0
 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
-Provides-Extra: all
-Provides-Extra: llm
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: SuperSuit (>=3.9.0,<4.0.0)
 Requires-Dist: accelerate (>=1.7.0,<2.0.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: deepspeed (>=0.17.1,<0.18.0)
 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)
@@ -27,12 +24,11 @@ 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[all] (==0.5.2)
+Requires-Dist: minari (>=0.5.2,<0.6.0)
 Requires-Dist: numpy (>=1.26.4,<2.0.0)
 Requires-Dist: omegaconf (>=2.3.0,<3.0.0)
-Requires-Dist: packaging (>=20.0)
 Requires-Dist: pandas (>=2.2.3,<3.0.0)
-Requires-Dist: peft (>=0.18.0,<0.19.0) ; extra == "llm" or extra == "all"
+Requires-Dist: peft (>=0.15.2,<0.16.0)
 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)
@@ -43,9 +39,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.57.1,<5.0.0) ; extra == "llm" or extra == "all"
+Requires-Dist: transformers (>=4.48.1,<5.0.0)
 Requires-Dist: ucimlrepo (>=0.0.3,<0.0.4)
-Requires-Dist: vllm (==0.10.0) ; extra == "llm" or extra == "all"
+Requires-Dist: vllm (==0.10.0)
 Requires-Dist: wandb (>=0.17.6,<0.18.0)
 Description-Content-Type: text/markdown
 
@@ -99,16 +95,6 @@ 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
@@ -167,12 +153,11 @@ 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 Fine-tuning Algorithms
+  ### LLM Reasoning 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 → agilerl-2.4.1.dev0}/README.md

@@ -48,16 +48,6 @@ 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
@@ -116,12 +106,11 @@ 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 Fine-tuning Algorithms
+  ### LLM Reasoning 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 → agilerl-2.4.1.dev0}/agilerl/algorithms/core/base.py

@@ -27,14 +27,20 @@ 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,
@@ -49,11 +55,7 @@ from agilerl.protocols import (
     EvolvableAttributeDict,
     EvolvableAttributeType,
     EvolvableModule,
-    LoraConfigProtocol,
     ModuleDict,
-    PeftModelProtocol,
-    PretrainedConfigProtocol,
-    PreTrainedModelProtocol,
 )
 from agilerl.typing import (
     ActionType,
@@ -72,7 +74,6 @@ from agilerl.typing import (
 )
 from agilerl.utils.algo_utils import (
     CosineLRScheduleConfig,
-    DummyOptimizer,
     VLLMConfig,
     check_supported_space,
     chkpt_attribute_to_device,
@@ -95,18 +96,11 @@ from agilerl.utils.evolvable_networks import (
     is_image_space,
     is_vector_space,
 )
-
-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,
-    )
+from agilerl.utils.llm_utils import (
+    DummyOptimizer,
+    create_model_from_name_or_path,
+    gather_if_zero3,
+)
 
 __all__ = ["EvolvableAlgorithm", "RLAlgorithm", "MultiAgentRLAlgorithm"]
 
@@ -607,16 +601,14 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
         )
         optimizer = opt.optimizer if hasattr(opt, "optimizer") else None
 
-        if isinstance(self, LLMAlgorithm):
-            if hasattr(self.actor, "optimizer"):
-                optimizer = getattr(
+        if isinstance(opt, DeepSpeedOptimizerWrapper):
+            if isinstance(opt.optimizer, DummyOptimizer):
+                opt = 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(
-                optimizer,
+                opt,
                 lr=getattr(self, config.lr),
                 accelerator=self.accelerator,
                 scheduler_config=self.cosine_lr_schedule_config,
@@ -1151,16 +1143,6 @@ 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.
@@ -1817,10 +1799,6 @@ 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__(
@@ -1835,10 +1813,10 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         seed: int,
         pad_token_id: int,
         pad_token: str,
-        lora_config: LoraConfigProtocol | None,
+        lora_config: LoraConfig | None,
         use_separate_reference_adapter: bool,
         model_name: str | None = None,
-        actor_network: PreTrainedModelProtocol | None = None,
+        actor_network: PreTrainedModel | None = None,
         micro_batch_size_per_gpu: int | None = None,
         cosine_lr_schedule_config: Optional[CosineLRScheduleConfig] = None,
         hp_config: Optional[HyperparameterConfig] = None,
@@ -1846,14 +1824,9 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         device: Union[str, torch.device] = "cpu",
         accelerator: Optional[Accelerator] = None,
         name: Optional[str] = None,
-        model_config: dict[str, Any] | PretrainedConfigProtocol | None = None,
+        model_config: dict[str, Any] | PretrainedConfig | 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."
@@ -1908,7 +1881,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                     )
                     lr = optim_lr
 
-        if lora_config is None and not isinstance(actor_network, PeftModelProtocol):
+        if lora_config is None and not isinstance(actor_network, PeftModel):
             warnings.warn(
                 "No LoRA config provided. AgileRL can only be used to finetune adapters at present. Using default LoRA configuration for RL finetuning."
             )
@@ -1925,21 +1898,15 @@ 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):
-            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
+        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
         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"
@@ -2074,7 +2041,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 device_map="auto"
             )
             tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
-            model = PeftModelProtocol.from_pretrained(base_model, path)
+            model = PeftModel.from_pretrained(base_model, path)
             """
         )
 
@@ -2186,11 +2153,6 @@ 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,
@@ -2214,8 +2176,10 @@ 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):
@@ -2250,8 +2214,8 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             input_args["wrap"] = False
             input_args["clone"] = True
 
-            actor: PeftModelProtocol = cast(
-                PeftModelProtocol,
+            actor: PeftModel = cast(
+                PeftModel,
                 (
                     self.accelerator.unwrap_model(self.actor)
                     if self.accelerator is not None
@@ -2443,12 +2407,12 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             self.reference_update_tracker += 1
 
     def _initialize_actors(
-        self, base_model: PreTrainedModelProtocol | None, add_adapters: bool = True
+        self, base_model: PreTrainedModel | None, add_adapters: bool = True
     ):
         """Initialize the actor network.
 
         :param base_model: Base model
-        :type base_model: PreTrainedModelProtocol
+        :type base_model: PreTrainedModel
         :param add_adapters: Flag to indicate if adapters should be added to the model, defaults to True
         :type add_adapters: bool, optional
         """
@@ -2458,7 +2422,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 self.pretrained_model_name_or_path
             )
 
-        if isinstance(base_model, PeftModelProtocol) and add_adapters:
+        if isinstance(base_model, PeftModel) 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())
@@ -2468,7 +2432,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             if "default" in list(base_model.peft_config.keys()):
                 base_model.peft_config.pop("default")
 
-        self.actor: PeftModelProtocol = (
+        self.actor: PeftModel = (
             get_peft_model(base_model, self.lora_config, adapter_name="actor")
             if add_adapters
             else base_model
@@ -2617,6 +2581,7 @@ 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)
@@ -2984,28 +2949,3 @@ 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 → agilerl-2.4.1.dev0}/agilerl/algorithms/core/optimizer_wrapper.py

@@ -2,27 +2,19 @@ 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.algo_utils import DummyOptimizer
-
-if HAS_LLM_DEPENDENCIES:
-    from peft import PeftModel
-
-    PeftModelType = PeftModel
-else:
-    PeftModelType = "PeftModel"
-
+from agilerl.utils.llm_utils import DummyOptimizer
 
 ModuleList = list[EvolvableModule]
 _Optimizer = Union[
     type[OptimizerType], dict[str, type[OptimizerType]], type[DummyOptimizer]
 ]
-_Module = Union[EvolvableModule, ModuleDict, ModuleList, PeftModelType]
+_Module = Union[EvolvableModule, ModuleDict, ModuleList, PeftModel]
 
 
 def init_from_multiple(

{agilerl-2.4.1 → agilerl-2.4.1.dev0}/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.algo_utils import DummyOptimizer
+from agilerl.utils.llm_utils import DummyOptimizer
 
 
 @dataclass

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

@@ -5,10 +5,11 @@ 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
@@ -24,7 +25,7 @@ class DPO(LLMAlgorithm):
     :param model_name: Model name
     :type model_name: str, optional
     :param actor_network: HuggingFace LLM
-    :type actor_network: PreTrainedModelProtocol
+    :type actor_network: PreTrainedModel
     :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
@@ -49,7 +50,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: LoraConfigProtocol, optional
+    :type lora_config: LoraConfig, 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
@@ -69,7 +70,7 @@ class DPO(LLMAlgorithm):
         pad_token_id: int,
         pad_token: str,
         model_name: str | None = None,
-        actor_network: PreTrainedModelProtocol | None = None,
+        actor_network: PreTrainedModel | None = None,
         model_config: dict[str, Any] | None = None,
         hp_config: HyperparameterConfig | None = None,
         index: int = 0,
@@ -82,7 +83,7 @@ class DPO(LLMAlgorithm):
         micro_batch_size_per_gpu: int | None = None,
         reduce_memory_peak: bool = False,
         device: str = "cpu",
-        lora_config: LoraConfigProtocol | None = None,
+        lora_config: LoraConfig | None = None,
         accelerator: Accelerator | None = None,
         wrap: bool = True,
         clone: bool = False,

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

@@ -1,18 +1,17 @@
 import gc
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 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,
@@ -24,8 +23,10 @@ from agilerl.utils.llm_utils import (
     ReasoningGym,
 )
 
-if HAS_LLM_DEPENDENCIES:
-    from transformers import GenerationConfig
+DeepSpeedOptimizerType = Union[
+    DeepSpeedZeroOptimizer,  # ZeRO Stage 1 & 2 optimizer
+    DeepSpeedZeroOptimizer_Stage3,  # ZeRO Stage 3 optimizer
+]
 
 
 class GRPO(LLMAlgorithm):
@@ -38,7 +39,7 @@ class GRPO(LLMAlgorithm):
     :param model_name: Model name
     :type model_name: str, optional
     :param actor_network: HuggingFace LLM
-    :type actor_network: PreTrainedModelProtocol
+    :type actor_network: PreTrainedModel
     :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.
@@ -76,7 +77,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: LoraConfigProtocol, optional
+    :type lora_config: LoraConfig, 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
@@ -104,7 +105,7 @@ class GRPO(LLMAlgorithm):
         pad_token_id: int,
         pad_token: str,
         model_name: str | None = None,
-        actor_network: PreTrainedModelProtocol | None = None,
+        actor_network: PreTrainedModel | None = None,
         model_config: dict[str, Any] | None = None,
         hp_config: Optional[HyperparameterConfig] = None,
         index: int = 0,
@@ -126,7 +127,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[LoraConfigProtocol] = None,
+        lora_config: Optional[LoraConfig] = None,
         cosine_lr_schedule_config: Optional[CosineLRScheduleConfig] = None,
         accelerator: Optional[Accelerator] = None,
         device: str = "cpu",
@@ -187,8 +188,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, (PeftModelProtocol, PreTrainedModelProtocol)
-            ), "Actor network must be a PeftModelProtocol or PreTrainedModelProtocol"
+                actor_network, (PeftModel, PreTrainedModel)
+            ), "Actor network must be a PeftModel or PreTrainedModel"
 
         self.clip_coef = clip_coef
         self.update_epochs = update_epochs

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

@@ -1223,20 +1223,6 @@ 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: