diff --git a/src/twinkle/cli/cli.py b/src/twinkle/cli/cli.py index 1467a68d..3e39a721 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -17,7 +17,7 @@ class ModelArgs: model_cls: str | None = None tokenizer_id: str | None = None mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16' - strategy: Literal['accelerate', 'native_fsdp'] = field( + strategy: Literal['accelerate', 'native_fsdp', 'deepspeed'] = field( default='accelerate', metadata={'aliases': ('use_megatron', )}) memory_efficient_init: bool = False gradient_checkpointing: bool = True @@ -33,6 +33,7 @@ class ModelArgs: variable_seq_lengths: bool = True ddp_config: dict[str, Any] | None = None fsdp_config: dict[str, Any] | None = None + deepspeed_config: dict[str, Any] | None = None grad_scaler_config: dict[str, Any] | None = None # Liger Kernel toggle: gates the fused-linear-CE loss in the cookbooks. # Off by default — opt in with --enable-liger / TWINKLE_ENABLE_LIGER. diff --git a/src/twinkle/model/transformers/strategy/__init__.py b/src/twinkle/model/transformers/strategy/__init__.py index 8ab90b18..64964f22 100644 --- a/src/twinkle/model/transformers/strategy/__init__.py +++ b/src/twinkle/model/transformers/strategy/__init__.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .accelerate import AccelerateStrategy +from .deepspeed import DeepSpeedStrategy from .native_fsdp import NativeFSDPStrategy -__all__ = ['AccelerateStrategy', 'NativeFSDPStrategy'] +__all__ = ['AccelerateStrategy', 'DeepSpeedStrategy', 'NativeFSDPStrategy'] diff --git a/src/twinkle/model/transformers/strategy/deepspeed.py b/src/twinkle/model/transformers/strategy/deepspeed.py new file mode 100644 index 00000000..fd64c4aa --- /dev/null +++ b/src/twinkle/model/transformers/strategy/deepspeed.py @@ -0,0 +1,253 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSpeed training strategy for twinkle. + +A peer of ``AccelerateStrategy`` and ``NativeFSDPStrategy``. All DeepSpeed +logic is self-contained here — the existing strategies are not modified. + +Design: construct ``DeepSpeedPlugin`` + ``Accelerator`` in ``__init__``. +``wrap_model`` applies ZeRO-3 compatibility patches (from +``twinkle.patch.deepspeed_patches``) around ``accelerator.prepare``. +``get_full_state_dict`` / ``gather_parameters`` use DeepSpeed's +``GatheredParameters`` for ZeRO-3 parameter gathering. +""" +import os +from contextlib import nullcontext +from datetime import timedelta +from typing import Any, Dict, Literal, Mapping, Optional + +from twinkle.patch import apply_context, apply_patch +# Re-export for convenience +from twinkle.patch.deepspeed_patches import DeepSpeedLeafModulesPatch # noqa: F401 +from twinkle.patch.deepspeed_patches import (DeepSpeedHookReorderPatch, DeepSpeedModulesToSavePatch, + DeepSpeedParamWrapperPatch) + + +class DeepSpeedStrategy: + """A training strategy backed by HuggingFace accelerate + DeepSpeed. + + Args: + mixed_precision: Mixed precision type. + deepspeed_config: DeepSpeed config dict. If ``None``, a minimal + ZeRO-2 config is used. + device_mesh: Model device mesh (unused by DeepSpeed but kept for + interface parity with other strategies). + """ + + def __init__( + self, + mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16', + deepspeed_config: Optional[Dict[str, Any]] = None, + device_mesh=None, + ): + from accelerate import Accelerator + from accelerate.utils import InitProcessGroupKwargs + + self.device_mesh = device_mesh + self.mixed_precision = mixed_precision + + config = deepspeed_config or {'zero_optimization': {'stage': 2}} + self._deepspeed_plugin = self._build_plugin(config) + + self.accelerator = Accelerator( + deepspeed_plugin=self._deepspeed_plugin, + mixed_precision=mixed_precision, + kwargs_handlers=[ + InitProcessGroupKwargs( + timeout=timedelta(seconds=int(os.environ.get('TWINKLE_DIST_TIMEOUT_SECONDS', '7200')), ), ), + ], + ) + + # -- plugin construction ------------------------------------------------ + + @staticmethod + def _build_plugin(config: Dict[str, Any]): + """Build a ``DeepSpeedPlugin`` from a raw config dict. + + Uses ``HfTrainerDeepSpeedConfig`` when available so that ``"auto"`` + values are resolved; falls back to a plain dict otherwise. + """ + from accelerate.utils import DeepSpeedPlugin + + try: + from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig + + hf_ds_config = HfTrainerDeepSpeedConfig(config) + return DeepSpeedPlugin(hf_ds_config=hf_ds_config) + except ImportError: + return DeepSpeedPlugin(hf_ds_config=config) + + # -- state queries ------------------------------------------------------ + + @property + def is_deepspeed_enabled(self) -> bool: + return True + + def is_zero_stage(self, stage: int) -> bool: + cfg = self._deepspeed_plugin.deepspeed_config + if not isinstance(cfg, dict): + cfg = getattr(cfg, 'config', {}) + return cfg.get('zero_optimization', {}).get('stage', 0) == stage + + @property + def zero_stage(self) -> int: + cfg = self._deepspeed_plugin.deepspeed_config + if not isinstance(cfg, dict): + cfg = getattr(cfg, 'config', {}) + return cfg.get('zero_optimization', {}).get('stage', 0) + + # -- twinkle Strategy interface ----------------------------------------- + + def pretrained_load_context(self): + """Context manager for loading pretrained weights. + + Under ZeRO-3, returns ``deepspeed.zero.Init`` so that weights are + partitioned immediately during ``from_pretrained``. For ZeRO-1/2, + weights fit in full on each rank, so no special context is needed. + """ + if self.is_zero_stage(3): + import deepspeed + + cfg = self._deepspeed_plugin.deepspeed_config + if not isinstance(cfg, dict): + cfg = getattr(cfg, 'config', cfg) + return deepspeed.zero.Init(config_dict_or_path=cfg) + return nullcontext() + + def capture_pre_ep_state_if_needed(self, model, *, enable_ep: bool) -> None: + pass + + def prepare_adapter_config(self, config_or_dir, *, enable_ep: bool): + return config_or_dir + + def wrap_model(self, model, *args): + """Wrap model with DeepSpeed via ``accelerator.prepare``. + + Patch ordering: + 1. ``DeepSpeedLeafModulesPatch`` (permanent) — ZeRO-3 MoE leaf modules + 2. ``DeepSpeedModulesToSavePatch`` (permanent) — PEFT modules_to_save + 3. ``DeepSpeedHookReorderPatch`` (temporary) — wrap ``deepspeed.initialize`` + 4. ``DeepSpeedParamWrapperPatch`` (temporary) — PEFT param metadata + + Steps 3-4 are active only during ``accelerator.prepare``; the wrapper + restores originals on exit. + """ + apply_patch(model, DeepSpeedLeafModulesPatch) + apply_patch(model, DeepSpeedModulesToSavePatch) + + with apply_context(model, DeepSpeedHookReorderPatch): + with apply_context(model, DeepSpeedParamWrapperPatch): + return self.accelerator.prepare(model, *args) + + def unwrap_model(self, model): + return self.accelerator.unwrap_model(model, keep_torch_compile=False) + + def load_peft_weights( + self, + model, + adapter_weights: Mapping[str, Any], + adapter_name: str, + ) -> None: + from peft.utils import set_peft_model_state_dict + + set_peft_model_state_dict(model, adapter_weights, adapter_name=adapter_name) + + def needs_wrapped_optimizer_state(self) -> bool: + return True + + def save_optimizer_checkpoint( + self, + model, + optimizer, + output_path: str, + ) -> None: + """Save optimizer state via DeepSpeed engine checkpoint. + + ``output_path`` is a file path (e.g. ``.../optimizer.pt``); DeepSpeed + ``save_checkpoint`` takes a directory, so we use the parent dir. + Both model and optimizer state are saved — the model portion is + redundant with ``save()`` but is the standard DeepSpeed format. + """ + engine = self.unwrap_model(model) + save_dir = os.path.dirname(output_path) + engine.save_checkpoint(save_dir) + + def load_optimizer_checkpoint( + self, + model, + optimizer, + input_path: str, + ) -> None: + """Load optimizer state from a DeepSpeed engine checkpoint.""" + from twinkle.utils import get_logger + + logger = get_logger() + engine = self.unwrap_model(model) + load_dir = os.path.dirname(input_path) + try: + engine.load_checkpoint(load_dir) + except Exception as e: + logger.warning(f'Failed to load deepspeed checkpoint: {e}') + + def get_full_state_dict(self, model) -> dict: + """Collect full state dict, gathering ZeRO-3 partitioned params.""" + unwrapped = self.unwrap_model(model) + if self.is_zero_stage(3): + from deepspeed.runtime.zero import GatheredParameters + + params = list(unwrapped.parameters()) + with GatheredParameters(params): + return {k: v.cpu() for k, v in unwrapped.named_parameters()} + return {k: v.cpu() for k, v in unwrapped.named_parameters()} + + def load_full_state_dict(self, model, state_dict) -> None: + """Load a full state dict into the (possibly partitioned) model.""" + unwrapped = self.unwrap_model(model) + if self.is_zero_stage(3): + from deepspeed.runtime.zero import GatheredParameters + + params = list(unwrapped.parameters()) + with GatheredParameters(params): + unwrapped.load_state_dict(state_dict, strict=False) + else: + unwrapped.load_state_dict(state_dict, strict=False) + + def get_adapter_state_dict(self, model, adapter_name: str) -> dict: + """Collect only LoRA adapter parameters, gathering under ZeRO-3.""" + unwrapped = self.unwrap_model(model) + adapter_suffix = f'.{adapter_name}.' + state_dict = {} + + if self.is_zero_stage(3): + from deepspeed.runtime.zero import GatheredParameters + + params = [p for n, p in unwrapped.named_parameters() if _is_lora_state_key(n) and adapter_suffix in n] + with GatheredParameters(params): + for n, p in unwrapped.named_parameters(): + if _is_lora_state_key(n) and adapter_suffix in n: + state_dict[n] = p.cpu() + else: + for n, p in unwrapped.named_parameters(): + if _is_lora_state_key(n) and adapter_suffix in n: + state_dict[n] = p.cpu() + return state_dict + + # -- DeepSpeed-specific ------------------------------------------------- + + def gather_parameters(self, params): + """Context manager to gather ZeRO-3 partitioned parameters. + + Usage:: + + with strategy.gather_parameters(list(model.parameters())): + # params are fully gathered here + ... + """ + if self.is_zero_stage(3): + from deepspeed.runtime.zero import GatheredParameters + + return GatheredParameters(params) + return nullcontext() + + +def _is_lora_state_key(name: str) -> bool: + return 'lora_A' in name or 'lora_B' in name or 'lora_embedding' in name diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index e9e86f96..bc08deb5 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -36,7 +36,7 @@ from twinkle.model.base import TwinkleModel from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus from twinkle.model.transformers.moe import apply_expert_parallel -from twinkle.model.transformers.strategy import AccelerateStrategy, NativeFSDPStrategy +from twinkle.model.transformers.strategy import AccelerateStrategy, DeepSpeedStrategy, NativeFSDPStrategy from twinkle.module.optimizer import GaLoreConfig, create_galore_param_groups from twinkle.patch import Patch, apply_context, apply_patch from twinkle.processor import InputProcessor @@ -220,11 +220,12 @@ def __init__( config: Optional[PretrainedConfig] = None, device_mesh: Optional[DeviceMesh] = None, mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16', - strategy: Literal['accelerate', 'native_fsdp'] = 'accelerate', + strategy: Literal['accelerate', 'native_fsdp', 'deepspeed'] = 'accelerate', ddp_config: Dict[str, Any] = None, fsdp_config: Dict[str, Any] = None, grad_scaler_config: Dict[str, Any] = None, memory_efficient_init: bool = False, + deepspeed_config: Dict[str, Any] = None, **kwargs): os.environ['TOKENIZERS_PARALLELISM'] = 'true' self._try_init_process_group() @@ -238,7 +239,7 @@ def __init__( self._memory_efficient_init = memory_efficient_init self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) self._router_replay_applied = False - self._decide_strategy(strategy) + self._decide_strategy(strategy, deepspeed_config=deepspeed_config) self.grad_scaler_config = grad_scaler_config if model_id is not None: model_id = HubOperation.download_model(model_id) @@ -294,32 +295,41 @@ def _init_empty_model_from_config(self, model_cls, **kwargs): model.tie_weights() return model - def _decide_strategy(self, strategy: Literal['accelerate', 'native_fsdp']): + def _decide_strategy(self, + strategy: Literal['accelerate', 'native_fsdp', 'deepspeed'], + deepspeed_config: Dict[str, Any] = None): self._expert_parallel_config = self._fsdp_config.pop('expert_parallel', None) self._enable_expert_parallel = self._should_enable_expert_parallel(self._expert_parallel_config, self.device_mesh) self._expert_parallel_applied = False - use_native_fsdp = self._enable_expert_parallel or strategy == 'native_fsdp' - if use_native_fsdp: - ep_size = (self._expert_parallel_config.get('ep_size') if self._expert_parallel_config else None) - if ep_size is None and self.device_mesh is not None: - ep_size = getattr(self.device_mesh, 'ep_size', None) - self.strategy = NativeFSDPStrategy( + if strategy == 'deepspeed': + self.strategy = DeepSpeedStrategy( mixed_precision=self.mixed_precision, - fsdp_config=self._fsdp_config, + deepspeed_config=deepspeed_config, device_mesh=self.device_mesh, - memory_efficient_init=self._memory_efficient_init, - enable_ep=self._enable_expert_parallel, - ep_size=ep_size, ) else: - self.strategy = AccelerateStrategy( - mixed_precision=self.mixed_precision, - ddp_config=self._ddp_config, - fsdp_config=self._fsdp_config, - device_mesh=self.device_mesh, - memory_efficient_init=self._memory_efficient_init) + use_native_fsdp = self._enable_expert_parallel or strategy == 'native_fsdp' + if use_native_fsdp: + ep_size = (self._expert_parallel_config.get('ep_size') if self._expert_parallel_config else None) + if ep_size is None and self.device_mesh is not None: + ep_size = getattr(self.device_mesh, 'ep_size', None) + self.strategy = NativeFSDPStrategy( + mixed_precision=self.mixed_precision, + fsdp_config=self._fsdp_config, + device_mesh=self.device_mesh, + memory_efficient_init=self._memory_efficient_init, + enable_ep=self._enable_expert_parallel, + ep_size=ep_size, + ) + else: + self.strategy = AccelerateStrategy( + mixed_precision=self.mixed_precision, + ddp_config=self._ddp_config, + fsdp_config=self._fsdp_config, + device_mesh=self.device_mesh, + memory_efficient_init=self._memory_efficient_init) # Sequence parallel ("ulysses") is derived from dp/fsdp ranks; it does not change world size. # We construct `sp_strategy` after the underlying HF model is initialized (see __init__). diff --git a/src/twinkle/patch/deepspeed_patches.py b/src/twinkle/patch/deepspeed_patches.py new file mode 100644 index 00000000..58628ea6 --- /dev/null +++ b/src/twinkle/patch/deepspeed_patches.py @@ -0,0 +1,221 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSpeed ZeRO-3 compatibility patches. + +Four patches that make PEFT and multimodal hooks work correctly under +DeepSpeed ZeRO-3 parameter partitioning. All inherit from twinkle's +``Patch`` base class and are applied via ``apply_patch`` / ``apply_context``. + +1. **DeepSpeedLeafModulesPatch** — marks MoE blocks as ZeRO-3 leaf modules + so that expert parameters are partitioned as a unit, not individually. + +2. **DeepSpeedParamWrapperPatch** — patches ``ParamWrapper.get_param`` so + that NOT_AVAILABLE ZeRO-3 params return a stride-0 placeholder with + correct metadata (shape/ndim/dtype) using O(1) memory instead of + gathering the full parameter. + +3. **DeepSpeedModulesToSavePatch** — patches ``ModulesToSaveWrapper.__setattr__`` + so that ``ds_grads_remaining`` is propagated to all ``modules_to_save`` + sub-modules, keeping ZeRO-3 gradient tracking consistent. + +4. **DeepSpeedHookReorderPatch** — wraps ``deepspeed.initialize`` so that + application-level forward pre-hooks (e.g. multimodal input preprocessing) + are moved to the end of ``_forward_pre_hooks`` after DeepSpeed adds its + own hooks, ensuring they execute last. +""" +from functools import wraps +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type + +from twinkle.patch import Patch + +if TYPE_CHECKING: + import torch.nn as nn + +# model_type -> (import_path, class_name) for MoE leaf module resolution. +_MOE_LEAF_MAP: Dict[str, Tuple[str, str]] = { + 'qwen3_vl_moe': ('transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe', 'Qwen3VLMoeTextSparseMoeBlock'), + 'qwen3_omni_moe': + ('transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe', 'Qwen3OmniMoeThinkerTextSparseMoeBlock'), + 'qwen2_moe': ('transformers.models.qwen2_moe.modeling_qwen2_moe', 'Qwen2MoeSparseMoeBlock'), + 'qwen3_moe': ('transformers.models.qwen3_moe.modeling_qwen3_moe', 'Qwen3MoeSparseMoeBlock'), + 'gemma4': ('transformers.models.gemma4.modeling_gemma4', 'Gemma4TextExperts'), + 'glm4_moe': ('transformers.models.glm4_moe.modeling_glm4_moe', 'Glm4MoeMoE'), + 'glm4_moe_lite': ('transformers.models.glm4_moe_lite.modeling_glm4_moe_lite', 'Glm4MoeLiteMoE'), + 'glm4v_moe': ('transformers.models.glm4v_moe.modeling_glm4v_moe', 'Glm4vMoeTextMoE'), + 'gpt_oss': ('transformers.models.gpt_oss.modeling_gpt_oss', 'GptOssMLP'), + 'llama4': ('transformers.models.llama4.modeling_llama4', 'Llama4TextMoe'), + 'qwen3_next': ('transformers.models.qwen3_next.modeling_qwen3_next', 'Qwen3NextSparseMoeBlock'), + 'olmoe': ('transformers.models.olmoe.modeling_olmoe', 'OlomoeSparseMoeBlock'), + 'qwen3_5_moe': ('transformers.models.qwen3_5_moe.modeling_qwen3_5_moe', 'Qwen3_5MoeSparseMoeBlock'), + 'glm_moe_dsa': ('transformers.models.glm_moe_dsa.modeling_glm_moe_dsa', 'GlmMoeDsaMoE'), +} + + +class DeepSpeedLeafModulesPatch(Patch): + """Mark MoE blocks as ZeRO-3 leaf modules. + + DeepSpeed ZeRO-3 partitions parameters at the leaf-module level. For MoE + models, the sparse MoE block must be a leaf module so that expert + parameters are partitioned as a group, not individually (which would + break expert routing). + + Resolves the MoE block class by ``model.config.model_type`` via a static + map. For trust_remote_code models not in the map, scans ``model.modules()`` + for classes whose name ends with ``MoE`` or ``SparseMoeBlock``. + + Permanent patch: ``unpatch`` is a no-op. + """ + + def __call__(self, module: 'nn.Module', *args, **kwargs): + try: + model_type = module.config.model_type + except Exception: + return module + + leaf_modules = self._resolve_leaf_modules(module, model_type) + if leaf_modules: + from deepspeed.utils import set_z3_leaf_modules + set_z3_leaf_modules(module, leaf_modules) + return module + + @staticmethod + def _resolve_leaf_modules( + model: 'nn.Module', + model_type: str, + ) -> Optional[List[Type]]: + entry = _MOE_LEAF_MAP.get(model_type) + if entry is not None: + import importlib + module_path, class_name = entry + try: + mod = importlib.import_module(module_path) + return [getattr(mod, class_name)] + except (ImportError, AttributeError): + pass + + # trust_remote_code fallback: scan for MoE block by class name + for sub in model.modules(): + cn = type(sub).__name__ + if cn.endswith('MoE') or cn.endswith('SparseMoeBlock'): + return [type(sub)] + return None + + def unpatch(self, module: 'nn.Module', *args, **kwargs): + pass + + +class DeepSpeedParamWrapperPatch(Patch): + """Patch ``ParamWrapper.get_param`` for ZeRO-3 compatibility. + + When a parameter is ``NOT_AVAILABLE`` in ZeRO-3, ``param.data`` is a + placeholder with wrong shape/ndim. Callers of ``get_param()`` only need + metadata (shape, ndim, dtype, device, requires_grad), so we use + ``ds_shape`` + ``expand`` to create a stride-0 tensor with correct + metadata using O(1) memory instead of gathering the full parameter. + + Temporary patch: ``unpatch`` restores the original method. + """ + + def __call__(self, module: 'nn.Module', *args, **kwargs): + try: + from peft.tuners.lora.layer import ParamWrapper + except ImportError: + return module + + self._origin = ParamWrapper.get_param + origin = self._origin + + def _get_param_patched(wrapper_self): + param = origin(wrapper_self) + if hasattr(param, 'ds_id'): + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + ds_shape = param.ds_shape + ones = tuple(1 for _ in ds_shape) + import torch + fake = torch.empty(ones, dtype=param.dtype, device=param.device) + if param.requires_grad and param.dtype.is_floating_point: + fake.requires_grad_(True) + return fake.expand(ds_shape) + return param + + ParamWrapper.get_param = _get_param_patched + return module + + def unpatch(self, module: 'nn.Module', *args, **kwargs): + try: + from peft.tuners.lora.layer import ParamWrapper + ParamWrapper.get_param = self._origin + except (ImportError, AttributeError): + pass + + +class DeepSpeedModulesToSavePatch(Patch): + """Patch ``ModulesToSaveWrapper.__setattr__`` for ZeRO-3 compatibility. + + Propagates ``ds_grads_remaining`` from the wrapper to all + ``modules_to_save`` sub-modules, so ZeRO-3 gradient tracking stays + consistent when modules are saved outside the main model. + + Permanent patch: ``_patched`` class flag prevents double-application. + ``unpatch`` is a no-op. + """ + + def __call__(self, module: 'nn.Module', *args, **kwargs): + from peft.utils import ModulesToSaveWrapper + + if getattr(ModulesToSaveWrapper, '_patched', False): + return module + + ModulesToSaveWrapper._patched = True + old_setattr = ModulesToSaveWrapper.__setattr__ + self._old_setattr = old_setattr + + def _patched_setattr(wrapper_self, name, value): + old_setattr(wrapper_self, name, value) + if name == 'ds_grads_remaining': + for sub in wrapper_self.modules_to_save.values(): + sub.ds_grads_remaining = value + + ModulesToSaveWrapper.__setattr__ = _patched_setattr + return module + + def unpatch(self, module: 'nn.Module', *args, **kwargs): + pass + + +class DeepSpeedHookReorderPatch(Patch): + """Reorder forward pre-hooks after ``deepspeed.initialize``. + + ``deepspeed.initialize`` adds its own forward pre-hooks to the model. + Application-level hooks (e.g. multimodal input preprocessing) registered + before initialization must run *after* DeepSpeed's hooks, so we record + pre-existing hook IDs, call the original ``deepspeed.initialize``, then + move those hooks to the end of ``_forward_pre_hooks``. + + Temporary patch: ``unpatch`` restores ``deepspeed.initialize``. + """ + + def __call__(self, module: 'nn.Module', *args, **kwargs): + import deepspeed + + self._model = module + self._origin_init = deepspeed.initialize + + model = module + origin_init = self._origin_init + + @wraps(origin_init) + def _initialize(*args, **kwargs): + pre_hook_ids = list(model._forward_pre_hooks.keys()) + res = origin_init(*args, **kwargs) + for hook_id in pre_hook_ids: + model._forward_pre_hooks.move_to_end(hook_id) + return res + + deepspeed.initialize = _initialize + return module + + def unpatch(self, module: 'nn.Module', *args, **kwargs): + import deepspeed + if hasattr(self, '_origin_init'): + deepspeed.initialize = self._origin_init