config.py 65 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505
  1. import enum
  2. import json
  3. import os
  4. from dataclasses import dataclass, field, fields
  5. from typing import (TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Tuple,
  6. Union)
  7. import torch
  8. from loguru import logger
  9. from transformers import PretrainedConfig
  10. from aphrodite.common.utils import (cuda_device_count_stateless,
  11. get_cpu_memory, is_cpu, is_hip, is_neuron,
  12. is_tpu, is_xpu,
  13. update_environment_variables)
  14. from aphrodite.modeling.models import ModelRegistry
  15. from aphrodite.quantization import QUANTIZATION_METHODS
  16. from aphrodite.transformers_utils.config import get_config, get_hf_text_config
  17. if TYPE_CHECKING:
  18. from ray.util.placement_group import PlacementGroup
  19. from aphrodite.modeling.model_loader.loader import BaseModelLoader
  20. # If true, will load models from ModelScope instead of Hugging Face Hub.
  21. APHRODITE_USE_MODELSCOPE = os.environ.get("APHRODITE_USE_MODELSCOPE",
  22. "False").lower() == "true"
  23. _GB = 1 << 30
  24. _EMBEDDING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
  25. class ModelConfig:
  26. """Configuration for the model.
  27. Args:
  28. model: Name or path of the huggingface model to use.
  29. tokenizer: Name or path of the huggingface tokenizer to use.
  30. tokenizer_mode: Tokenizer mode. "auto" will use the fast tokenizer if
  31. available, and "slow" will always use the slow tokenizer.
  32. trust_remote_code: Trust remote code (e.g., from HuggingFace) when
  33. downloading the model and tokenizer.
  34. dtype: Data type for model weights and activations. The "auto" option
  35. will use FP16 precision for FP32 and FP16 models, and BF16 precision
  36. for BF16 models.
  37. seed: Random seed for reproducibility.
  38. revision: The specific model version to use. It can be a branch name,
  39. a tag name, or a commit id. If unspecified, will use the default
  40. version.
  41. code_revision: The specific revision to use for the model code on
  42. Hugging Face Hub. It can be a branch name, a tag name, or a
  43. commit id. If unspecified, will use the default version.
  44. rope_scaling: Dictionary containing the scaling configuration for the
  45. RoPE embeddings. When using this flag, don't update
  46. `max_position_embeddings` to the expected new maximum.
  47. tokenizer_revision: The specific tokenizer version to use. It can be a
  48. branch name, a tag name, or a commit id. If unspecified, will use
  49. the default version.
  50. max_model_len: Maximum length of a sequence (including prompt and
  51. output). If None, will be derived from the model.
  52. quantization: Quantization method that was used to quantize the model
  53. weights. If None, we assume the model weights are not quantized.
  54. load_in_4bit: Whether to load the FP16 model in AutoQuant 4bit
  55. format. Works with AWQ models as well as FP16.
  56. load_in_8bit: Whether to load the FP16 model in 8bit format. Slower
  57. than load_in_smooth in terms of throughput.
  58. load_in_smooth: Whether to load the FP16 model in smoothquant format.
  59. deepspeed_fp_bits: Number of bits to use for DeepSpeed FP quantization.
  60. Supported number of bits are: 4, 6, 8, 12.
  61. quantization_param_path: Path to JSON file containing scaling factors.
  62. Used to load KV cache scaling factors into the model when KV cache
  63. type is FP8_E4M3 on ROCm (AMD GPU). In the future these will also
  64. be used to load activation and weight scaling factors when the
  65. model dtype is FP8_E4M3 on ROCm.
  66. enforce_eager: Whether to enforce eager execution. If True, we will
  67. disable CUDA graph and always execute the model in eager mode.
  68. If False, we will use CUDA graph and eager execution in hybrid.
  69. max_context_len_to_capture: Maximum context len covered by CUDA graphs.
  70. When a sequence has context length larger than this, we fall back
  71. to eager mode (DEPRECATED. Use max_seq_len_to_capture instead).
  72. max_seq_len_to_capture: Maximum sequence len covered by CUDA graphs.
  73. When a sequence has context length larger than this, we fall back
  74. to eager mode
  75. disable_sliding_window: Whether to disable sliding window. If True,
  76. we will disable the sliding window functionality of the model.
  77. If the model does not support sliding window, this argument is
  78. ignored.
  79. skip_tokenizer_init: If true, skip initialization of tokenizer and
  80. detokenizer.
  81. """
  82. def __init__(
  83. self,
  84. model: str,
  85. tokenizer: str,
  86. tokenizer_mode: str,
  87. trust_remote_code: bool,
  88. dtype: Union[str, torch.dtype],
  89. seed: int,
  90. revision: Optional[str] = None,
  91. code_revision: Optional[str] = None,
  92. rope_scaling: Optional[dict] = None,
  93. rope_theta: Optional[float] = None,
  94. tokenizer_revision: Optional[str] = None,
  95. max_model_len: Optional[int] = None,
  96. quantization: Optional[str] = None,
  97. load_in_4bit: bool = False,
  98. load_in_8bit: bool = False,
  99. load_in_smooth: bool = False,
  100. deepspeed_fp_bits: Optional[int] = None,
  101. quantization_param_path: Optional[str] = None,
  102. enforce_eager: bool = True,
  103. max_context_len_to_capture: Optional[int] = None,
  104. max_seq_len_to_capture: Optional[int] = None,
  105. max_logprobs: int = 5,
  106. disable_sliding_window: bool = False,
  107. skip_tokenizer_init: bool = False,
  108. ) -> None:
  109. self.model = model
  110. self.tokenizer = tokenizer
  111. self.tokenizer_mode = tokenizer_mode
  112. self.trust_remote_code = trust_remote_code
  113. self.seed = seed
  114. self.revision = revision
  115. self.code_revision = code_revision
  116. self.rope_scaling = rope_scaling
  117. self.rope_theta = rope_theta
  118. # The tokenizer version is consistent with the model version by default.
  119. if tokenizer_revision is None:
  120. self.tokenizer_revision = revision
  121. else:
  122. self.tokenizer_revision = tokenizer_revision
  123. self.quantization = quantization
  124. self.load_in_4bit = load_in_4bit
  125. self.load_in_8bit = load_in_8bit
  126. self.load_in_smooth = load_in_smooth
  127. self.deepspeed_fp_bits = deepspeed_fp_bits
  128. self.quantization_param_path = quantization_param_path
  129. self.enforce_eager = enforce_eager
  130. self.max_context_len_to_capture = max_context_len_to_capture
  131. if self.max_context_len_to_capture is not None:
  132. raise ValueError("`max_context_len_to_capture` is deprecated. "
  133. "Use `max_seq_len_to_capture` instead.")
  134. self.max_seq_len_to_capture = (max_seq_len_to_capture
  135. or max_context_len_to_capture)
  136. self.max_logprobs = max_logprobs
  137. self.disable_sliding_window = disable_sliding_window
  138. self.skip_tokenizer_init = skip_tokenizer_init
  139. self.hf_config = get_config(self.model, trust_remote_code, revision,
  140. code_revision, rope_scaling, rope_theta)
  141. self.hf_text_config = get_hf_text_config(self.hf_config)
  142. self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
  143. self.max_model_len = _get_and_verify_max_len(
  144. hf_config=self.hf_text_config,
  145. max_model_len=max_model_len,
  146. disable_sliding_window=self.disable_sliding_window,
  147. sliding_window_len=self.get_hf_config_sliding_window(),
  148. rope_scaling_arg=self.rope_scaling)
  149. if not self.skip_tokenizer_init:
  150. self._verify_tokenizer_mode()
  151. self._verify_embedding_mode()
  152. self._verify_quantization()
  153. self._verify_cuda_graph()
  154. def _verify_tokenizer_mode(self) -> None:
  155. tokenizer_mode = self.tokenizer_mode.lower()
  156. if tokenizer_mode not in ["auto", "slow"]:
  157. raise ValueError(
  158. f"Unknown tokenizer mode: {self.tokenizer_mode}. Must be "
  159. "either 'auto' or 'slow'.")
  160. self.tokenizer_mode = tokenizer_mode
  161. def _verify_embedding_mode(self) -> None:
  162. architectures = getattr(self.hf_config, "architectures", [])
  163. self.embedding_mode = any(
  164. ModelRegistry.is_embedding_model(arch) for arch in architectures)
  165. def _parse_quant_hf_config(self):
  166. quant_cfg = getattr(self.hf_config, "quantization_config", None)
  167. if quant_cfg is None:
  168. # compress-tensors uses a "compression_config" key
  169. quant_cfg = getattr(self.hf_config, "compression_config", None)
  170. return quant_cfg
  171. def _verify_quantization(self) -> None:
  172. supported_quantization = [*QUANTIZATION_METHODS]
  173. rocm_supported_quantization = ["gptq", "squeezellm"]
  174. if self.quantization is not None:
  175. self.quantization = self.quantization.lower()
  176. # Parse quantization method from the HF model config, if available.
  177. quant_cfg = self._parse_quant_hf_config()
  178. if quant_cfg is not None:
  179. quant_method = quant_cfg.get("quant_method", "").lower()
  180. # Detect which checkpoint is it
  181. for _, method in QUANTIZATION_METHODS.items():
  182. quantization_override = method.override_quantization_method(
  183. quant_cfg, self.quantization)
  184. if quantization_override:
  185. quant_method = quantization_override
  186. self.quantization = quantization_override
  187. break
  188. # Verify quantization configurations.
  189. if self.quantization is None:
  190. self.quantization = quant_method
  191. elif self.quantization != quant_method:
  192. raise ValueError(
  193. "Quantization method specified in the model config "
  194. f"({quant_method}) does not match the quantization "
  195. f"method specified in the `quantization` argument "
  196. f"({self.quantization}).")
  197. if self.load_in_4bit:
  198. # the kernels seem to not work with 4bit weight_only
  199. if torch.cuda.get_device_capability(0)[0] < 8:
  200. raise ValueError(
  201. "load_in_4bit quantization is not supported on GPUs with "
  202. "compute capability less than 8.0.")
  203. if self.quantization is None:
  204. self.quantization = "autoquant"
  205. self.hf_config.quantization_config = {
  206. "bits": 4,
  207. "quant_mode": "weight_only",
  208. "quant_method": "autoquant",
  209. "group_size": 128,
  210. "zero_point": True,
  211. "from_float": True
  212. }
  213. elif self.quantization == "awq":
  214. logger.warning("AWQ model is being loaded in 4bit autoquant "
  215. "format.")
  216. self.quantization = "autoquant"
  217. self.hf_config.quantization_config = {
  218. "zero_point": True,
  219. "q_group_size": 128,
  220. "w_bit": 4,
  221. "version": "gemm"
  222. }
  223. elif self.quantization != "autoquant":
  224. raise ValueError("4bit quantization is not supported in "
  225. f"{self.quantization}.")
  226. if self.load_in_8bit:
  227. if self.quantization is None:
  228. self.quantization = "autoquant"
  229. elif self.quantization != "autoquant":
  230. raise ValueError("8bit quantization is not supported in "
  231. f"{self.quantization}.")
  232. self.hf_config.quantization_config = {
  233. "bits": 8,
  234. "quant_mode": "llm_int8",
  235. "quant_method": "autoquant",
  236. "group_size": 128,
  237. "zero_point": True,
  238. "from_float": True
  239. }
  240. self.enforce_eager = True
  241. if self.load_in_smooth:
  242. if self.quantization is None:
  243. self.quantization = "autoquant"
  244. elif self.quantization != "autoquant":
  245. raise ValueError("Smooth quantization is not supported in "
  246. f"{self.quantization}.")
  247. self.hf_config.quantization_config = {
  248. "bits": 8,
  249. "quant_mode": "smoothquant",
  250. "quant_method": "autoquant",
  251. "group_size": 128,
  252. "zero_point": True,
  253. "from_float": True
  254. }
  255. self.enforce_eager = True
  256. if self.quantization == "deepspeedfp":
  257. gs = 32 if self.deepspeed_fp_bits == 4 else 128
  258. self.hf_config.quantization_config = {
  259. "bits": self.deepspeed_fp_bits,
  260. "group_size": int(os.environ.get("DEEPSPEED_GROUP_SIZE", gs)),
  261. "quant_method": "deepspeedfp"
  262. }
  263. if self.quantization is not None:
  264. if self.quantization not in supported_quantization:
  265. raise ValueError(
  266. f"Unknown quantization method: {self.quantization}. Must "
  267. f"be one of {supported_quantization}.")
  268. if is_hip(
  269. ) and self.quantization not in rocm_supported_quantization:
  270. raise ValueError(
  271. f"{self.quantization} quantization is currently not "
  272. "supported in ROCm.")
  273. if (self.quantization
  274. not in ["fp8", "marlin", "gptq_marlin_24", "gptq_marlin"]):
  275. logger.warning(
  276. f"{self.quantization} quantization is not fully "
  277. "optimized yet. The speed can be slower than "
  278. "non-quantized models.")
  279. if self.quantization == "deepspeedfp" and self.deepspeed_fp_bits \
  280. is None:
  281. raise ValueError(
  282. "deepspeed_fp_bits must be specified when using "
  283. "deepspeedfp quantization.")
  284. def _verify_cuda_graph(self) -> None:
  285. if self.max_seq_len_to_capture is None:
  286. self.max_seq_len_to_capture = self.max_model_len
  287. self.max_seq_len_to_capture = min(self.max_seq_len_to_capture,
  288. self.max_model_len)
  289. def verify_with_parallel_config(
  290. self,
  291. parallel_config: "ParallelConfig",
  292. ) -> None:
  293. total_num_attention_heads = getattr(self.hf_text_config,
  294. "num_attention_heads", 0)
  295. tensor_parallel_size = parallel_config.tensor_parallel_size
  296. if total_num_attention_heads % tensor_parallel_size != 0:
  297. raise ValueError(
  298. f"Total number of attention heads ({total_num_attention_heads})"
  299. " must be divisible by tensor parallel size "
  300. f"({tensor_parallel_size}).")
  301. total_num_hidden_layers = getattr(self.hf_text_config,
  302. "num_hidden_layers", 0)
  303. pipeline_parallel_size = parallel_config.pipeline_parallel_size
  304. if total_num_hidden_layers % pipeline_parallel_size != 0:
  305. raise ValueError(
  306. f"Total number of hidden layers ({total_num_hidden_layers}) "
  307. "must be divisible by pipeline parallel size "
  308. f"({pipeline_parallel_size}).")
  309. if self.quantization == "bitsandbytes" and (
  310. parallel_config.tensor_parallel_size > 1
  311. or parallel_config.pipeline_parallel_size > 1):
  312. raise ValueError(
  313. "BitAndBytes quantization with TP or PP is not supported yet.")
  314. def get_hf_config_sliding_window(self) -> Optional[int]:
  315. """Get the sliding window size, or None if disabled.
  316. """
  317. # Some models, like Qwen2 and Qwen1.5, use `use_sliding_window` in
  318. # addition to sliding window size. We check if that field is present
  319. # and if it's False, return None.
  320. if (hasattr(self.hf_text_config, "use_sliding_window")
  321. and not self.hf_text_config.use_sliding_window):
  322. return None
  323. return getattr(self.hf_text_config, "sliding_window", None)
  324. def get_sliding_window(self) -> Optional[int]:
  325. """Get the sliding window size, or None if disabled.
  326. """
  327. # If user disables sliding window, return None.
  328. if self.disable_sliding_window:
  329. return None
  330. # Otherwise get the value from the hf config.
  331. return self.get_hf_config_sliding_window()
  332. def get_vocab_size(self) -> int:
  333. return self.hf_text_config.vocab_size
  334. def get_hidden_size(self) -> int:
  335. return self.hf_text_config.hidden_size
  336. def get_head_size(self) -> int:
  337. if hasattr(self.hf_text_config, "head_dim"):
  338. return self.hf_text_config.head_dim
  339. # FIXME: This may not be true for all models.
  340. return (self.hf_text_config.hidden_size //
  341. self.hf_text_config.num_attention_heads)
  342. def get_total_num_kv_heads(self) -> int:
  343. """Returns the total number of KV heads."""
  344. # For GPTBigCode & Falcon:
  345. # NOTE: for falcon, when new_decoder_architecture is True, the
  346. # multi_query flag is ignored and we use n_head_kv for the number of
  347. # KV heads.
  348. falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
  349. new_decoder_arch_falcon = (
  350. self.hf_config.model_type in falcon_model_types
  351. and getattr(self.hf_config, "new_decoder_architecture", False))
  352. if not new_decoder_arch_falcon and getattr(self.hf_text_config,
  353. "multi_query", False):
  354. # Multi-query attention, only one KV head.
  355. # Currently, tensor parallelism is not supported in this case.
  356. return 1
  357. # For DBRX and MPT
  358. if self.hf_config.model_type == "mpt":
  359. if "kv_n_heads" in self.hf_config.attn_config:
  360. return self.hf_config.attn_config["kv_n_heads"]
  361. return self.hf_config.num_attention_heads
  362. if self.hf_config.model_type == "dbrx":
  363. return getattr(self.hf_config.attn_config, "kv_n_heads",
  364. self.hf_config.num_attention_heads)
  365. attributes = [
  366. # For Falcon:
  367. "n_head_kv",
  368. "num_kv_heads",
  369. # For LLaMA-2:
  370. "num_key_value_heads",
  371. # For ChatGLM:
  372. "multi_query_group_num",
  373. ]
  374. for attr in attributes:
  375. num_kv_heads = getattr(self.hf_text_config, attr, None)
  376. if num_kv_heads is not None:
  377. return num_kv_heads
  378. # For non-grouped-query attention models, the number of KV heads is
  379. # equal to the number of attention heads.
  380. return self.hf_text_config.num_attention_heads
  381. def get_num_kv_heads(self, parallel_config: "ParallelConfig") -> int:
  382. """Returns the number of KV heads per GPU."""
  383. total_num_kv_heads = self.get_total_num_kv_heads()
  384. # If tensor parallelism is used, we divide the number of KV heads by
  385. # the tensor parallel size. We will replicate the KV heads in the
  386. # case where the number of KV heads is smaller than the tensor
  387. # parallel size so each GPU has at least one KV head.
  388. return max(1,
  389. total_num_kv_heads // parallel_config.tensor_parallel_size)
  390. def get_num_attention_heads(self,
  391. parallel_config: "ParallelConfig") -> int:
  392. num_heads = getattr(self.hf_text_config, "num_attention_heads", 0)
  393. return num_heads // parallel_config.tensor_parallel_size
  394. def get_num_layers(self, parallel_config: "ParallelConfig") -> int:
  395. total_num_hidden_layers = self.hf_text_config.num_hidden_layers
  396. return total_num_hidden_layers // parallel_config.pipeline_parallel_size
  397. class CacheConfig:
  398. """Configuration for the KV cache.
  399. Args:
  400. block_size: Size of a cache block in number of tokens.
  401. gpu_memory_utilization: Fraction of GPU memory to use for the
  402. Aphrodite execution.
  403. swap_space: Size of the CPU swap space per GPU (in GiB).
  404. cache_dtype: Data type for kv cache storage.
  405. num_gpu_blocks_override: Number of GPU blocks to use. This overrides the
  406. profiled num_gpu_blocks if specified. Does nothing if None.
  407. """
  408. def __init__(
  409. self,
  410. block_size: int,
  411. gpu_memory_utilization: float,
  412. swap_space: int,
  413. cache_dtype: str,
  414. num_gpu_blocks_override: Optional[int] = None,
  415. sliding_window: Optional[int] = None,
  416. enable_prefix_caching: bool = False,
  417. ) -> None:
  418. self.block_size = block_size
  419. self.gpu_memory_utilization = gpu_memory_utilization
  420. self.swap_space_bytes = swap_space * _GB
  421. self.num_gpu_blocks_override = num_gpu_blocks_override
  422. self.cache_dtype = cache_dtype
  423. self.sliding_window = sliding_window
  424. self.enable_prefix_caching = enable_prefix_caching
  425. self._verify_args()
  426. self._verify_cache_dtype()
  427. self._verify_prefix_caching()
  428. # Will be set after profiling.
  429. self.num_gpu_blocks = None
  430. self.num_cpu_blocks = None
  431. def metrics_info(self):
  432. # convert cache_config to dict(key: str, value: str) for prometheus
  433. # metrics info
  434. return {key: str(value) for key, value in self.__dict__.items()}
  435. def _verify_args(self) -> None:
  436. if self.gpu_memory_utilization > 1.0:
  437. raise ValueError(
  438. "GPU memory utilization must be less than 1.0. Got "
  439. f"{self.gpu_memory_utilization}.")
  440. def _verify_cache_dtype(self) -> None:
  441. if self.cache_dtype == "auto":
  442. pass
  443. elif self.cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2"):
  444. logger.info(
  445. "Using fp8 data type to store kv cache. It reduces the GPU "
  446. "memory footprint and boosts the performance. "
  447. "Meanwhile, it may cause accuracy drop without a proper "
  448. "scaling factor")
  449. else:
  450. raise ValueError(f"Unknown kv cache dtype: {self.cache_dtype}")
  451. def _verify_prefix_caching(self) -> None:
  452. if not self.enable_prefix_caching:
  453. return
  454. if self.sliding_window is not None:
  455. raise NotImplementedError(
  456. "Prefix caching is not supported with sliding window. "
  457. "Run with --disable-sliding-window to use prefix caching.")
  458. if self.cache_dtype == "fp8":
  459. raise NotImplementedError(
  460. "Prefix caching is not supported for fp8 cache_dtype. "
  461. "Run with --kv-cache-dtype auto to use prefix caching.")
  462. def verify_with_parallel_config(
  463. self,
  464. parallel_config: "ParallelConfig",
  465. ) -> None:
  466. total_cpu_memory = get_cpu_memory()
  467. # FIXME: Here, it is assumed that the GPUs in a tensor parallel
  468. # group are in the same node. However, the GPUs may span multiple nodes.
  469. num_gpus_per_node = parallel_config.tensor_parallel_size
  470. cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node
  471. msg = (f"{cpu_memory_usage / _GB:.2f} GiB out of "
  472. f"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is "
  473. "allocated for the swap space.")
  474. if cpu_memory_usage > 0.7 * total_cpu_memory:
  475. raise ValueError("Too large swap space. " + msg)
  476. elif cpu_memory_usage > 0.4 * total_cpu_memory:
  477. logger.warning("Possibly too large swap space. " + msg)
  478. @dataclass
  479. class TokenizerPoolConfig:
  480. """Configuration for the tokenizer pool.
  481. Args:
  482. pool_size: Number of tokenizer workers in the pool.
  483. pool_type: Type of the pool.
  484. extra_config: Additional config for the pool.
  485. The way the config will be used depends on the
  486. pool type.
  487. """
  488. pool_size: int
  489. pool_type: str
  490. extra_config: dict
  491. def __post_init__(self):
  492. if self.pool_type not in ("ray", ):
  493. raise ValueError(f"Unknown pool type: {self.pool_type}")
  494. if not isinstance(self.extra_config, dict):
  495. raise ValueError("extra_config must be a dictionary.")
  496. @classmethod
  497. def create_config(
  498. cls, tokenizer_pool_size: int, tokenizer_pool_type: str,
  499. tokenizer_pool_extra_config: Optional[Union[str, dict]]
  500. ) -> Optional["TokenizerPoolConfig"]:
  501. """Create a TokenizerPoolConfig from the given parameters.
  502. If tokenizer_pool_size is 0, return None.
  503. Args:
  504. tokenizer_pool_size: Number of tokenizer workers in the pool.
  505. tokenizer_pool_type: Type of the pool.
  506. tokenizer_pool_extra_config: Additional config for the pool.
  507. The way the config will be used depends on the
  508. pool type. This can be a JSON string (will be parsed).
  509. """
  510. if tokenizer_pool_size:
  511. if isinstance(tokenizer_pool_extra_config, str):
  512. tokenizer_pool_extra_config_parsed = json.loads(
  513. tokenizer_pool_extra_config)
  514. else:
  515. tokenizer_pool_extra_config_parsed = (
  516. tokenizer_pool_extra_config or {})
  517. tokenizer_pool_config = cls(tokenizer_pool_size,
  518. tokenizer_pool_type,
  519. tokenizer_pool_extra_config_parsed)
  520. else:
  521. tokenizer_pool_config = None
  522. return tokenizer_pool_config
  523. class LoadFormat(str, enum.Enum):
  524. AUTO = "auto"
  525. PT = "pt"
  526. SAFETENSORS = "safetensors"
  527. NPCACHE = "npcache"
  528. DUMMY = "dummy"
  529. TENSORIZER = "tensorizer"
  530. SHARDED_STATE = "sharded_state"
  531. BITSANDBYTES = "bitsandbytes"
  532. @dataclass
  533. class LoadConfig:
  534. """
  535. download_dir: Directory to download and load the weights, default to the
  536. default cache directory of huggingface.
  537. load_format: The format of the model weights to load:
  538. "auto" will try to load the weights in the safetensors format and
  539. fall back to the pytorch bin format if safetensors format is
  540. not available.
  541. "pt" will load the weights in the pytorch bin format.
  542. "safetensors" will load the weights in the safetensors format.
  543. "npcache" will load the weights in pytorch format and store
  544. a numpy cache to speed up the loading.
  545. "dummy" will initialize the weights with random values, which is
  546. mainly for profiling.
  547. "tensorizer" will use CoreWeave's tensorizer library for
  548. fast weight loading.
  549. """
  550. load_format: Union[str, LoadFormat, "BaseModelLoader"] = LoadFormat.AUTO
  551. download_dir: Optional[str] = None
  552. model_loader_extra_config: Optional[Union[str, dict]] = field(
  553. default_factory=dict)
  554. def __post_init__(self):
  555. model_loader_extra_config = self.model_loader_extra_config or {}
  556. if isinstance(model_loader_extra_config, str):
  557. self.model_loader_extra_config = json.loads(
  558. model_loader_extra_config)
  559. self._verify_load_format()
  560. def _verify_load_format(self) -> None:
  561. if not isinstance(self.load_format, str):
  562. return
  563. load_format = self.load_format.lower()
  564. self.load_format = LoadFormat(load_format)
  565. rocm_not_supported_load_format: List[str] = []
  566. if is_hip() and load_format in rocm_not_supported_load_format:
  567. rocm_supported_load_format = [
  568. f for f in LoadFormat.__members__
  569. if (f not in rocm_not_supported_load_format)
  570. ]
  571. raise ValueError(
  572. f"load format '{load_format}' is not supported in ROCm. "
  573. f"Supported load formats are "
  574. f"{rocm_supported_load_format}")
  575. class ParallelConfig:
  576. """Configuration for the distributed execution.
  577. Args:
  578. pipeline_parallel_size: Number of pipeline parallel groups.
  579. tensor_parallel_size: Number of tensor parallel groups.
  580. worker_use_ray: Deprecated, use distributed_executor_backend instead.
  581. max_parallel_loading_workers: Maximum number of multiple batches
  582. when load model sequentially. To avoid RAM OOM when using tensor
  583. parallel and large models.
  584. disable_custom_all_reduce: Disable the custom all-reduce kernel and
  585. fall back to NCCL.
  586. tokenizer_pool_config: Config for the tokenizer pool.
  587. If None, will use synchronous tokenization.
  588. ray_workers_use_nsight: Whether to profile Ray workers with nsight, see
  589. https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html#profiling-nsight-profiler.
  590. placement_group: ray distributed model workers placement group.
  591. distributed_executor_backend: Backend to use for distributed model
  592. workers, either "ray" or "mp" (multiprocessing). If either
  593. pipeline_parallel_size or tensor_parallel_size is greater than 1,
  594. will default to "ray" if Ray is installed or "mp" otherwise.
  595. """
  596. def __init__(
  597. self,
  598. pipeline_parallel_size: int,
  599. tensor_parallel_size: int,
  600. worker_use_ray: Optional[bool] = None,
  601. max_parallel_loading_workers: Optional[int] = None,
  602. disable_custom_all_reduce: bool = False,
  603. tokenizer_pool_config: Optional[TokenizerPoolConfig] = None,
  604. ray_workers_use_nsight: bool = False,
  605. placement_group: Optional["PlacementGroup"] = None,
  606. distributed_executor_backend: Optional[str] = None,
  607. ) -> None:
  608. self.pipeline_parallel_size = pipeline_parallel_size
  609. self.tensor_parallel_size = tensor_parallel_size
  610. self.distributed_executor_backend = distributed_executor_backend
  611. self.max_parallel_loading_workers = max_parallel_loading_workers
  612. self.disable_custom_all_reduce = disable_custom_all_reduce
  613. self.tokenizer_pool_config = tokenizer_pool_config
  614. self.ray_workers_use_nsight = ray_workers_use_nsight
  615. self.placement_group = placement_group
  616. self.world_size = pipeline_parallel_size * self.tensor_parallel_size
  617. if worker_use_ray:
  618. if self.distributed_executor_backend is None:
  619. self.distributed_executor_backend = "ray"
  620. elif self.distributed_executor_backend != "ray":
  621. raise ValueError(f"worker-use-ray can't be used with "
  622. f"distributed executor backend "
  623. f"'{self.distributed_executor_backend}'.")
  624. if self.distributed_executor_backend is None and self.world_size > 1:
  625. # We use multiprocessing by default if world_size fits on the
  626. # current node and we aren't in a ray placement group.
  627. from aphrodite.executor import ray_utils
  628. backend = "mp"
  629. ray_found = ray_utils.ray is not None
  630. if cuda_device_count_stateless() < self.world_size:
  631. if not ray_found:
  632. raise ValueError("Unable to load Ray which is "
  633. "required for multi-node inference")
  634. backend = "ray"
  635. elif ray_found:
  636. if self.placement_group:
  637. backend = "ray"
  638. else:
  639. from ray import is_initialized as ray_is_initialized
  640. if ray_is_initialized():
  641. from ray.util import get_current_placement_group
  642. if get_current_placement_group():
  643. backend = "ray"
  644. self.distributed_executor_backend = backend
  645. logger.info(
  646. f"Defaulting to use {backend} for distributed inference.")
  647. # If CUDA_VISIBLE_DEVICES is set on ROCm prior to Aphrodite init,
  648. # propagate changes to HIP_VISIBLE_DEVICES (conversion handled by
  649. # the update_environment_variables function)
  650. if is_hip() and os.getenv("CUDA_VISIBLE_DEVICES"):
  651. update_environment_variables({
  652. "CUDA_VISIBLE_DEVICES":
  653. os.getenv("CUDA_VISIBLE_DEVICES", "")
  654. })
  655. self._verify_args()
  656. def _verify_args(self) -> None:
  657. if self.pipeline_parallel_size > 1:
  658. raise NotImplementedError(
  659. "Pipeline parallelism is not supported yet.")
  660. if self.distributed_executor_backend not in ("ray", "mp", None):
  661. raise ValueError(
  662. "Unrecognized distributed executor backend. Supported values "
  663. "are 'ray' or 'mp'.")
  664. if not self.disable_custom_all_reduce and self.world_size > 1:
  665. if is_hip():
  666. self.disable_custom_all_reduce = True
  667. logger.info(
  668. "Disabled the custom all-reduce kernel because it is not "
  669. "supported on AMD GPUs.")
  670. elif self.pipeline_parallel_size > 1:
  671. self.disable_custom_all_reduce = True
  672. logger.info(
  673. "Disabled the custom all-reduce kernel because it is not "
  674. "supported with pipeline parallelism.")
  675. if self.ray_workers_use_nsight and (
  676. not self.distributed_executor_backend == "ray"):
  677. raise ValueError("Unable to use nsight profiling unless workers "
  678. "run with Ray.")
  679. class SchedulerConfig:
  680. """Scheduler configuration.
  681. Args:
  682. max_num_batched_tokens: Maximum number of tokens to be processed in
  683. a single iteration.
  684. max_num_seqs: Maximum number of sequences to be processed in a single
  685. iteration.
  686. max_model_len: Maximum length of a sequence (including prompt
  687. and generated text).
  688. use_v2_block_manager: Whether to use the BlockSpaceManagerV2 or not.
  689. num_lookahead_slots: The number of slots to allocate per sequence per
  690. step, beyond the known token ids. This is used in speculative
  691. decoding to store KV activations of tokens which may or may not be
  692. accepted.
  693. delay_factor: Apply a delay (of delay factor multiplied by previous
  694. prompt latency) before scheduling next prompt.
  695. enable_chunked_prefill: If True, prefill requests can be chunked based
  696. on the remaining max_num_batched_tokens.
  697. embedding_mode: Whether the running model is for embedding.
  698. preemption_mode: Whether to perform preemption by swapping or
  699. recomputation. If not specified, we determine the mode as follows:
  700. We use recomputation by default since it incurs lower overhead than
  701. swapping. However, when the sequence group has multiple sequences
  702. (e.g., beam search), recomputation is not currently supported. In
  703. such a case, we use swapping instead.
  704. """
  705. def __init__(self,
  706. max_num_batched_tokens: Optional[int],
  707. max_num_seqs: int,
  708. max_model_len: int,
  709. use_v2_block_manager: bool = False,
  710. num_lookahead_slots: int = 0,
  711. delay_factor: float = 0.0,
  712. enable_chunked_prefill: bool = False,
  713. embedding_mode: Optional[bool] = False,
  714. preemption_mode: Optional[str] = None) -> None:
  715. if max_num_batched_tokens is not None:
  716. self.max_num_batched_tokens = max_num_batched_tokens
  717. else:
  718. if enable_chunked_prefill:
  719. # For chunked prefill, choose the well-tuned batch size.
  720. self.max_num_batched_tokens = 768
  721. elif embedding_mode:
  722. # For embedding, choose specific value for higher throughput
  723. self.max_num_batched_tokens = max(
  724. max_model_len, _EMBEDDING_MODEL_MAX_NUM_BATCHED_TOKENS)
  725. else:
  726. # If max_model_len is too short, use 2048 as the default value
  727. # for higher throughput.
  728. self.max_num_batched_tokens = max(max_model_len, 2048)
  729. if enable_chunked_prefill:
  730. logger.info("Chunked prefill is enabled (EXPERIMENTAL).")
  731. self.max_num_seqs = max_num_seqs
  732. self.max_model_len = max_model_len
  733. self.use_v2_block_manager = use_v2_block_manager
  734. self.num_lookahead_slots = num_lookahead_slots
  735. self.delay_factor = delay_factor
  736. self.chunked_prefill_enabled = enable_chunked_prefill
  737. self.embedding_mode = embedding_mode
  738. self.preemption_mode = preemption_mode
  739. self._verify_args()
  740. def _verify_args(self) -> None:
  741. if (self.max_num_batched_tokens < self.max_model_len
  742. and not self.chunked_prefill_enabled):
  743. raise ValueError(
  744. f"max_num_batched_tokens ({self.max_num_batched_tokens}) is "
  745. f"smaller than max_model_len ({self.max_model_len}). "
  746. "This effectively limits the maximum sequence length to "
  747. "max_num_batched_tokens and makes Aphrodite reject longer "
  748. "sequences. Please increase max_num_batched_tokens or "
  749. "decrease max_model_len.")
  750. if self.max_num_batched_tokens < self.max_num_seqs:
  751. raise ValueError(
  752. f"max_num_batched_tokens ({self.max_num_batched_tokens}) must "
  753. "be greater than or equal to max_num_seqs "
  754. f"({self.max_num_seqs}).")
  755. if self.num_lookahead_slots < 0:
  756. raise ValueError(
  757. "num_lookahead_slots "
  758. f"({self.num_lookahead_slots}) must be greater than or "
  759. "equal to 0.")
  760. class DeviceConfig:
  761. def __init__(self, device: str = "auto") -> None:
  762. if device == "auto":
  763. # Automated device type detection
  764. if is_neuron():
  765. self.device_type = "neuron"
  766. elif is_tpu():
  767. self.device_type = "tpu"
  768. elif is_cpu():
  769. self.device_type = "cpu"
  770. elif is_xpu():
  771. self.device_type = "xpu"
  772. else:
  773. # We don't call torch.cuda.is_available() here to
  774. # avoid initializing CUDA before workers are forked
  775. self.device_type = "cuda"
  776. else:
  777. # Device type is assigned explicitly
  778. self.device_type = device
  779. # Some device types require processing inputs on CPU
  780. if self.device_type in ["neuron"]:
  781. self.device = torch.device("cpu")
  782. elif self.device_type in ["tpu"]:
  783. self.device = None
  784. else:
  785. # Set device with device type
  786. self.device = torch.device(self.device_type)
  787. class SpeculativeConfig:
  788. """Configuration for speculative decoding.
  789. The configuration is currently specialized to draft-model speculative
  790. decoding with top-1 proposals.
  791. """
  792. @staticmethod
  793. def maybe_create_spec_config(
  794. target_model_config: ModelConfig,
  795. target_parallel_config: ParallelConfig,
  796. target_dtype: str,
  797. speculative_model: Optional[str],
  798. speculative_draft_tensor_parallel_size: Optional[int],
  799. num_speculative_tokens: Optional[int],
  800. speculative_max_model_len: Optional[int],
  801. enable_chunked_prefill: bool,
  802. use_v2_block_manager: bool,
  803. speculative_disable_by_batch_size: Optional[int],
  804. ngram_prompt_lookup_max: Optional[int],
  805. ngram_prompt_lookup_min: Optional[int],
  806. ) -> Optional["SpeculativeConfig"]:
  807. """Create a SpeculativeConfig if possible, else return None.
  808. This function attempts to create a SpeculativeConfig object based on the
  809. provided parameters. If the necessary conditions are met, it returns an
  810. instance of SpeculativeConfig. Otherwise, it returns None.
  811. Args:
  812. target_model_config (ModelConfig): The configuration of the target
  813. model.
  814. target_parallel_config (ParallelConfig): The parallel configuration
  815. for the target model.
  816. target_dtype (str): The data type used for the target model.
  817. speculative_model (Optional[str]): The name of the speculative
  818. model, if provided.
  819. num_speculative_tokens (Optional[int]): The number of speculative
  820. tokens, if provided. Will default to the number in the draft
  821. model config if present, otherwise is required.
  822. speculative_draft_tensor_parallel_size (Optional[int]): The degree
  823. of the tensor parallelism for the draft model.
  824. speculative_max_model_len (Optional[int]): The maximum model len of
  825. the speculative model. Used when testing the ability to skip
  826. speculation for some sequences.
  827. enable_chunked_prefill (bool): Whether Aphrodite is configured to
  828. use chunked prefill or not. Used for raising an error since its
  829. not yet compatible with spec decode.
  830. use_v2_block_manager (bool): Whether Aphrodite is configured to
  831. use the v2 block manager or not. Used for raising an error
  832. since the v2 block manager is required with spec decode.
  833. speculative_disable_by_batch_size (Optional[int]): Disable
  834. speculative decoding for new incoming requests when the number
  835. of enqueue requests is larger than this value, if provided.
  836. ngram_prompt_lookup_max (Optional[int]): Max size of ngram token
  837. window, if provided.
  838. ngram_prompt_lookup_min (Optional[int]): Min size of ngram token
  839. window, if provided.
  840. Returns:
  841. Optional["SpeculativeConfig"]: An instance of SpeculativeConfig if
  842. the necessary conditions are met, else None.
  843. """
  844. if speculative_model is None:
  845. if num_speculative_tokens is not None:
  846. raise ValueError("num_speculative_tokens was provided without "
  847. "speculative_model.")
  848. return None
  849. if (speculative_disable_by_batch_size is not None
  850. and speculative_disable_by_batch_size < 2):
  851. raise ValueError("Expected the batch size threshold of disabling "
  852. "speculative decoding is > 1, but got "
  853. f"{speculative_disable_by_batch_size=}")
  854. if enable_chunked_prefill:
  855. raise ValueError(
  856. "Speculative decoding and chunked prefill are "
  857. f"currently mutually exclusive ({enable_chunked_prefill=}).")
  858. if not use_v2_block_manager:
  859. raise ValueError(
  860. "Speculative decoding requires usage of the V2 "
  861. "block manager. Enable it with --use-v2-block-manager.")
  862. # TODO: The user should be able to specify revision/quantization/max
  863. # model len for the draft model. It is not currently supported.
  864. draft_revision = None
  865. draft_code_revision = None
  866. draft_quantization = None
  867. if speculative_model == "[ngram]":
  868. if ngram_prompt_lookup_min is None:
  869. ngram_prompt_lookup_min = 1
  870. if ngram_prompt_lookup_max is None or ngram_prompt_lookup_max < 1:
  871. raise ValueError(f"{ngram_prompt_lookup_max=} must be > 0")
  872. if ngram_prompt_lookup_min < 1:
  873. raise ValueError(f"{ngram_prompt_lookup_min=} must be > 0")
  874. if ngram_prompt_lookup_min > ngram_prompt_lookup_max:
  875. raise ValueError(f"{ngram_prompt_lookup_min=} cannot be "
  876. f"larger than {ngram_prompt_lookup_max=}")
  877. # TODO: current we still need extract vocab_size from target model
  878. # config, in future, we may try refactoring it out, and set
  879. # draft related config as None here.
  880. draft_model_config = target_model_config
  881. draft_parallel_config = target_parallel_config
  882. else:
  883. ngram_prompt_lookup_max = 0
  884. ngram_prompt_lookup_min = 0
  885. draft_model_config = ModelConfig(
  886. model=speculative_model,
  887. tokenizer=target_model_config.tokenizer,
  888. tokenizer_mode=target_model_config.tokenizer_mode,
  889. trust_remote_code=target_model_config.trust_remote_code,
  890. dtype=target_model_config.dtype,
  891. seed=target_model_config.seed,
  892. revision=draft_revision,
  893. code_revision=draft_code_revision,
  894. tokenizer_revision=target_model_config.tokenizer_revision,
  895. max_model_len=None,
  896. quantization=draft_quantization,
  897. enforce_eager=target_model_config.enforce_eager,
  898. max_seq_len_to_capture=target_model_config.
  899. max_seq_len_to_capture,
  900. max_logprobs=target_model_config.max_logprobs,
  901. )
  902. draft_hf_config = draft_model_config.hf_config
  903. if (draft_hf_config.model_type == "mlp_speculator"
  904. and target_parallel_config.world_size != 1):
  905. # MLPSpeculator TP support will be added very soon
  906. raise ValueError(
  907. "Speculative decoding with mlp_speculator models does not "
  908. "yet support distributed inferencing (TP > 1).")
  909. if (num_speculative_tokens is not None
  910. and hasattr(draft_hf_config, "num_lookahead_tokens")):
  911. draft_hf_config.num_lookahead_tokens = num_speculative_tokens
  912. n_predict = getattr(draft_hf_config, "n_predict", None)
  913. if n_predict is not None:
  914. if num_speculative_tokens is None:
  915. # Default to max value defined in draft model config.
  916. num_speculative_tokens = n_predict
  917. elif num_speculative_tokens > n_predict:
  918. # Verify provided value doesn't exceed the maximum
  919. # supported by the draft model.
  920. raise ValueError(
  921. "Expected both speculative_model and "
  922. "num_speculative_tokens to be provided, but found "
  923. f"{speculative_model=} and {num_speculative_tokens=}.")
  924. draft_model_config.max_model_len = (
  925. SpeculativeConfig._maybe_override_draft_max_model_len(
  926. speculative_max_model_len,
  927. draft_model_config.max_model_len,
  928. target_model_config.max_model_len,
  929. ))
  930. draft_parallel_config = (
  931. SpeculativeConfig.create_draft_parallel_config(
  932. target_parallel_config,
  933. speculative_draft_tensor_parallel_size))
  934. if num_speculative_tokens is None:
  935. raise ValueError(
  936. "num_speculative_tokens must be provided with "
  937. "speculative_model unless the draft model config contains an "
  938. "n_predict parameter.")
  939. return SpeculativeConfig(draft_model_config, draft_parallel_config,
  940. num_speculative_tokens,
  941. speculative_disable_by_batch_size,
  942. ngram_prompt_lookup_max,
  943. ngram_prompt_lookup_min)
  944. @staticmethod
  945. def _maybe_override_draft_max_model_len(
  946. speculative_max_model_len: Optional[int],
  947. draft_max_model_len: int,
  948. target_max_model_len: int,
  949. ) -> int:
  950. """Determine the max sequence len for the draft model. This is usually
  951. the draft_max_model_len, but may be the target_max_model_len if it is
  952. less than the draft_max_model_len, or may be speculative_max_model_len
  953. if it is specified.
  954. This is necessary so that sequences do not exceed the capacity of the
  955. draft model or the target model.
  956. speculative_max_model_len is mainly used for testing that sequences can
  957. skip speculation.
  958. """
  959. if speculative_max_model_len is not None:
  960. if speculative_max_model_len > draft_max_model_len:
  961. raise ValueError(f"{speculative_max_model_len=} cannot be "
  962. f"larger than {draft_max_model_len=}")
  963. if speculative_max_model_len > target_max_model_len:
  964. raise ValueError(f"{speculative_max_model_len=} cannot be "
  965. f"larger than {target_max_model_len=}")
  966. return speculative_max_model_len
  967. return min(
  968. draft_max_model_len,
  969. target_max_model_len,
  970. )
  971. @staticmethod
  972. def create_draft_parallel_config(
  973. target_parallel_config: ParallelConfig,
  974. speculative_draft_tensor_parallel_size: Optional[int]
  975. ) -> ParallelConfig:
  976. """Create a parallel config for use by the draft worker.
  977. This is mostly a copy of the target parallel config, except the tp_size.
  978. """
  979. if speculative_draft_tensor_parallel_size is None:
  980. speculative_draft_tensor_parallel_size = \
  981. target_parallel_config.tensor_parallel_size
  982. elif speculative_draft_tensor_parallel_size != 1:
  983. # TODO: allow tp values larger than 1
  984. raise ValueError(
  985. f"{speculative_draft_tensor_parallel_size=} cannot be"
  986. f"other value than 1")
  987. draft_parallel_config = ParallelConfig(
  988. pipeline_parallel_size=target_parallel_config.
  989. pipeline_parallel_size,
  990. tensor_parallel_size=speculative_draft_tensor_parallel_size,
  991. distributed_executor_backend=target_parallel_config.
  992. distributed_executor_backend,
  993. max_parallel_loading_workers=target_parallel_config.
  994. max_parallel_loading_workers,
  995. disable_custom_all_reduce=target_parallel_config.
  996. disable_custom_all_reduce,
  997. tokenizer_pool_config=target_parallel_config.tokenizer_pool_config,
  998. ray_workers_use_nsight=target_parallel_config.
  999. ray_workers_use_nsight,
  1000. placement_group=target_parallel_config.placement_group,
  1001. )
  1002. return draft_parallel_config
  1003. def __init__(
  1004. self,
  1005. draft_model_config: ModelConfig,
  1006. draft_parallel_config: ParallelConfig,
  1007. num_speculative_tokens: int,
  1008. speculative_disable_by_batch_size: Optional[int],
  1009. ngram_prompt_lookup_max: Optional[int],
  1010. ngram_prompt_lookup_min: Optional[int],
  1011. ):
  1012. """Create a SpeculativeConfig object.
  1013. Args:
  1014. draft_model_config: ModelConfig for the draft model.
  1015. draft_parallel_config: ParallelConfig for the draft model.
  1016. num_speculative_tokens: The number of tokens to sample from the
  1017. draft model before scoring with the target model.
  1018. speculative_disable_by_batch_size: Disable speculative
  1019. decoding for new incoming requests when the number of
  1020. enqueue requests is larger than this value.
  1021. ngram_prompt_lookup_max: Max size of ngram token window.
  1022. ngram_prompt_lookup_min: Min size of ngram token window.
  1023. """
  1024. self.draft_model_config = draft_model_config
  1025. self.draft_parallel_config = draft_parallel_config
  1026. self.num_speculative_tokens = num_speculative_tokens
  1027. self.speculative_disable_by_batch_size = \
  1028. speculative_disable_by_batch_size
  1029. self.ngram_prompt_lookup_max = ngram_prompt_lookup_max or 0
  1030. self.ngram_prompt_lookup_min = ngram_prompt_lookup_min or 0
  1031. self._verify_args()
  1032. def _verify_args(self) -> None:
  1033. if self.num_speculative_tokens <= 0:
  1034. raise ValueError("Expected num_speculative_tokens to be greater "
  1035. f"than zero ({self.num_speculative_tokens}).")
  1036. if self.draft_model_config:
  1037. self.draft_model_config.verify_with_parallel_config(
  1038. self.draft_parallel_config)
  1039. @property
  1040. def num_lookahead_slots(self) -> int:
  1041. """The number of additional slots the scheduler should allocate per
  1042. step, in addition to the slots allocated for each known token.
  1043. This is equal to the number of speculative tokens, as each speculative
  1044. token must be scored.
  1045. """
  1046. return self.num_speculative_tokens
  1047. def __repr__(self) -> str:
  1048. if self.ngram_prompt_lookup_max > 0:
  1049. draft_model = "[ngram]"
  1050. else:
  1051. draft_model = self.draft_model_config.model
  1052. num_spec_tokens = self.num_speculative_tokens
  1053. return f"SpeculativeConfig({draft_model=}, {num_spec_tokens=})"
  1054. @dataclass
  1055. class LoRAConfig:
  1056. max_lora_rank: int
  1057. max_loras: int
  1058. fully_sharded_loras: bool = False
  1059. max_cpu_loras: Optional[int] = None
  1060. lora_dtype: Optional[torch.dtype] = None
  1061. lora_extra_vocab_size: int = 256
  1062. # This is a constant.
  1063. lora_vocab_padding_size: ClassVar[int] = 256
  1064. long_lora_scaling_factors: Optional[Tuple[float]] = None
  1065. def __post_init__(self):
  1066. # Keep this in sync with kernels/punica/bgmv/bgmv_config.h
  1067. possible_max_ranks = (8, 16, 32, 64)
  1068. possible_lora_extra_vocab_size = (0, 256, 512)
  1069. if self.max_lora_rank not in possible_max_ranks:
  1070. raise ValueError(
  1071. f"max_lora_rank ({self.max_lora_rank}) must be one of "
  1072. f"{possible_max_ranks}.")
  1073. if self.lora_extra_vocab_size not in possible_lora_extra_vocab_size:
  1074. raise ValueError(
  1075. f"lora_extra_vocab_size ({self.lora_extra_vocab_size}) "
  1076. f"must be one of {possible_lora_extra_vocab_size}.")
  1077. if self.max_loras < 1:
  1078. raise ValueError(f"max_loras ({self.max_loras}) must be >= 1.")
  1079. if self.max_cpu_loras is None:
  1080. self.max_cpu_loras = self.max_loras
  1081. elif self.max_cpu_loras < self.max_loras:
  1082. raise ValueError(
  1083. f"max_cpu_loras ({self.max_cpu_loras}) must be >= "
  1084. f"max_loras ({self.max_loras})")
  1085. def verify_with_model_config(self, model_config: ModelConfig):
  1086. if self.lora_dtype in (None, "auto"):
  1087. self.lora_dtype = model_config.dtype
  1088. elif isinstance(self.lora_dtype, str):
  1089. self.lora_dtype = getattr(torch, self.lora_dtype)
  1090. if model_config.quantization and model_config.quantization not in [
  1091. "awq", "gptq"
  1092. ]:
  1093. # TODO support all other quants
  1094. logger.warning(f"{model_config.quantization} quantization is not "
  1095. "tested with LoRA yet.")
  1096. def verify_with_scheduler_config(self, scheduler_config: SchedulerConfig):
  1097. if scheduler_config.max_num_batched_tokens > 65528:
  1098. raise ValueError(
  1099. "Due to limitations of the custom LoRA CUDA kernel, "
  1100. "max_num_batched_tokens must be <= 65528 when "
  1101. "LoRA is enabled.")
  1102. if scheduler_config.chunked_prefill_enabled:
  1103. raise ValueError("LoRA is not supported with chunked prefill yet.")
  1104. @dataclass
  1105. class VisionLanguageConfig:
  1106. """Configs the input data format and how models should run for
  1107. vision language models."""
  1108. class ImageInputType(enum.Enum):
  1109. """Image input type into the vision language model.
  1110. An image roughly goes through the following transformation:
  1111. Raw image --> pixel values --> image features --> image embeddings.
  1112. The difference between different image input types is where the
  1113. image encoder (pixel values --> image features) is run.
  1114. Different image input types also correspond to different tensor shapes.
  1115. For example, for Llava, PIXEL_VALUES: (1, 3, 336, 336).
  1116. IMAGE_FEATURES: (1, 576, 1024).
  1117. """
  1118. PIXEL_VALUES = enum.auto()
  1119. IMAGE_FEATURES = enum.auto()
  1120. image_input_type: ImageInputType
  1121. # The input id corresponding to image token.
  1122. image_token_id: int
  1123. # Used for running `run_prefill_max_token`.
  1124. # For models that support varying resolution, this corresponds to
  1125. # worst case scenario (biggest supported resolution).
  1126. image_input_shape: tuple
  1127. image_feature_size: int
  1128. # The image processor to load from HuggingFace
  1129. image_processor: Optional[str]
  1130. image_processor_revision: Optional[str]
  1131. @classmethod
  1132. def get_image_input_enum_type(
  1133. cls, value: str) -> "VisionLanguageConfig.ImageInputType":
  1134. """Get the image input type from a string."""
  1135. try:
  1136. return cls.ImageInputType[value.upper()]
  1137. except KeyError as e:
  1138. raise ValueError(f"{value} is not a valid choice. "
  1139. f"Expecting to choose from "
  1140. f"{[x.name for x in cls.ImageInputType]}.") from e
  1141. def as_cli_args_dict(self) -> Dict[str, Any]:
  1142. """Flatten vision language config to pure args.
  1143. Compatible with what llm entrypoint expects.
  1144. """
  1145. result: Dict[str, Any] = {}
  1146. for f in fields(self):
  1147. value = getattr(self, f.name)
  1148. if isinstance(value, enum.Enum):
  1149. result[f.name] = value.name.lower()
  1150. elif isinstance(value, tuple):
  1151. result[f.name] = ",".join([str(item) for item in value])
  1152. else:
  1153. result[f.name] = value
  1154. result["disable_image_processor"] = self.image_processor is None
  1155. return result
  1156. _STR_DTYPE_TO_TORCH_DTYPE = {
  1157. "half": torch.float16,
  1158. "float16": torch.float16,
  1159. "float": torch.float32,
  1160. "float32": torch.float32,
  1161. "bfloat16": torch.bfloat16,
  1162. }
  1163. _ROCM_NOT_SUPPORTED_DTYPE = ["float", "float32"]
  1164. def _get_and_verify_dtype(
  1165. config: PretrainedConfig,
  1166. dtype: Union[str, torch.dtype],
  1167. ) -> torch.dtype:
  1168. # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
  1169. # because config.torch_dtype can be None.
  1170. config_dtype = getattr(config, "torch_dtype", None)
  1171. if config_dtype is None:
  1172. config_dtype = torch.float32
  1173. if isinstance(dtype, str):
  1174. dtype = dtype.lower()
  1175. if dtype == "auto":
  1176. if config_dtype == torch.float32:
  1177. # Following the common practice, we use float16 for float32
  1178. # models.
  1179. torch_dtype = torch.float16
  1180. else:
  1181. torch_dtype = config_dtype
  1182. else:
  1183. if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:
  1184. raise ValueError(f"Unknown dtype: {dtype}")
  1185. torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
  1186. elif isinstance(dtype, torch.dtype):
  1187. torch_dtype = dtype
  1188. else:
  1189. raise ValueError(f"Unknown dtype: {dtype}")
  1190. if is_hip() and torch_dtype == torch.float32:
  1191. rocm_supported_dtypes = [
  1192. k for k, v in _STR_DTYPE_TO_TORCH_DTYPE.items()
  1193. if (k not in _ROCM_NOT_SUPPORTED_DTYPE)
  1194. ]
  1195. raise ValueError(f"dtype '{dtype}' is not supported in ROCm. "
  1196. f"Supported dtypes are {rocm_supported_dtypes}")
  1197. # Verify the dtype.
  1198. if torch_dtype != config_dtype:
  1199. if torch_dtype == torch.float32:
  1200. # Upcasting to float32 is allowed.
  1201. pass
  1202. elif config_dtype == torch.float32:
  1203. # Downcasting from float32 to float16 or bfloat16 is allowed.
  1204. pass
  1205. else:
  1206. # Casting between float16 and bfloat16 is allowed with a warning.
  1207. logger.warning(f"Casting {config_dtype} to {torch_dtype}.")
  1208. return torch_dtype
  1209. def _get_and_verify_max_len(
  1210. hf_config: PretrainedConfig,
  1211. max_model_len: Optional[int],
  1212. disable_sliding_window: bool,
  1213. sliding_window_len: Optional[int],
  1214. rope_scaling_arg: Optional[Dict[str, Any]],
  1215. ) -> int:
  1216. """Get and verify the model's maximum length."""
  1217. derived_max_model_len = float("inf")
  1218. possible_keys = [
  1219. # Cohere: needs to prioritize this over "max_position_embeddings"
  1220. "model_max_length",
  1221. # OPT
  1222. "max_position_embeddings",
  1223. # GPT-2
  1224. "n_positions",
  1225. # MPT
  1226. "max_seq_len",
  1227. # ChatGLM2
  1228. "seq_length",
  1229. # Command-R
  1230. "model_max_length",
  1231. # Others
  1232. "max_sequence_length",
  1233. "max_seq_length",
  1234. "seq_len",
  1235. ]
  1236. # Choose the smallest "max_length" from the possible keys.
  1237. max_len_key = None
  1238. for key in possible_keys:
  1239. max_len = getattr(hf_config, key, None)
  1240. if max_len is not None:
  1241. max_len_key = key if max_len < derived_max_model_len \
  1242. else max_len_key
  1243. derived_max_model_len = min(derived_max_model_len, max_len)
  1244. # If sliding window is manually disabled, max_length should be less
  1245. # than the sliding window length in the model config.
  1246. if disable_sliding_window and sliding_window_len is not None:
  1247. max_len_key = "sliding_window" \
  1248. if sliding_window_len < derived_max_model_len else max_len_key
  1249. derived_max_model_len = min(derived_max_model_len, sliding_window_len)
  1250. # If none of the keys were found in the config, use a default and
  1251. # log a warning.
  1252. if derived_max_model_len == float("inf"):
  1253. if max_model_len is not None:
  1254. # If max_model_len is specified, we use it.
  1255. return max_model_len
  1256. default_max_len = 2048
  1257. logger.warning(
  1258. "The model's config.json does not contain any of the following "
  1259. "keys to determine the original maximum length of the model: "
  1260. f"{possible_keys}. Assuming the model's maximum length is "
  1261. f"{default_max_len}.")
  1262. derived_max_model_len = default_max_len
  1263. rope_scaling = getattr(hf_config, "rope_scaling", None)
  1264. if rope_scaling is not None:
  1265. rope_type = rope_scaling.get("type", rope_scaling.get("rope_type"))
  1266. if rope_type not in {"su", "longrope", "llama3"}:
  1267. if disable_sliding_window:
  1268. # TODO: Find a model that supports rope_scaling
  1269. # with sliding window to see if this case should be allowed.
  1270. raise NotImplementedError(
  1271. "Disabling sliding window is not supported for models "
  1272. "with rope_scaling. Please raise an issue so we can "
  1273. "investigate.")
  1274. assert "factor" in rope_scaling
  1275. scaling_factor = rope_scaling["factor"]
  1276. if rope_type == "yarn":
  1277. derived_max_model_len = rope_scaling[
  1278. "original_max_position_embeddings"]
  1279. derived_max_model_len *= scaling_factor
  1280. if max_model_len is None:
  1281. max_model_len = derived_max_model_len
  1282. elif max_model_len > derived_max_model_len and rope_scaling_arg is None:
  1283. raise ValueError(
  1284. f"User-specified max_model_len {max_model_len} is higher than "
  1285. f"the original {derived_max_model_len}. "
  1286. "Please provide a rope_scaling dict to scale the model.")
  1287. elif max_model_len > derived_max_model_len and rope_scaling_arg is not None:
  1288. # hope this works
  1289. logger.warning(
  1290. f"User-specified max_model_len {max_model_len} is higher than "
  1291. f"the original {derived_max_model_len}. "
  1292. "Attempting to use RoPE scaling with the provided rope_scaling "
  1293. "dict.")
  1294. derived_max_model_len = max_model_len
  1295. return int(max_model_len)
  1296. @dataclass
  1297. class DecodingConfig:
  1298. """Dataclass which contains the decoding strategy of the engine"""
  1299. # Which guided decoding algo to use. 'outlines' / 'lm-format-enforcer'
  1300. guided_decoding_backend: str = 'outlines'
  1301. def __post_init__(self):
  1302. valid_guided_backends = ['outlines', 'lm-format-enforcer']
  1303. backend = self.guided_decoding_backend
  1304. if backend not in valid_guided_backends:
  1305. raise ValueError(f"Invalid guided_decoding_backend '{backend},"
  1306. f"must be one of {valid_guided_backends}")
  1307. @dataclass(frozen=True)
  1308. class EngineConfig:
  1309. """Dataclass which contains all engine-related configuration. This
  1310. simplifies passing around the distinct configurations in the codebase.
  1311. """
  1312. model_config: ModelConfig
  1313. cache_config: CacheConfig
  1314. parallel_config: ParallelConfig
  1315. scheduler_config: SchedulerConfig
  1316. device_config: DeviceConfig
  1317. load_config: LoadConfig
  1318. lora_config: Optional[LoRAConfig]
  1319. vision_language_config: Optional[VisionLanguageConfig]
  1320. speculative_config: Optional[SpeculativeConfig]
  1321. decoding_config: Optional[DecodingConfig]
  1322. def __post_init__(self):
  1323. """Verify configs are valid & consistent with each other.
  1324. """
  1325. self.model_config.verify_with_parallel_config(self.parallel_config)
  1326. self.cache_config.verify_with_parallel_config(self.parallel_config)
  1327. if self.lora_config:
  1328. self.lora_config.verify_with_model_config(self.model_config)
  1329. self.lora_config.verify_with_scheduler_config(
  1330. self.scheduler_config)
  1331. def to_dict(self):
  1332. """Return the configs as a dictionary, for use in **kwargs.
  1333. """
  1334. return dict(
  1335. (field.name, getattr(self, field.name)) for field in fields(self))