project-llm-trainer 0.12.3__py3-none-any.whl

llm_trainer/loss.py

@@ -0,0 +1,268 @@
+from typing import List, Optional
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class LMLoss(nn.Module):
+    """
+    llm loss
+    """
+    def __init__(
+            self,
+            ignore_index: int = -100,
+            *,
+            critical_tokens: Optional[List[int]] = None,
+            critical_alpha: float = 1.0,
+            vocab_size: int = 0
+    ):
+        super().__init__()
+        self.ignore_index = ignore_index
+        self.critical_tokens = critical_tokens
+        self.critical_alpha = critical_alpha
+
+        if critical_tokens and vocab_size > 0:
+            self.register_buffer('weights', torch.ones(vocab_size))
+            # 为关键token设置权重
+            self.weights[self.critical_tokens] = critical_alpha
+
+
+    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        # logits shape (batch, seq_len, vocab_size)
+        # labels shape (batch, seq_len)
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        logits = shift_logits.reshape(-1, logits.shape[-1])
+        targets = shift_labels.reshape(-1)
+
+        ce_loss = F.cross_entropy(
+            logits,
+            targets,
+            ignore_index=self.ignore_index,
+            weight=self.weights.to(logits.device, dtype=logits.dtype) if self.critical_tokens else None
+        )
+
+        # 添加额外惩罚项（可选）
+        # if self.critical_tokens:
+        #     crit_mask = torch.isin(targets, torch.tensor(self.critical_tokens).to(targets.device))
+        #     crit_logits = logits[crit_mask]
+        #     crit_targets = targets[crit_mask]
+        #     extra_loss = F.cross_entropy(crit_logits, crit_targets, ignore_index=self.ignore_index)
+        #     return ce_loss + extra_loss * (self.critical_alpha - 1)  # 增强惩罚
+
+        return ce_loss
+
+
+class KDLoss(nn.Module):
+    """
+    Language Model Knowledge Distillation Loss
+    https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/loss.py#L266
+    """
+    def __init__(self, ignore_index: int = -100):
+        super().__init__()
+        self.ignore_index = ignore_index
+
+    def forward(self, logits: torch.Tensor, teacher_logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        teacher_probs = F.softmax(teacher_logits, dim=-1, dtype=torch.float32)
+        inf_mask = torch.isinf(logits)
+
+        logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
+        prod_probs = torch.masked_fill(teacher_probs * logprobs, inf_mask, 0)
+
+        x = torch.sum(prod_probs, dim=-1).view(-1)
+        mask = (labels != self.ignore_index).int()
+
+        distil_loss = -torch.sum(x * mask.view(-1), dim=0) / torch.sum(mask.view(-1), dim=0)
+
+        return distil_loss
+
+
+class DPOLoss(nn.Module):
+    def __init__(
+            self,
+            beta: float,
+            label_smoothing: float = 0.0,
+            ipo: bool = False
+    ):
+        super().__init__()
+        self.beta = beta
+        self.label_smoothing = label_smoothing
+        self.ipo = ipo
+
+    def forward(
+            self,
+            policy_chosen_logps: torch.Tensor,
+            policy_reject_logps: torch.Tensor,
+            ref_chosen_logps: torch.Tensor,
+            ref_reject_logps: torch.Tensor
+    ) -> torch.Tensor:
+        pi_logratios = policy_chosen_logps - policy_reject_logps
+        ref_logratios = ref_chosen_logps - ref_reject_logps
+        logits = pi_logratios - ref_logratios
+
+        if self.ipo:
+            losses = (logits - 1 / (2 * self.beta)) ** 2  # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf
+        else:
+            # Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf)
+            losses = (
+                    -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
+                    - F.logsigmoid(-self.beta * logits) * self.label_smoothing
+            )
+
+        loss = losses.mean()
+
+        # chosen_rewards = self.beta * (policy_chosen_probs - ref_chosen_probs).detach()
+        # rejected_rewards = self.beta * (policy_reject_probs - ref_reject_probs).detach()
+
+        return loss
+
+
+class PPOLoss(nn.Module):
+    """
+    PPO (Proximal Policy Optimization) 损失函数。
+    这个类统一计算 Actor 和 Value 的损失。
+    """
+
+    def __init__(
+            self,
+            clip_eps: float,
+            vf_coef: float,
+    ):
+        """
+        初始化PPO损失函数。
+        :param clip_eps: PPO裁剪范围的epsilon值。
+        :param vf_coef: 价值函数损失的系数。
+        """
+        super().__init__()
+        self.clip_eps = clip_eps
+        self.vf_coef = vf_coef
+
+    def forward(
+            self,
+            log_probs: torch.Tensor,
+            old_log_probs: torch.Tensor,
+            values: torch.Tensor,
+            old_values: torch.Tensor,
+            returns: torch.Tensor,
+            advantages: torch.Tensor,
+            mask: torch.Tensor
+    ):
+        """
+        计算PPO的总损失、Actor损失和Value损失。
+
+        :param log_probs: 当前策略的log probabilities, 形状: [batch_size, seq_len]
+        :param old_log_probs: 生成rollout时的旧策略的log probabilities, 形状: [batch_size, seq_len]
+        :param values: 当前评论家模型输出的价值, 形状: [batch_size, seq_len]
+        :param old_values: 生成rollout时的旧价值, 形状: [batch_size, seq_len]
+        :param returns: GAE计算出的回报, 形状: [batch_size, seq_len]
+        :param advantages: GAE计算出的优势, 形状: [batch_size, seq_len]
+        :param mask: 掩码，只计算生成部分的损失, 形状: [batch_size, seq_len]
+        :return: (总损失, Actor损失, Value损失, Entropy)
+        """
+        # Value Loss (价值损失) with clipping
+        values_clipped = old_values + torch.clamp(values - old_values, -self.clip_eps, self.clip_eps)
+        vf_loss_unclipped = F.mse_loss(values, returns, reduction='none')
+        vf_loss_clipped = F.mse_loss(values_clipped, returns, reduction='none')
+        value_loss = torch.max(vf_loss_unclipped, vf_loss_clipped)
+        # Apply mask and average
+        value_loss = 0.5 * (value_loss * mask).sum() / mask.sum().clamp(min=1.0)
+        value_loss = value_loss * self.vf_coef
+
+        # Actor Loss (策略损失)
+        # 计算新旧策略的概率比 r_t = exp(log_prob_new - log_prob_old)
+        # ratio 形状: [batch_size, seq_len]
+        ratio = torch.exp(log_probs - old_log_probs)
+
+        # PPO裁剪替代目标（Clipped Surrogate Objective）
+        # surr1 形状: [batch_size, seq_len]
+        surr1 = ratio * advantages
+        # surr2 形状: [batch_size, seq_len]
+        surr2 = torch.clamp(ratio, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * advantages
+
+        # 取两者中较小的一个，并加负号（因为我们要最大化这个目标，所以最小化它的负值）
+        # 我们只关心生成部分（由mask标记）的损失
+        actor_loss = -torch.sum(torch.min(surr1, surr2) * mask) / torch.sum(mask).clamp(min=1.0)
+
+        # 总损失
+        total_loss = actor_loss + value_loss
+
+        with torch.no_grad():
+            # 计算近似KL散度
+            logratios = log_probs - old_log_probs
+            approx_kl = torch.sum(((torch.exp(logratios) - 1) - logratios) * mask) / mask.sum().clamp(min=1.0)
+
+            # 计算裁剪比例
+            clipped = ratio.gt(1.0 + self.clip_eps) | ratio.lt(1.0 - self.clip_eps)
+            clip_frac = torch.sum(clipped.float() * mask) / mask.sum().clamp(min=1.0)
+
+        return total_loss, actor_loss, value_loss, approx_kl, clip_frac
+
+
+class GRPOLoss(nn.Module):
+    def __init__(
+            self,
+            beta: float,
+            clip_eps_low: float,
+            clip_eps_high: Optional[float] = None,
+            delta: Optional[float] = None,
+            importance_sampling_level: str = 'token',
+            loss_type: str = 'grpo',
+            gen_max_new_tokens: Optional[float] = None
+    ):
+        super().__init__()
+
+        self.beta = beta
+        self.clip_eps_low = clip_eps_low
+        self.clip_eps_high = clip_eps_high if clip_eps_high else clip_eps_low
+        self.delta = delta
+        self.importance_sampling_level = importance_sampling_level
+        self.loss_type = loss_type
+        self.gen_max_new_tokens = gen_max_new_tokens
+
+    def forward(
+            self,
+            log_probs: torch.Tensor,
+            old_log_probs: torch.Tensor,
+            ref_log_probs: torch.Tensor,
+            completion_mask: torch.Tensor,
+            advantages: torch.Tensor
+    ) -> torch.Tensor:
+
+        if self.beta != 0.0:
+            per_token_kl = torch.exp(ref_log_probs - log_probs) - (ref_log_probs - log_probs) - 1
+        else:
+            per_token_kl = None
+
+        log_ratio = log_probs - old_log_probs
+        if self.importance_sampling_level == "seq":
+            # GSPO
+            log_importance_weights = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
+            log_importance_weights = log_importance_weights.unsqueeze(-1)
+        else:
+            # GRPO
+            log_importance_weights = log_ratio
+
+        coef_1 = torch.exp(log_importance_weights)
+        coef_2 = torch.clamp(coef_1, 1 - self.clip_eps_low, 1 + self.clip_eps_high)
+
+        # Two-sided clipping
+        if self.delta is not None:
+            coef_1 = torch.clamp(coef_1, max=self.delta)
+
+        per_token_loss1 = coef_1 * advantages
+        per_token_loss2 = coef_2 * advantages
+        per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
+
+        if self.beta != 0.0:
+            per_token_loss = per_token_loss + self.beta * per_token_kl
+
+        if self.loss_type == "bnpo":
+            loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
+        elif self.loss_type == "dr_grpo":
+            assert self.gen_max_new_tokens is not None
+            loss = (per_token_loss * completion_mask).sum() / (per_token_loss.size(0) * self.gen_max_new_tokens)
+        else:
+            loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean()
+
+        return loss

llm_trainer/parallel.py

@@ -0,0 +1,220 @@
+import os
+from typing import Optional, Tuple
+from abc import ABC, abstractmethod
+
+import torch
+from torch import nn
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+try:
+    import deepspeed
+except: ...
+
+from .log import Logger
+
+
+class Parallel(ABC):
+    def __init__(
+            self,
+            _init_process_group: bool = True,
+            _use_parallel: bool = True
+    ):
+        self._initialize(_init_process_group, _use_parallel)
+
+    def _initialize(
+            self,
+            _init_process_group: bool,
+            _use_parallel: bool
+    ):
+        self._global_rank: int = int(os.environ.get('RANK', -1))
+        self._local_rank: int = int(os.environ.get('LOCAL_RANK', -1))
+        self._use_parallel: bool = _use_parallel and self._global_rank != -1
+
+        self._sampler: Optional[DistributedSampler] = None
+        self.model: Optional[nn.Module] = None
+
+        try:
+            torch.set_float32_matmul_precision('high')
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+        except:
+            pass
+
+        if self._use_parallel:
+            if _init_process_group:
+                dist.init_process_group(backend='nccl')
+
+            self.device: str = f'cuda:{self._local_rank}'
+            self.device_type: str = 'cuda'
+            torch.cuda.set_device(self.device)
+
+            Logger.std_log(f'global_rank={self._global_rank}, local_rank={self._local_rank}, world_size={self.world_size}')
+        else:
+            device = "cpu"
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+
+            self.device: str = device
+            self.device_type: str = device
+
+    @abstractmethod
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]: ...
+
+    def process_dataloader(
+            self,
+            dataset: Dataset,
+            data_loader_kwargs: dict,
+            sampler_kwargs: Optional[dict]=None
+    ) -> DataLoader:
+        """
+        :param dataset:
+        :param data_loader_kwargs
+                "batch_size" int,
+                "pin_memory" bool,
+                "collate_fn" collate_fn,
+                "num_workers" int
+                "shuffle" bool
+                "drop_last" bool
+        :param sampler_kwargs:
+                "shuffle" bool
+                "drop_last" bool
+        :return:
+        """
+
+        if self._use_parallel:
+            self._sampler = DistributedSampler(dataset=dataset, **sampler_kwargs)
+            return DataLoader(dataset=dataset, sampler=self._sampler, **data_loader_kwargs)
+
+        return DataLoader(dataset=dataset, **data_loader_kwargs)
+
+    def on_epoch_start(self, epoch):
+        if self._sampler:
+            self._sampler.set_epoch(epoch)
+
+    def on_epoch_end(self, epoch): ...
+
+    def synchronize(self):
+        if self._use_parallel:
+            torch.cuda.synchronize(device=self.device)
+
+    def destroy(self):
+        if self._use_parallel:
+            dist.destroy_process_group()
+
+    @property
+    def parallel_train(self) -> bool:
+        return self._use_parallel
+
+    @property
+    def is_main_process(self) -> bool:
+        if self._use_parallel:
+            return self._global_rank == 0
+
+        return True
+
+    @property
+    def world_size(self) -> int:
+        if self._use_parallel:
+            return dist.get_world_size()
+        return 1
+
+    def wait(self, msg=None):
+        if self.world_size == 1:
+            return
+
+        msg = f' for {msg}' if msg else ''
+        Logger.std_log(f'wait at {self.device}{msg}')
+        dist.barrier()
+        Logger.std_log(f'continue at {self.device}{msg}')
+
+
+class DsParallel(Parallel):
+    def __init__(self):
+        deepspeed.init_distributed(dist_backend='nccl')
+        super().__init__(_init_process_group=False)
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        """
+            :param model:
+            :param optimizer:
+            :param kwargs:
+                参考deepspeed配置
+            :param save_instance
+            :return:
+        """
+        model, optim, _, _ = deepspeed.initialize(
+            model=model,
+            optimizer=optimizer,
+            dist_init_required=False,
+            config_params=kwargs
+        )
+
+        if save_instance:
+            self.model = model
+
+        return model, optim
+
+    def synchronize(self): ...
+
+    def destroy(self): ...
+
+
+class DdpParallel(Parallel):
+    def __init__(self):
+        super().__init__()
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        model.to(self.device)
+
+        if self._use_parallel:
+            # self.model = DDP(module=model, broadcast_buffers=False, find_unused_parameters=True)
+            model = DDP(module=model, device_ids=[self._local_rank], output_device=self._local_rank)
+        else:
+            model = model
+
+        if save_instance:
+            self.model = model
+
+        return model, optimizer
+
+
+class NoneParallel(Parallel):
+    def __init__(self):
+        super().__init__(_use_parallel=False)
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        model.to(self.device)
+
+        if save_instance:
+            self.model = model
+
+        return model, optimizer

llm_trainer/partition_utils.py

@@ -0,0 +1,219 @@
+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
+import torch.distributed as dist
+
+from .tools import TrainerTools
+from .parallel import DsParallel, 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):
+    """
+        必须在所有rank上调用，非rank0, _to 可以设置为None.
+        当前函数不适用于_to是一个zero3模型
+    """
+    if isinstance(TrainerTools().parallel, DsParallel):
+        state_dict = _get_ds_model_params(_from, only_rank0=_to is None)
+    elif isinstance(_from, DDP):
+        state_dict = _from.module.state_dict()
+    else:
+        state_dict = _from.state_dict()
+
+    if not _to or not state_dict:
+        return
+
+    unwrap_to_model = unwrap_model(_to)
+    if mixup_alpha == 1.0:
+        # 直接覆盖
+        unwrap_to_model.load_state_dict(state_dict, strict=False)
+    else:
+        # 混合参数
+        for param_name, target_param in unwrap_to_model.named_parameters():
+            if param_name in state_dict:
+                from_param_tensor = state_dict[param_name]
+                target_param.data.mul_(1.0 - mixup_alpha).add_(
+                    from_param_tensor.data.to(target_param.device),
+                    alpha=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 _get_ds_full_state_dict_on_rank0(model: nn.Module) -> Optional[dict]:
+    """
+        需要在所有rank上调用，然后只有rank0有值
+    """
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+    if model.zero_optimization_stage() != 3:
+        if TrainerTools().parallel.is_main_process:
+            return {k: v.cpu().clone() for k, v in model.module.state_dict().items()}
+        return None
+
+    # --- ZeRO-3 ---
+    # 只调用一次 GatheredParameters，传入所有参数
+    with deepspeed.zero.GatheredParameters(model.parameters(), modifier_rank=0):
+        if TrainerTools().parallel.is_main_process:
+            # 在这个 'with' 代码块内，rank 0 上的 model.module 拥有完整的参数
+            # 所以我们可以像操作普通模型一样直接调用 state_dict()
+            full_state_dict = model.module.state_dict()
+
+            # 将其克隆到 CPU 并返回
+            return {k: v.cpu().clone() for k, v in full_state_dict.items()}
+
+    # 其他 rank 执行到这里时，上下文结束，直接返回 None
+    return None
+
+
+def _get_ds_model_params(model: nn.Module, only_rank0=False):
+    """
+        从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
+        兼容 ZeRO Stages 0, 1, 2, 3。
+        包含了对 ZeRO-3 中分片参数的正确处理。
+    """
+
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+    state_dict = _get_ds_full_state_dict_on_rank0(model)
+
+    # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
+    # 我们需要将其广播给所有进程。
+    if not only_rank0 and TrainerTools().parallel.world_size > 1:
+        # 准备一个列表，rank 0 有数据，其他 rank 是占位符
+        object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
+        # 执行广播，这个操作是阻塞的，会同步所有进程
+        dist.broadcast_object_list(object_list, src=0)
+        # 所有进程从列表中获取广播后的 state_dict 副本
+        state_dict = object_list[0]
+
+    return state_dict
+
+
+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):
+            # why only 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())