fix cyclic imports

Author: zigzagcai

internlm/__init__.py

@@ -1,8 +1,7 @@
 from .initialize.initialize_trainer import initialize_trainer
-from .initialize.launch import get_default_parser, launch_from_slurm, launch_from_torch
+from .initialize.launch import launch_from_slurm, launch_from_torch
 
 __all__ = [
-    "get_default_parser",
     "initialize_trainer",
     "launch_from_slurm",
     "launch_from_torch",

internlm/checkpoint/checkpoint_manager.py

@@ -16,8 +16,8 @@
     auto_resume_sanity_check,
     ckpt_info_sanity_check,
 )
-from internlm.model.base_model import BaseModel
-from internlm.model.registry import model_initializer
+from internlm.model.model_implementations.registry import model_initializer
+from internlm.model.model_implementations.transformers.base_model import BaseModel
 from internlm.monitor import send_alert_message
 from internlm.solver.optimizer import HybridZeroOptimizer, HybridZeroOptimizer_v2
 from internlm.utils.common import get_current_device

internlm/checkpoint/components.py

@@ -9,7 +9,7 @@
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc
 from internlm.core.trainer import TrainState
-from internlm.model.moe import MoE
+from internlm.model.model_ops.moe import MoE
 from internlm.solver.optimizer import HybridZeroOptimizer, HybridZeroOptimizer_v2
 from internlm.utils.common import get_current_device
 from internlm.utils.lazy import LazyObject

internlm/checkpoint/load_funcs.py

@@ -1,8 +1,8 @@
 # Copyright (c) InternLM. All rights reserved.
 
-from internlm.model.modeling_internlm import InternLM1
-from internlm.model.modeling_internlm2 import InternLM2
-from internlm.model.modeling_llama import Llama2
+from internlm.model.model_implementations.transformers.modeling_internlm import InternLM1
+from internlm.model.model_implementations.transformers.modeling_internlm2 import InternLM2
+from internlm.model.model_implementations.transformers.modeling_llama import Llama2
 from internlm.utils.logger import get_logger
 
 logger = get_logger(__file__)

internlm/core/context/__init__.py

@@ -1,11 +1,11 @@
+from .config import Config
 from .parallel_context import (
     IS_REPLICA_EXPERT_DATA_PARALLEL,
     IS_REPLICA_ZERO_PARALLEL,
     IS_TENSOR_EXPERT_DATA_PARALLEL,
     IS_TENSOR_ZERO_PARALLEL,
     IS_WEIGHT_EXPERT_DATA_PARALLEL,
     IS_WEIGHT_ZERO_PARALLEL,
-    Config,
     ParallelContext,
     global_context,
 )

internlm/core/context/config.py

@@ -0,0 +1,95 @@
+import inspect
+import sys
+from importlib.machinery import SourceFileLoader
+from pathlib import Path
+
+
+class Config(dict):
+    """This is a wrapper class for dict objects so that values of which can be
+    accessed as attributes.
+
+    Args:
+        config (dict): The dict object to be wrapped.
+    """
+
+    def __init__(self, config: dict = None):  # pylint: disable=W0231
+        if config is not None:
+            for k, v in config.items():
+                self._add_item(k, v)
+
+    def __missing__(self, key):
+        raise KeyError(key)
+
+    def __getattr__(self, key):
+        try:
+            value = super().__getitem__(key)
+            return value
+        except KeyError:
+            raise AttributeError(key)
+
+    def __setattr__(self, key, value):
+        super().__setitem__(key, value)
+
+    def _add_item(self, key, value):
+        if isinstance(value, dict):
+            self.__setattr__(key, Config(value))
+        else:
+            self.__setattr__(key, value)
+
+    def update(self, config):
+        assert isinstance(config, (Config, dict)), "can only update dictionary or Config objects."
+        for k, v in config.items():
+            self._add_item(k, v)
+        return self
+
+    @staticmethod
+    def from_file(filename: str):
+        """Reads a python file and constructs a corresponding :class:`Config` object.
+
+        Args:
+            filename (str): Name of the file to construct the return object.
+
+        Returns:
+            :class:`Config`: A :class:`Config` object constructed with information in the file.
+
+        Raises:
+            AssertionError: Raises an AssertionError if the file does not exist, or the file is not .py file
+        """
+
+        # check config path
+        if isinstance(filename, str):
+            filepath = Path(filename).absolute()
+        elif isinstance(filename, Path):
+            filepath = filename.absolute()
+
+        assert filepath.exists(), f"{filename} is not found, please check your configuration path"
+
+        # check extension
+        extension = filepath.suffix
+        assert extension == ".py", "only .py files are supported"
+
+        # import the config as module
+        remove_path = False
+        if filepath.parent not in sys.path:
+            sys.path.insert(0, (filepath))
+            remove_path = True
+
+        module_name = filepath.stem
+        source_file = SourceFileLoader(fullname=str(module_name), path=str(filepath))
+        module = source_file.load_module()  # pylint: disable=W4902,E1120,W1505
+
+        # load into config
+        config = Config()
+
+        for k, v in module.__dict__.items():
+            if k.startswith("__") or inspect.ismodule(v) or inspect.isclass(v):
+                continue
+            else:
+                config._add_item(k, v)
+
+        # remove module
+        del sys.modules[module_name]
+        if remove_path:
+            sys.path.pop(0)
+
+        return config

internlm/core/context/parallel_context.py

@@ -3,19 +3,16 @@
 
 # adopted from https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/context
 
-import inspect
 import random
 import socket
-import sys
-from importlib.machinery import SourceFileLoader
-from pathlib import Path
 from typing import Union
 
 import numpy as np
 import torch
 import torch.distributed as dist
 
 from internlm.accelerator import get_accelerator
+from internlm.core.context import Config
 from internlm.utils.common import SingletonMeta
 from internlm.utils.logger import get_logger
 from internlm.utils.timeout import LLM_NCCL_TIMEOUT
@@ -46,97 +43,6 @@
 internlm_accelerator = get_accelerator()
 
 
-class Config(dict):
-    """This is a wrapper class for dict objects so that values of which can be
-    accessed as attributes.
-
-    Args:
-        config (dict): The dict object to be wrapped.
-    """
-
-    def __init__(self, config: dict = None):  # pylint: disable=W0231
-        if config is not None:
-            for k, v in config.items():
-                self._add_item(k, v)
-
-    def __missing__(self, key):
-        raise KeyError(key)
-
-    def __getattr__(self, key):
-        try:
-            value = super().__getitem__(key)
-            return value
-        except KeyError:
-            raise AttributeError(key)
-
-    def __setattr__(self, key, value):
-        super().__setitem__(key, value)
-
-    def _add_item(self, key, value):
-        if isinstance(value, dict):
-            self.__setattr__(key, Config(value))
-        else:
-            self.__setattr__(key, value)
-
-    def update(self, config):
-        assert isinstance(config, (Config, dict)), "can only update dictionary or Config objects."
-        for k, v in config.items():
-            self._add_item(k, v)
-        return self
-
-    @staticmethod
-    def from_file(filename: str):
-        """Reads a python file and constructs a corresponding :class:`Config` object.
-
-        Args:
-            filename (str): Name of the file to construct the return object.
-
-        Returns:
-            :class:`Config`: A :class:`Config` object constructed with information in the file.
-
-        Raises:
-            AssertionError: Raises an AssertionError if the file does not exist, or the file is not .py file
-        """
-
-        # check config path
-        if isinstance(filename, str):
-            filepath = Path(filename).absolute()
-        elif isinstance(filename, Path):
-            filepath = filename.absolute()
-
-        assert filepath.exists(), f"{filename} is not found, please check your configuration path"
-
-        # check extension
-        extension = filepath.suffix
-        assert extension == ".py", "only .py files are supported"
-
-        # import the config as module
-        remove_path = False
-        if filepath.parent not in sys.path:
-            sys.path.insert(0, (filepath))
-            remove_path = True
-
-        module_name = filepath.stem
-        source_file = SourceFileLoader(fullname=str(module_name), path=str(filepath))
-        module = source_file.load_module()  # pylint: disable=W4902,E1120,W1505
-
-        # load into config
-        config = Config()
-
-        for k, v in module.__dict__.items():
-            if k.startswith("__") or inspect.ismodule(v) or inspect.isclass(v):
-                continue
-            else:
-                config._add_item(k, v)
-
-        # remove module
-        del sys.modules[module_name]
-        if remove_path:
-            sys.path.pop(0)
-
-        return config
-
-
 class ParallelContext(metaclass=SingletonMeta):
     """This class provides interface functions for users to get the parallel context,
     such as the global rank, the local rank, the world size, etc. of each device.

internlm/core/parallel/comm/isp.py

@@ -25,9 +25,8 @@
     expandKVPacked,
     reduce_scatter_raw,
 )
-from internlm.model.modules.embedding import Embedding1D
-from internlm.model.modules.linear import ParallelLinearWithCommExt
-from internlm.model.modules.utils import is_moe_param
+from internlm.model.model_ops.modules.linear import ParallelLinearWithCommExt
+from internlm.model.model_ops.modules.utils import is_moe_param
 from internlm.utils.common import SchedulerHook, UniqueChainMap, get_current_device
 from internlm.utils.utils import (
     CuSeqlenType,
@@ -179,14 +178,19 @@ class EmbeddingWeightParallelCommunicator:
     """
 
     def __init__(self, parallel_mode: ParallelMode) -> None:
+        from internlm.model.model_ops.modules.embedding import Embedding1D
+
+        self.embedding1d_cls = Embedding1D
         self.parallel_mode = parallel_mode
         self.gather_dim = 0
 
         self._cur_micro_step = 0
         self._num_micro_step = gpc.config.data.micro_num
 
-    def register_module_hook(self, module: Embedding1D) -> None:
-        assert isinstance(module, Embedding1D), "Embbeding weight parallel communicator is only support Embedding1D"
+    def register_module_hook(self, module: nn.Module) -> None:
+        assert isinstance(
+            module, self.embedding1d_cls
+        ), "Embbeding weight parallel communicator is only support Embedding1D"
 
         module.weight.evo_tensor = None
         self.gather_dim = 0 if module.vocab_parallel else 1

internlm/core/parallel/comm/tensor.py

@@ -23,8 +23,8 @@
     reduce_scatter_raw,
     split_forward_gather_backward,
 )
-from internlm.model.modules.embedding import Embedding1D
-from internlm.model.moe.moe import MoE
+from internlm.model.model_ops.modules.embedding import Embedding1D
+from internlm.model.model_ops.moe.moe import MoE
 
 # input gather dim
 _GATHER_DIM = 1  # shape: [batch, seqlen, dim] or [1, packlen, dim]

internlm/core/parallel/comm/zero.py

@@ -12,8 +12,8 @@
 from internlm.core.context import global_context as gpc
 from internlm.core.naive_amp import unwrap_naive_amp
 from internlm.core.parallel.comm.isp import ISPCommunicatorWrapper
-from internlm.model.modules.embedding import Embedding1D
-from internlm.model.modules.linear import ScaleColumnParallelLinear
+from internlm.model.model_ops.modules.embedding import Embedding1D
+from internlm.model.model_ops.modules.linear import ScaleColumnParallelLinear
 from internlm.solver.optimizer.utils import flatten
 
 

internlm/core/trainer_builder.py

@@ -17,14 +17,14 @@
 from internlm.data.train_state import get_train_state
 from internlm.eval.evaluation import evaluate_on_val_dls
 from internlm.initialize.initialize_trainer import initialize_trainer
-from internlm.model.losses.ce_loss import InternLoss
-from internlm.model.metrics import AccPerplex
+from internlm.model.model_ops.losses.ce_loss import InternLoss
+from internlm.model.model_ops.metrics import AccPerplex
+from internlm.model_inject.inject import inject_model
 from internlm.monitor.monitor import send_alert_message
 from internlm.train.pipeline import (
     get_scheduler_hooks,
     initialize_llm_profile,
     initialize_optimizer,
-    inject_model,
     load_new_batch,
     record_current_batch_training_metrics,
 )

internlm/env.py

@@ -0,0 +1 @@
+VERSION = "0.5.3"

internlm/eval/evaluation.py

@@ -9,7 +9,7 @@
 from internlm.core.context import global_context as gpc
 from internlm.core.parallel.shard import split_data_for_sequence_parallel
 from internlm.core.scheduler.pipeline_scheduler_1f1b import get_tensor_shape
-from internlm.model.metrics import AccPerplex, SchedulerMetricHook
+from internlm.model.model_ops.metrics import AccPerplex, SchedulerMetricHook
 from internlm.utils.common import get_current_device
 from internlm.utils.parallel import is_using_isp
 

internlm/initialize/__init__.py

@@ -1,14 +1,12 @@
 from .initialize_trainer import initialize_trainer
 from .launch import (
-    get_default_parser,
     initialize_distributed_env,
     launch_from_slurm,
     launch_from_torch,
     try_bind_numa,
 )
 
 __all__ = [
-    "get_default_parser",
     "initialize_trainer",
     "launch_from_slurm",
     "launch_from_torch",