project-llm-trainer 0.4.15__py3-none-any.whl → 0.5.1__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -6,35 +6,11 @@ from torch.optim import Optimizer
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from .parallel_ds import DsParallel
-from .parallel_fsdp import FsdpParallel
-from .parallel_ddp import DdpParallel
 from .scheduler import LRScheduler
 from .tools import TrainerTools
 
-try:
-    from .dcp import save_dcp, load_dcp, convert_dcp_to_pth
-except:
-    os.environ['ENABLE_DCP'] = "0"
-
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-
-# https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html
-
 DEFAULT_CHECKPOINT_NAME = "checkpoint.pth"
 
-
-def _can_use_dcp(model: nn.Module) -> bool:
-    if os.environ.get('ENABLE_DCP', '1') != '1':
-        return False
-
-    # 如果是fsdp或者ddp，才能使用dcp保存
-    if (isinstance(TrainerTools().parallel, FsdpParallel)
-            or isinstance(TrainerTools().parallel, DdpParallel)):
-        return True
-
-    return False
-
-
 def save_checkpoint(
         model: nn.Module,
         optimizer: Optional[Optimizer] = None,
@@ -43,11 +19,6 @@ def save_checkpoint(
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import save_ds_checkpoint
         save_ds_checkpoint(model, suffix)
-    elif _can_use_dcp(model):
-        save_dcp(model, optimizer, suffix)
-    elif isinstance(model, FSDP):
-        from .fsdp_checkpoint import save_fsdp_checkpoint
-        save_fsdp_checkpoint(model, optimizer, suffix)
     else:
         if TrainerTools().parallel.is_main_process:
             checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
@@ -73,11 +44,6 @@ def load_checkpoint(
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import load_ds_checkpoint
         load_ds_checkpoint(model, load_module_only=load_module_only, suffix=suffix)
-    elif _can_use_dcp(model):
-        load_dcp(model, optimizer, suffix)
-    elif isinstance(model, FSDP):
-        from .fsdp_checkpoint import load_fsdp_checkpoint
-        load_fsdp_checkpoint(model, optimizer, device, suffix)
     else:
         checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
         if suffix:
@@ -99,22 +65,6 @@ def load_checkpoint_for_eval(
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import load_ds_checkpoint_for_eval
         load_ds_checkpoint_for_eval(model)
-    elif _can_use_dcp(model):
-        checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-
-        # load_dcp方式在cpu上会报错，所以改为先将ckpt转换为pth，然后再加载pth
-        # load_dcp(model, optimizer)
-        pth_name = os.environ.get('EVAL_CHECKPOINT_NAME', checkpoint_name)
-        if suffix:
-            pth_name = f'{pth_name}_{suffix}'
-
-        convert_dcp_to_pth(pth_name)
-
-        if os.path.exists(pth_name):
-            ckpt = torch.load(pth_name, map_location=device, weights_only=True)
-            model.load_state_dict(ckpt['app']['model_state_dict'])
-            # 使用完删除
-            os.remove(pth_name)
     else:
         load_checkpoint(model, None, device, suffix=suffix)
 

llm_trainer/dpo_trainer.py

@@ -12,7 +12,7 @@ from .dataset import DPODataset
 from .loss import DPOLoss
 from .tools import TrainerTools
 from .utils import get_dpo_collate_fn
-from .model_params import copy_model_params
+from .partition_utils import sync_model_params
 
 from .checkpoint import (
     save_checkpoint,
@@ -38,7 +38,6 @@ class DPOTrainer(Trainer):
 
     def _init_reference_model(self):
         reference_model = self._new_model(self.train_config)
-        copy_model_params(_from=self.train_model, _to=reference_model)
 
         reference_model, _ = TrainerTools().parallel.process(
             model=reference_model,
@@ -51,6 +50,11 @@ class DPOTrainer(Trainer):
         for param in reference_model.parameters():
             param.requires_grad = False
 
+        sync_model_params(
+            _from=self.train_model,
+            _to=reference_model
+        )
+
         return reference_model
 
     def _init_loss(self):
@@ -210,7 +214,6 @@ class DPOTrainer(Trainer):
                         if need_update_grad:
                             loss_tensor = torch.tensor(loss_accumulation, device=TrainerTools().parallel.device)
 
-                            # todo check all_reduce??
                             if TrainerTools().parallel.parallel_train:
                                 dist.all_reduce(loss_tensor, dist.ReduceOp.AVG)
 

llm_trainer/eval.py

@@ -5,16 +5,14 @@ from .log import get_log_dir
 from .tools import TrainerTools
 from .train_configs import EvalConfig
 
-
-def _eval_task(
+def submit_gen_task(
         eval_model: torch.nn.Module,
         eval_config: EvalConfig,
         tag,
         prompt,
         pixel_values,
         max_position_embeddings,
-        tokens_per_image,
-        device
+        tokens_per_image
 ):
     log_dir = get_log_dir()
 
@@ -28,33 +26,8 @@ def _eval_task(
         p=eval_config.top_p,
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
-        device=device
+        device=TrainerTools().parallel.device
     )
 
     with open(f'{log_dir}gen.txt', 'a') as f:
         f.write(f"{tag}, gen->{gen_result}\n")
-
-
-def submit_gen_task(
-        eval_model: torch.nn.Module,
-        eval_config: EvalConfig,
-        tag,
-        prompt,
-        pixel_values,
-        max_position_embeddings,
-        tokens_per_image
-):
-    eval_model.to(TrainerTools().parallel.device)
-    _eval_task(
-        eval_model=eval_model,
-        eval_config=eval_config,
-        tag=tag,
-        prompt=prompt,
-        pixel_values=pixel_values,
-        max_position_embeddings=max_position_embeddings,
-        tokens_per_image=tokens_per_image,
-        device=TrainerTools().parallel.device
-    )
-    eval_model.to('cpu')
-
-    # threading.Thread(target=_eval_task, args=args).start()

llm_trainer/generate_utils.py

@@ -1,7 +1,6 @@
 from typing import Union, Optional, List
 from contextlib import nullcontext
 import torch
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from llm_model import VlmModel, KVCache
 from .tools import TrainerTools
 from .utils import batch_repeat_image_tok
@@ -108,8 +107,7 @@ def _generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int],
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int]
 ):
     """
     :param model:
@@ -131,8 +129,7 @@ def _generate(
         device_type=device,
         dtype=TrainerTools().dtype,
         enabled=True,
-        # fsdp模式，需要将cache_enabled设置为false
-        cache_enabled=False if isinstance(model, FSDP) else None
+        cache_enabled=None
     ) if TrainerTools().use_amp else nullcontext()
 
     if isinstance(model, VlmModel):
@@ -144,28 +141,6 @@ def _generate(
     kv_cache: Optional[KVCache] = None
     generate_tokens = tokens.clone()
 
-    reasoning_start = TrainerTools().tokenizer.reasoning_start
-    reasoning_end = TrainerTools().tokenizer.reasoning_end
-
-    # --- 状态初始化 ---
-    in_reasoning_block = False
-    reasoning_step_count = 0
-    # “冷静期”标志位。当强制结束思考后，在下一步抑制<reasoning>的生成。
-    suppress_reasoning_start_next = False
-
-    if reasoning_budget is not None:
-        prompt_tokens = tokens[0]
-        start_indices = (prompt_tokens == reasoning_start).nonzero(as_tuple=True)[0]
-        end_indices = (prompt_tokens == reasoning_end).nonzero(as_tuple=True)[0]
-
-        last_start_idx = start_indices[-1].item() if len(start_indices) > 0 else -1
-        last_end_idx = end_indices[-1].item() if len(end_indices) > 0 else -1
-
-        if last_start_idx > last_end_idx:
-            in_reasoning_block = True
-            reasoning_step_count = len(prompt_tokens) - 1 - last_start_idx
-
-    model.eval()
     with torch.inference_mode():
         for _ in range(max_new_tokens):
             # 是否需要截取？？
@@ -185,23 +160,6 @@ def _generate(
             # (batch, vocab_size)
             logits = logits[:, -1, :]
 
-            # --- 推理预算逻辑 ---
-            force_end_reasoning_token = False
-            if reasoning_budget is not None:
-                # 检查是否需要在此步抑制 <reasoning>
-                should_suppress_this_step = suppress_reasoning_start_next
-                suppress_reasoning_start_next = False  # 立即重置标志位
-
-                # 修改: 检查是否超出预算
-                if in_reasoning_block and reasoning_step_count >= reasoning_budget:
-                    force_end_reasoning_token = True
-                    # 设置标志位，在下一步抑制 <reasoning>
-                    suppress_reasoning_start_next = True
-
-                # 如果上一轮设置了抑制标志，则在此轮执行抑制
-                if should_suppress_this_step:
-                    logits[:, reasoning_start] = -float("inf")
-
             # 抑制特殊token输出
             if suppress_tokens and len(suppress_tokens) != 0:
                 logits = _suppress_warper(logits, suppress_tokens)
@@ -217,10 +175,6 @@ def _generate(
             if p and 0 < p <= 1:
                 logits = _top_p_warper(logits, p)
 
-            if force_end_reasoning_token:
-                logits[:] = -float("inf")
-                logits[:, reasoning_end] = 0.0
-
             if multinomial:
                 prob = logits.softmax(dim=-1)
                 # 返回下标
@@ -229,18 +183,6 @@ def _generate(
                 # 返回下标
                 next_token = logits.argmax(dim=-1, keepdim=True)
 
-            if reasoning_budget is not None:
-                current_token_id = next_token.item()
-                if not in_reasoning_block and current_token_id == reasoning_start:
-                    in_reasoning_block = True
-                    reasoning_step_count = 0
-                elif in_reasoning_block:
-                    if current_token_id == reasoning_end:
-                        in_reasoning_block = False
-                        reasoning_step_count = 0
-                    else:
-                        reasoning_step_count += 1
-
             # token, is_full_result
             yield next_token, False
 
@@ -269,8 +211,7 @@ def _streaming_generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int] = None,
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int] = None
 ):
     device = TrainerTools().parallel.device if not device else device
     encoded_tokens = TrainerTools().tokenizer.encode(prompt, unsqueeze=True, covert_tensor=True).to(device)
@@ -286,8 +227,7 @@ def _streaming_generate(
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
         suppress_tokens=suppress_tokens,
-        device=device,
-        reasoning_budget=reasoning_budget
+        device=device
     )
 
     for (token, is_full_result) in generate_text_iterator:
@@ -306,8 +246,7 @@ def streaming_generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int] = None,
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int] = None
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -320,8 +259,7 @@ def streaming_generate(
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
         suppress_tokens=suppress_tokens,
-        device=device,
-        reasoning_budget=reasoning_budget
+        device=device
     )
 
     for (token, is_full_result) in text_iterator:
@@ -341,8 +279,7 @@ def generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int] = None,
-        reasoning_budget: Optional[int] = None
+        device: Union[str, torch.device, int] = None
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -355,8 +292,7 @@ def generate(
         suppress_tokens=suppress_tokens,
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
-        device=device,
-        reasoning_budget=reasoning_budget
+        device=device
     )
 
     for (token, is_full_result) in text_iterator:
@@ -386,7 +322,7 @@ def batch_generate(
         device_type=device,
         dtype=TrainerTools().dtype,
         enabled=True,
-        cache_enabled=False if isinstance(model, FSDP) else None
+        cache_enabled=None
     ) if TrainerTools().use_amp else nullcontext()
 
     if isinstance(model, VlmModel):
@@ -403,7 +339,6 @@ def batch_generate(
     end_token = TrainerTools().tokenizer.end
     done = torch.zeros(batch_size, dtype=torch.bool, device=device)
 
-    model.eval()
     with torch.inference_mode():
         for _ in range(max_new_tokens):
             # 只处理未完成的样本

llm_trainer/grpo_trainer.py

@@ -1,5 +1,4 @@
 import time
-import copy
 from typing import Tuple, List, Union, Callable, Optional
 import torch
 from torch.utils.data import Dataset
@@ -15,7 +14,11 @@ from .loss import GRPOLoss
 from .tools import TrainerTools
 from .generate_utils import batch_generate
 from .log import log
-from .model_params import copy_model_params
+
+from .partition_utils import (
+    sync_model_params,
+    unwrap_model_for_generation
+)
 
 from .checkpoint import (
     save_checkpoint,
@@ -39,7 +42,6 @@ class GRPOTrainer(Trainer):
 
         self.reward_func = reward_func
         self.reference_model = self._init_reference_model()
-        self.generate_model = self._init_generate_model()
 
         # 默认使用torch提供的pad_sequence
         # 如果pad_sequence不支持padding_side参数，则将改参数置为False，使用反转的方式
@@ -47,17 +49,20 @@ class GRPOTrainer(Trainer):
 
     def _init_reference_model(self):
         reference_model = self._new_model(self.train_config)
-        reference_model.to('cpu')
-        reference_model.eval()
 
+        reference_model, _ = TrainerTools().parallel.process(
+            model=reference_model,
+            optimizer=None,
+            kwargs=self._init_reference_args(),
+            save_instance=False
+        )
+
+        reference_model.eval()
         for param in reference_model.parameters():
             param.requires_grad = False
 
         return reference_model
 
-    def _init_generate_model(self):
-        return copy.deepcopy(self.reference_model)
-
     def _init_loss(self):
         criterion = GRPOLoss(
             clip_eps=self.train_config.grpo_config.clip_eps,
@@ -163,7 +168,7 @@ class GRPOTrainer(Trainer):
         # [batch*group_size, 1]
         return advantages.unsqueeze(1)  # Add dimension for token-wise operations
 
-    def _generate_completions(self, prompts, group_size: int):
+    def _generate_completions(self, model, prompts, group_size: int):
         pad_token_id = TrainerTools().tokenizer.pad
         device = TrainerTools().parallel.device
 
@@ -181,7 +186,7 @@ class GRPOTrainer(Trainer):
 
         # [batch*group_size, max_prompt_len+max_gen_len]
         outputs: torch.Tensor = batch_generate(
-            model=self.generate_model,
+            model=model,
             tokens=prompt_ids,
             pad_token_id=pad_token_id,
             attention_mask=prompt_masks,
@@ -201,7 +206,7 @@ class GRPOTrainer(Trainer):
 
         return prompt_ids, prompt_masks, completion_ids, completion_masks
 
-    def _generate_rollout_data(self, batch_data: List[dict]):
+    def _generate_rollout_data(self, generate_model, batch_data: List[dict]):
         prompts = [item["prompt"] for item in batch_data]
         answers = [item["answer"] for item in batch_data]
         group_size = self.train_config.grpo_config.group_size
@@ -210,13 +215,13 @@ class GRPOTrainer(Trainer):
         # 修复问题：Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal
         with torch.no_grad():
         # with torch.inference_mode():
-            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_completions(prompts, group_size)
+            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_completions(generate_model, prompts, group_size)
             input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
             attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
             logits_to_keep = completion_ids.shape[1]
 
             # Compute old_log_probs from the current model, with gradients disabled.
-            old_log_probs, _ = self._compute_log_probabilities(self.generate_model, input_ids, attention_mask, logits_to_keep)
+            old_log_probs, _ = self._compute_log_probabilities(generate_model, input_ids, attention_mask, logits_to_keep)
 
             # Compute ref_log_probs from the reference model, which remains static.
             ref_log_probs, _ = self._compute_log_probabilities(self.reference_model, input_ids, attention_mask, logits_to_keep)
@@ -275,12 +280,15 @@ class GRPOTrainer(Trainer):
     def train(self):
         global_steps = 0
         skipping_train = False
-        device = TrainerTools().parallel.device
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
 
         for epoch in range(self.train_config.n_epochs):
-            copy_model_params(_from=self.train_model, _to=self.reference_model)
-            self.train_model.train()
+            sync_model_params(
+                _from=self.train_model,
+                _to=self.reference_model,
+                mixup_alpha=self.train_config.grpo_config.mixup_alpha
+            )
+
             file_count = len(self.train_config.file_dataset)
 
             for file_idx in range(file_count):
@@ -307,22 +315,13 @@ class GRPOTrainer(Trainer):
                     skipping_train = False
 
                     # start generate
-                    # 使用单独的模型生成数据， 原因是在deepspeed并行训练时，使用train_model生成数据会卡死
-                    self.generate_model.to(device)
-                    self.reference_model.to(device)
-
                     if TrainerTools().parallel.is_main_process:
                         log(f'start generate for batch {batch}/{batch_count_per_file}')
 
                     # 生成数据
-                    with torch.no_grad():
-                        # 保存了train_model checkpoint后，这里保证生成模型使用的参数是最新
-                        copy_model_params(_from=self.train_model, _to=self.generate_model)
-                        rollout_data = self._generate_rollout_data(batch_data)
-
-                    # 卸载到cpu上，等待下次使用时再to gpu
-                    self.generate_model.to('cpu')
-                    self.reference_model.to('cpu')
+                    with unwrap_model_for_generation(self.train_model) as generate_model:
+                        rollout_data = self._generate_rollout_data(generate_model, batch_data)
+
                     torch.cuda.empty_cache()
                     # end generate
 

llm_trainer/loss.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional
 import torch
 from torch import nn
 import torch.nn.functional as F

llm_trainer/partition_utils.py

@@ -0,0 +1,146 @@
+from typing import Optional
+from contextlib import contextmanager
+import itertools
+from packaging import version
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .tools import TrainerTools
+from .parallel_ds import DsParallel
+from .parallel_ddp import DdpParallel
+
+
+@contextmanager
+def unwrap_model_for_generation(model: nn.Module):
+    """
+    Context manager to unwrap distributed or accelerated models for generation tasks.
+
+    Args:
+        model:
+            Model to be unwrapped.
+    Yields:
+        Unwrapped model.
+
+    Example:
+    ```python
+    with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
+        generated_outputs = unwrapped_model.generate(input_ids)
+    ```
+    """
+    if isinstance(TrainerTools().parallel, DsParallel):
+        import deepspeed
+        assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+        if model.zero_optimization_stage() == 3:
+            with deepspeed.zero.GatheredParameters(model.parameters()):
+                _remove_hooks(model)
+                yield unwrap_model(model)
+                _add_hooks(model)
+        else:
+            yield unwrap_model(model)
+    elif isinstance(TrainerTools().parallel, DdpParallel):
+        yield unwrap_model(model)
+    else:
+        yield model
+
+
+def sync_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    if isinstance(TrainerTools().parallel, DsParallel):
+        _sync_ds_model_params(_from, _to, mixup_alpha)
+    elif isinstance(TrainerTools().parallel, DdpParallel):
+        _sync_ddp_model_params(_from, _to, mixup_alpha)
+    else:
+        _copy_params(_from, _to, mixup_alpha)
+
+
+def unwrap_model(model) -> nn.Module:
+    try:
+        import deepspeed
+        if isinstance(model, deepspeed.DeepSpeedEngine):
+            return model.module
+    except: ...
+
+    if isinstance(model, DDP):
+        return model.module
+
+    return model
+
+
+def _copy_params(model, target_model, mixup_alpha):
+    for target_param, copy_param in zip(target_model.parameters(), model.parameters()):
+        target_param.data.mul_(1.0 - mixup_alpha).add_(copy_param.data, alpha=mixup_alpha)
+
+
+def _sync_ds_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    import deepspeed
+    assert isinstance(_from, deepspeed.DeepSpeedEngine)
+
+    origin_from = unwrap_model(_from)
+
+    if _from.zero_optimization_stage() == 3:
+        with deepspeed.zero.GatheredParameters(list(origin_from.parameters()) + list(_to.parameters()), modifier_rank=0):
+            if TrainerTools().parallel.is_main_process:
+                _copy_params(origin_from, _to, mixup_alpha)
+    else:
+        _copy_params(origin_from, _to, mixup_alpha)
+
+
+def _sync_ddp_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    assert isinstance(_from, DDP)
+
+    origin_from = unwrap_model(_from)
+    _copy_params(origin_from, _to, mixup_alpha)
+
+
+def _add_hooks(model: nn.Module) -> None:
+    """Adds the optimizer hooks from a DeepSpeed ZeRO-3 model."""
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+    if not hasattr(model, "optimizer"):  # before the first training step, the model has no optimizer
+        return
+    if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"):
+        optimizer_offload = model.optimizer.parameter_offload
+    elif model.optimizer is not None:
+        optimizer_offload = model.optimizer
+    else:
+        raise RuntimeError("The model optimizer is None, which is not yet supported.")
+    if version.parse(deepspeed.__version__) >= version.parse("0.16.4"):
+        # Account for renaming in https://github.com/deepspeedai/DeepSpeed/pull/6847
+        optimizer_offload._register_deepspeed_module(optimizer_offload.module)
+    else:
+        optimizer_offload._register_hooks_recursively(optimizer_offload.module)
+
+
+def _remove_hooks(model: nn.Module) -> None:
+    """Removes the optimizer hooks from a DeepSpeed ZeRO-3 model."""
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+    if not hasattr(model, "optimizer"):  # before the first training step, the model has no optimizer
+        return
+    if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"):
+        optimizer_offload = model.optimizer.parameter_offload
+    elif model.optimizer is not None:
+        optimizer_offload = model.optimizer
+    else:
+        raise RuntimeError("The model optimizer is None, which is not yet supported.")
+
+    for param in _iter_params(optimizer_offload.module, recurse=True):
+        param.ds_active_sub_modules.clear()
+
+    for hook in optimizer_offload.forward_hooks:
+        hook.remove()
+    for hook in optimizer_offload.backward_hooks:
+        hook.remove()
+
+    optimizer_offload.forward_hooks = []
+    optimizer_offload.backward_hooks = []
+
+
+def _iter_params(module, recurse=False):
+    return [param for _, param in _get_all_parameters(module, recurse)]
+
+
+def _get_all_parameters(sub_module, recurse=False):
+    return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters())