fix unit test (#121)

Co-authored-by: hermanyu <[email protected]>
Tencent · Jan 23, 2024 · 817a1b2 · 817a1b2
1 parent c703d03
commit 817a1b2
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tencentpretrain/mpu/layers.py b/tencentpretrain/mpu/layers.py
@@ -36,8 +36,6 @@
 from .utils import split_tensor_along_last_dim
 from .utils import VocabUtility
 #from megatron import get_args
-import deepspeed.runtime.activation_checkpointing.checkpointing as ds_checkpointing
-
 
 _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {
     "tensor_model_parallel": False,
@@ -83,6 +81,7 @@ def maybe_copy(attribute):
 def _initialize_affine_weight_gpu(weight, init_method, partition_dim, stride=1):
     """Initialize affine weight for model parallel on GPU."""
 
+    import deepspeed.runtime.activation_checkpointing.checkpointing as ds_checkpointing
     set_tensor_model_parallel_attributes(
         tensor=weight, is_parallel=True, dim=partition_dim, stride=stride
     )