Repository: hpcaitech/ColossalAI
File: colossalai/inference/tensor_parallel/engine.py (Python)

Code before:
                enable_tensor_parallelism=False,
                enable_fused_normalization=False,
                enable_all_optimization=False,
                enable_flash_attention=False,
                enable_jit_fused=False,
                inference_only=True,
            )
        else:
            shard_config.inference_only = True
            shard_config.pipeline_stage_manager = None
            if shard_config.enable_tensor_parallelism:
                self.tp_size = shard_config.tensor_parallel_size
        self._init_manager()

        return shard_config

    def _shard_model_by(self, shardformer: ShardFormer, model: nn.Module) -> None:
        """Shard original model by the given ShardFormer and store the sharded model."""
        assert (
            self.tp_size == shardformer.shard_config.tensor_parallel_size
        ), "Discrepancy between the tp size of TPInferEngine and the tp size of shard config"
        model_name = model.__class__.__name__
        assert model_name in self.supported_models, f"Unsupported model cls {model_name} for TP inference."

        if self.shard_config.inference_gptq:
            policy = get_autopolicy(model.model, inference_only=True)
            self.model, _ = shardformer.optimize(model.model, policy)
            self._post_init_gptq_buffer(self.model)
        else:
            policy = get_autopolicy(model, inference_only=True)
            self.model, _ = shardformer.optimize(model, policy)

        self.model = self.model.cuda()

    @property
    def supported_models(self) -> List[str]:
        return _supported_models

    def generate(self, input_tokens: Union[BatchEncoding, dict, list, torch.Tensor], **generate_kwargs) -> torch.Tensor:
        """Generate token sequence.

        Args:
            input_tokens: could be one of the following types
                1. BatchEncoding or dict (e.g. tokenizer batch_encode)
                2. list of input token ids (e.g. appended result of tokenizer encode)
                3. torch.Tensor (e.g. tokenizer encode with return_tensors='pt')
        Returns:
            torch.Tensor: The returned sequence is given inputs + generated_tokens.
        """
        if isinstance(input_tokens, torch.Tensor):
            input_tokens = dict(input_ids=input_tokens, attention_mask=torch.ones_like(input_tokens, dtype=torch.bool))

Review comment (performance):
The old version is clearer and simpler. You can `to_optimize_model = model.model is self.shard_config.inference_gptq else model`, then follow the old version

Code after:
                enable_tensor_parallelism=False,
                enable_fused_normalization=False,
                enable_all_optimization=False,
                enable_flash_attention=False,
                enable_jit_fused=False,
                inference_only=True,
            )
        else:
            shard_config.inference_only = True
            shard_config.pipeline_stage_manager = None
            if shard_config.enable_tensor_parallelism:
                self.tp_size = shard_config.tensor_parallel_size
        self._init_manager()

        return shard_config

    def _shard_model_by(self, shardformer: ShardFormer, model: nn.Module) -> None:
        """Shard original model by the given ShardFormer and store the sharded model."""
        assert (
            self.tp_size == shardformer.shard_config.tensor_parallel_size
        ), "Discrepancy between the tp size of TPInferEngine and the tp size of shard config"
        model_name = model.__class__.__name__
        assert model_name in self.supported_models, f"Unsupported model cls {model_name} for TP inference."

        model = model.model if self.shard_config.inference_gptq else model

        policy = get_autopolicy(model, inference_only=True)
        self.model, _ = shardformer.optimize(model, policy)

        if self.shard_config.inference_gptq:
            self._post_init_gptq_buffer(self.model)

        self.model = self.model.cuda()

    @property
    def supported_models(self) -> List[str]:
        return _supported_models

    def generate(self, input_tokens: Union[BatchEncoding, dict, list, torch.Tensor], **generate_kwargs) -> torch.Tensor:
        """Generate token sequence.

        Args:
            input_tokens: could be one of the following types
                1. BatchEncoding or dict (e.g. tokenizer batch_encode)
                2. list of input token ids (e.g. appended result of tokenizer encode)
                3. torch.Tensor (e.g. tokenizer encode with return_tensors='pt')
        Returns:
            torch.Tensor: The returned sequence is given inputs + generated_tokens.
        """
        if isinstance(input_tokens, torch.Tensor):
            input_tokens = dict(input_ids=input_tokens, attention_mask=torch.ones_like(input_tokens, dtype=torch.bool))