fix(tests): fix ci pipeline test error

InternLM · Jan 26, 2024 · 8170641 · 8170641
1 parent 3d7402d
commit 8170641
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/tests/test_core/test_pipeline.py b/tests/test_core/test_pipeline.py
@@ -118,7 +118,7 @@ def exam_pipeline_parallel(args):
         )
         output_list.append(output)
 
-    engine.step()
+    # engine.step()
 
     # torch related
     if gpc.is_last_rank(ParallelMode.PIPELINE):

diff --git a/tests/test_core/utils.py b/tests/test_core/utils.py
@@ -13,7 +13,7 @@
 from internlm.core.scheduler import InterleavedPipelineScheduler, NonPipelineScheduler, PipelineScheduler
 from internlm.model.metrics import SchedulerMetricHook
 from internlm.solver.pipeline_utils import partition_uniform
-from internlm.train import initialize_optimizer, set_parallel_attr_for_param_groups
+from internlm.train import initialize_optimizer
 
 
 class MlpModel(nn.Module):
@@ -67,8 +67,6 @@ def init_model_and_optim(
     pp_model = _build_generic_model_1d(num_layers=num_layers, num_chunks=num_chunks, embedding=embedding)
     pp_model = pp_model.to(dtype)
 
-    set_parallel_attr_for_param_groups(pp_model)
-
     # pp scheduler
     scheduler_hooks = [
         SchedulerMetricHook(skip=True),

diff --git a/tests/test_utils/test_model_checkpoint.py b/tests/test_utils/test_model_checkpoint.py
@@ -16,6 +16,8 @@
     LOCAL_SAVE_PATH,
     del_tmp_file,
     init_config,
+    init_dist_and_model,
+    reset_singletons,
 )
 
 # (TOTAL_STEP, CKPT_EVERY, SNPASHOT_EVERY)