code format

Author: hushenwei2000

paddleformers/trainer/trainer.py

@@ -201,7 +201,6 @@
 DEFAULT_CALLBACKS = [DefaultFlowCallback]
 DEFAULT_PROGRESS_CALLBACK = ProgressCallback
 
-
 if is_datasets_available():
     import datasets
 

paddleformers/trainer/trainer_callback.py

@@ -25,11 +25,10 @@
 from typing import Dict, List, Optional, Union
 
 import numpy as np
+from moe_utils import offload, reload
 from tqdm.auto import tqdm
 
-from paddleformers.transformers.moe_utils import offload, reload
-from paddleformers.utils.log import logger
-
+from ..utils.log import logger
 from .trainer_utils import IntervalStrategy, has_length
 from .training_args import TrainingArguments
 

paddleformers/trainer/training_args.py

@@ -1575,7 +1575,6 @@ def is_segment_parallel_supported():
                 fleet.init(is_collective=True, strategy=strategy)
                 logger.info(strategy)
 
-                # if os.getenv("DSV3_FAST_PRETRAIN", "False"):
                 if self.dsv3_fast_pretrain:
                     if self.expert_parallel_degree > 1:
                         self.add_moe_comm_group()

paddleformers/transformers/moe_utils.py

@@ -99,18 +99,3 @@ def unpermute(
         include_self=True,
     )
     return output_tokens
-
-
-def offload(tensor):
-    if paddle.is_compiled_with_cuda():
-        place = paddle.CUDAPinnedPlace()
-    else:
-        place = paddle.CPUPlace()
-
-    new_tensor = to_device(tensor, place)
-    assert new_tensor is tensor, "to_device must be inplace operation"
-
-
-def reload(tensor):
-    new_tensor = to_device(tensor)
-    assert new_tensor is tensor, "to_device must be inplace operation"