Repository: hpcaitech/ColossalAI
File: colossalai/inference/executor/rpc_worker.py (Python)

Code before:
            kv_max_split_num=kv_max_split_num,
            head_dim=head_dim,
            dtype=self.dtype,
            device=get_accelerator().get_current_device(),
        )

        self.fd_inter_tensor = fd_inter_tensor

    def _init_model(self, model_or_path: Union[nn.Module, str], model_policy: Policy = None):
        """
        Shard model or/and Load weight

        Shard model: When we set tp_size > 1, we will shard the model by given model_policy.
        Load Weight: If we pass a local model path, we will load the model weight by checkpoint_io. If it is a remote-transformer url, we will use `AutoModel.from_pretrained` api of transformers lib

        Args:
            model_or_path Union[nn.Module, str]: path to the checkpoint or model of transformer format.
            model_policy (Policy): the policy to replace the model
        """

        if isinstance(model_or_path, str):
            is_local = os.path.isdir(model_or_path)
            try:
                hf_config = AutoConfig.from_pretrained(model_or_path, trust_remote_code=True, torch_dtype=self.dtype)
                arch = getattr(hf_config, "architectures")[0]
                if is_local:
                    with LazyInitContext(default_device="cuda"):
                        model = _SUPPORTED_MODELS[arch](hf_config)
                else:
                    # load the real checkpoint
                    model = _SUPPORTED_MODELS[arch].from_pretrained(
                        model_or_path, trust_remote_code=True, torch_dtype=self.dtype
                    )
            except Exception as e:
                logger.error(
                    f"An exception occurred during loading model: {e}, model should be loaded by transformers\n"
                )
        else:
            model = model_or_path

        self.model_config = model.config

        torch.cuda.empty_cache()
        init_gpu_memory = torch.cuda.mem_get_info()[0]

        self.device = get_accelerator().get_current_device()
        torch.cuda.set_device(self.device)
        if self.verbose:
            logger.info(f"the device is {self.device}")

        model = model.to(dtype=self.dtype, non_blocking=False).eval()

Review comment (suggestion):
For this method (`_init_model`), please also keep up to date with that (`init_model`) in `engine.py`. For example, the handling of condition of no policy is provided.

Code after:
            head_dim=head_dim,
            dtype=self.dtype,
            device=get_accelerator().get_current_device(),
        )

        self.fd_inter_tensor = fd_inter_tensor

    def _init_model(self, model_or_path: Union[nn.Module, str], model_policy: Policy = None):
        """
        Shard model or/and Load weight

        Shard model: When we set tp_size > 1, we will shard the model by given model_policy.
        Load Weight: If we pass a local model path, we will load the model weight by checkpoint_io. If it is a remote-transformer url, we will use `AutoModel.from_pretrained` api of transformers lib

        Args:
            model_or_path Union[nn.Module, str]: path to the checkpoint or model of transformer format.
            model_policy (Policy): the policy to replace the model
        """

        pretrained_path = None
        if isinstance(model_or_path, str):
            import colossalai.interface.pretrained as pretrained_utils

            try:
                hf_config = AutoConfig.from_pretrained(model_or_path, trust_remote_code=True, torch_dtype=self.dtype)
                arch = getattr(hf_config, "architectures")[0]
                if arch is "BaichuanForCausalLM":
                    self.logger.warning(
                        "Attention ! We use lazy init by default, which could be faster for model loading. For baichuan model, the output maybe have a slight difference with transformers"
                    )
                ctx = LazyInitContext(default_device="cuda")
                with ctx:
                    model = _SUPPORTED_MODELS[arch].from_pretrained(
                        model_or_path, trust_remote_code=True, torch_dtype=self.dtype
                    )
                pretrained_path = pretrained_utils.get_pretrained_path(model)
            except Exception as e:
                logger.error(
                    f"An exception occurred during loading model: {e}, model should be loaded by transformers\n"
                )
        else:
            model = model_or_path

        self.model_config = model.config

        torch.cuda.empty_cache()
        init_gpu_memory = torch.cuda.mem_get_info()[0]

        self.device = get_accelerator().get_current_device()
        torch.cuda.set_device(self.device)
        if self.verbose: