From 4dbb52750f8fd70b28030ec1d2e7846333a6c85b Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Thu, 9 Nov 2023 16:41:32 +0800
Subject: [PATCH] fix #1452

Former-commit-id: 0e86527d7fae9c9fe0df89d6fbd89035c9d83fe3
---
 src/llmtuner/tuner/dpo/trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/llmtuner/tuner/dpo/trainer.py b/src/llmtuner/tuner/dpo/trainer.py
index 8a9f8dd6..647bcee2 100644
--- a/src/llmtuner/tuner/dpo/trainer.py
+++ b/src/llmtuner/tuner/dpo/trainer.py
@@ -1,5 +1,4 @@
 import torch
-import deepspeed # type: ignore
 from copy import deepcopy
 from collections import defaultdict
 from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union
@@ -76,6 +75,8 @@ class CustomDPOTrainer(DPOTrainer):
         # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)
         if config_kwargs["zero_optimization"]["stage"] != 3:
             config_kwargs["zero_optimization"]["stage"] = 0
+        # lazy load
+        import deepspeed # type: ignore
         model, *_ = deepspeed.initialize(model=model, config=config_kwargs)
         model.eval()
         return model