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

Code before:
_ALLOWED_DTYPES = [torch.float16, torch.bfloat16, torch.float32]


_DEFAULT_PROMPT_TEMPLATES = {
    "llama": "
 
\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n
\n{input_text}
",
    "vicuna": "USER: {input_text}\n\nASSISTANT: ",
}


@dataclass
class InferenceConfig:
    """The inference configuration.

    Args:
        max_batch_size (int): Maximum batch size, defaults to 8.
        max_output_len (int): Maximum output length, defaults to 256.
        max_input_len (int): Maximum input length, defaults to 256.
        dtype (Union[str, torch.dtype]): The data type for weights and activations.
        prompt_template (Optional[str]): The prompt template for generation, defaults to None.
        do_sample (bool): Whether to use sampling for generation, defaults to False.
        beam_width (int): The maximum beam width used to initialize KV Cache, defaults to 1.
            During generation, the beam width provided as sampling parameter should be less than or equivalent to this value.
        prefill_ratio (Optional[float]): A controling ratio for prefill and decoding in running list, defaults to 1.2. We will do a step of prefill
            when the actual value exceeds this ratio.
        pad_input: Whether to pad all inputs to the max length.
        return_token_ids: Whether to return output token ids.
        early_stopping (Optional[bool]): Whether to stop the generation when all beam hypotheses have finished or not, defaults to False.
        top_k (Optional[int]): The number of highest probability vocabulary tokens to keep for top-k-filtering, defaults to None.
        top_p (Optional[float]): The cumulative probability threshold for retaining tokens with a total probability above it, defaults to None.
        min_p (Optional[float]): The minimum probability to keep for top-p filtering, defaults to None.
        block_size (int): The number of blocks in a logical block, defaults to 16.
        tp_size (int): Tensor parallel size, defaults to 1.
        pp_size (int): Pipeline parallel size, defaults to 1.
        micro_batch_size (int): the micro batch size, defaults to 1. Only useful when `pp_size` > 1.
        micro_batch_buffer_size (int): the buffer size for micro batch. Normally, it should be the same as the number of pipeline stages.

    """

    # NOTE: arrange configs according to their importance and frequency of usage

    # runtime limit
    max_batch_size: int = 8
    max_output_len: int = 256
    max_input_len: int = 256

    # general configs
    dtype: Union[str, torch.dtype] = torch.float16  # use fp16 by default

    # generation configs
    prompt_template: Optional[str] = None
    do_sample: bool = False

Review comment (suggestion):
`return_token_ids` is not a really a configuration. It means whether we want to have this control flow to return extra outputs to the user, so it is more associated with the code logic than generation configurations. I would recommend to do this instead:
```python
class Engine:

    def generate(self, ..., return_token_ids=False):
          ...

          if return_token_ids:
                  return outputs, token_ids
          else:
                  return outputs
```

Code after:
_ALLOWED_DTYPES = [torch.float16, torch.bfloat16, torch.float32]


_DEFAULT_PROMPT_TEMPLATES = {
    "llama": "
 
\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n
\n{input_text}
",
    "vicuna": "USER: {input_text}\n\nASSISTANT: ",
}


@dataclass
class InferenceConfig:
    """The inference configuration.

    Args:
        max_batch_size (int): Maximum batch size, defaults to 8.
        max_output_len (int): Maximum output length, defaults to 256.
        max_input_len (int): Maximum input length, defaults to 256.
        dtype (Union[str, torch.dtype]): The data type for weights and activations.
        prompt_template (Optional[str]): The prompt template for generation, defaults to None.
        do_sample (bool): Whether to use sampling for generation, defaults to False.
        beam_width (int): The maximum beam width used to initialize KV Cache, defaults to 1.
            During generation, the beam width provided as sampling parameter should be less than or equivalent to this value.
        prefill_ratio (Optional[float]): A controling ratio for prefill and decoding in running list, defaults to 1.2. We will do a step of prefill
            when the actual value exceeds this ratio.
        pad_input: Whether to pad all inputs to the max length.
        early_stopping (Optional[bool]): Whether to stop the generation when all beam hypotheses have finished or not, defaults to False.
        top_k (Optional[int]): The number of highest probability vocabulary tokens to keep for top-k-filtering, defaults to None.
        top_p (Optional[float]): The cumulative probability threshold for retaining tokens with a total probability above it, defaults to None.
        min_p (Optional[float]): The minimum probability to keep for top-p filtering, defaults to None.
        block_size (int): The number of blocks in a logical block, defaults to 16.
        tp_size (int): Tensor parallel size, defaults to 1.
        pp_size (int): Pipeline parallel size, defaults to 1.
        micro_batch_size (int): the micro batch size, defaults to 1. Only useful when `pp_size` > 1.
        micro_batch_buffer_size (int): the buffer size for micro batch. Normally, it should be the same as the number of pipeline stages.

    """

    # NOTE: arrange configs according to their importance and frequency of usage

    # runtime limit
    max_batch_size: int = 8
    max_output_len: int = 256
    max_input_len: int = 256

    # general configs
    dtype: Union[str, torch.dtype] = torch.float16  # use fp16 by default

    # generation configs
    prompt_template: Optional[str] = None
    do_sample: bool = False
    beam_width: int = 1  # TODO: beam search is not support for now