Function decode_with_vllm

src/MaxText/vllm_decode.py:90–185

Decode using vLLM with a MaxText model implementation. Args: model_name: Name of the model for MaxText. hf_model_name: Path to the Hugging Face model. hf_config_path: Path to the local Hugging Face model config. load_parameters_path: Path to load model parameters from. ici_dat

(
    model_name: str,
    hf_model_name: str,
    hf_config_path: str,
    load_parameters_path: str,
    ici_data_parallelism: int,
    ici_tensor_parallelism: int,
    ici_expert_parallelism: int,
    max_prefill_length: int,
    max_target_length: int,
    gpu_memory_utilization: float,
    enable_expert_parallel: bool,
    prompt: str,
    decode_sampling_temperature: float,
    decode_sampling_nucleus_p: float,
    decode_sampling_top_k: float,
)

def decode_with_vllm(
    model_name: str,
    hf_model_name: str,
    hf_config_path: str,
    load_parameters_path: str,
    ici_data_parallelism: int,
    ici_tensor_parallelism: int,
    ici_expert_parallelism: int,
    max_prefill_length: int,
    max_target_length: int,
    gpu_memory_utilization: float,
    enable_expert_parallel: bool,
    prompt: str,
    decode_sampling_temperature: float,
    decode_sampling_nucleus_p: float,
    decode_sampling_top_k: float,
) -> None:
  """Decode using vLLM with a MaxText model implementation.

  Args:
    model_name: Name of the model for MaxText.
    hf_model_name: Path to the Hugging Face model.
    hf_config_path: Path to the local Hugging Face model config.
    load_parameters_path: Path to load model parameters from.
    ici_data_parallelism: Size of the data parallelism dimension.
    ici_tensor_parallelism: Size of the non-expert tensor parallelism dimension.
    ici_expert_parallelism: Size of the MoE expert parallelism dimension.
    max_prefill_length: Maximum prefill length.
    max_target_length: Maximum total context length (MCL).
    gpu_memory_utilization: Fraction of GPU memory to be used for the model executor.
    enable_expert_parallel: Whether to enable expert parallelism.
    prompt: The prompt to decode.
    decode_sampling_temperature: Temperature for sampling.
    decode_sampling_nucleus_p: Nucleus sampling probability.
    decode_sampling_top_k: Top-k sampling probability.
  """

  # Prepare vLLM Arguments
  vllm_args = {}
  vllm_args["additional_config"] = {}

  # Core vLLM Arguments
  vllm_args["model"] = hf_model_name
  vllm_args["max_model_len"] = max_target_length
  vllm_args["tensor_parallel_size"] = ici_tensor_parallelism
  vllm_args["data_parallel_size"] = ici_data_parallelism
  vllm_args["enable_expert_parallel"] = enable_expert_parallel
  vllm_args["hf_config_path"] = hf_config_path
  vllm_args["gpu_memory_utilization"] = gpu_memory_utilization

  if load_parameters_path is None:
    vllm_args["load_format"] = "dummy"

  # Prepare MaxText and sharding configs (Parallelism is dynamic)
  vllm_args["additional_config"]["maxtext_config"] = {
      "model_name": model_name,
      "max_target_length": max_target_length,
      "weight_dtype": "bfloat16",