diff --git a/truss/base/constants.py b/truss/base/constants.py index 8708df1dc..ebe52f219 100644 --- a/truss/base/constants.py +++ b/truss/base/constants.py @@ -27,6 +27,7 @@ BEI_TRTLLM_CLIENT_BATCH_SIZE = 128 BEI_MAX_CONCURRENCY_TARGET_REQUESTS = 2048 BEI_REQUIRED_MAX_NUM_TOKENS = 16384 +BEI_TORCH_REQUIRED_MAX_NUM_TOKENS = 32768 TRTLLM_MIN_MEMORY_REQUEST_GI = 10 HF_MODELS_API_URL = "https://huggingface.co/api/models" diff --git a/truss/base/trt_llm_config.py b/truss/base/trt_llm_config.py index 200c891ed..a8d80a9e0 100644 --- a/truss/base/trt_llm_config.py +++ b/truss/base/trt_llm_config.py @@ -50,6 +50,11 @@ class TrussTRTLLMModel(str, Enum): # the encoder_bert setting will specfically optimize for thoughput and cold-start latency of small models (<4B parameters) # supports also splade and colbert style models or ModernBert. ENCODER_BERT = "encoder_bert" + # BEI torch backend (text-embeddings-router + vLLM). Serves embedding, reranker, and + # classification models whose architecture the TRT-LLM `encoder` path cannot compile + # (e.g. LlamaBidirectional / Ministral3-Embed, Gemma2Embedding, jina-v3). No engine + # build step; the checkpoint is served directly by the torch backend. + ENCODER_TORCH = "encoder_torch" # Decoder will launch the backend that is optimized for decoder only models such as LLama3ForCausalLM, Qwen3MoeForCausalLM etc. DECODER = "decoder" # a ERROR will be raised if you push one of the below models. Don't use @@ -365,7 +370,31 @@ def uses_lora(self) -> bool: def _bei_specfic_migration(self): """performs embedding specfic optimizations (no kv-cache, high batch size)""" - if self.base_model == TrussTRTLLMModel.ENCODER: + if self.base_model == TrussTRTLLMModel.ENCODER_TORCH: + # BEI-torch has no engine build, so build-time quantization is not applicable. + # Users wanting quantized weights must point checkpoint_repository at a + # pre-quantized HF checkpoint. + if self.quantization_type != TrussTRTLLMQuantizationType.NO_QUANT: + raise ValueError( + f"base_model=encoder_torch does not support build-time quantization; " + f"you set build.quantization_type={self.quantization_type.value}. " + "Point checkpoint_repository at a pre-quantized HF checkpoint instead." + ) + # Speculative decoding is a decoder concept; embedding/reranker serving via + # BEI-torch has no draft/target loop. + if self.speculator is not None: + raise ValueError( + "base_model=encoder_torch does not support speculative decoding; " + "remove trt_llm.build.speculator." + ) + # LoRA adapters are a TRT-LLM build-time concept; vLLM serves the + # checkpoint as-is, so adapters would be silently ignored. + if self.lora_adapters is not None: + raise ValueError( + "base_model=encoder_torch does not support lora_adapters; " + "remove trt_llm.build.lora_adapters." + ) + elif self.base_model == TrussTRTLLMModel.ENCODER: # Encoder specific settings if self.max_seq_len: logger.info( @@ -572,6 +601,7 @@ class VersionsOverrides(PydanticTrTBaseModel): briton_version: Optional[str] = None bei_version: Optional[str] = None bei_bert_version: Optional[str] = None + bei_torch_version: Optional[str] = None v2_llm_version: Optional[str] = None @model_validator(mode="before") @@ -581,6 +611,7 @@ def version_must_start_with_number(cls, data): "briton_version", "bei_version", "bei_bert_version", + "bei_torch_version", ]: v = data.get(field) if v is not None and (not v or not v[0].isdigit()): @@ -596,6 +627,10 @@ class ImageVersions(PydanticTrTBaseModel): # INTERNAL bei_image: str beibert_image: str + # Base image for the BEI torch backend (text-embeddings-router + vLLM). Backend + # inserts a resolved image reference; the default here is a placeholder so + # local unit tests can construct ImageVersions without wiring core-product. + bei_torch_image: str = "baseten/bei:torch-dev-placeholder" briton_image: str v2_llm_image: str @@ -618,7 +653,11 @@ def model_post_init(self, __context): self.runtime.enable_chunked_context and ( self.build.base_model - not in (TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT) + not in ( + TrussTRTLLMModel.ENCODER, + TrussTRTLLMModel.ENCODER_BERT, + TrussTRTLLMModel.ENCODER_TORCH, + ) ) and not ( self.build.plugin_configuration.use_paged_context_fmha @@ -636,7 +675,11 @@ def model_post_init(self, __context): if ( self.runtime.webserver_default_route is None and self.build.base_model - in (TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT) + in ( + TrussTRTLLMModel.ENCODER, + TrussTRTLLMModel.ENCODER_BERT, + TrussTRTLLMModel.ENCODER_TORCH, + ) and not ENGINE_BUILDER_TRUSS_RUNTIME_MIGRATION ): if hf_cfg is not None: @@ -672,9 +715,9 @@ def model_post_init(self, __context): f"but you set `trt_llm.build.base_model` to `decoder`. " f"Please set it to `encoder_bert`." ) - if ( - "ForCausalLM" in arch - and self.build.base_model != TrussTRTLLMModel.DECODER + if "ForCausalLM" in arch and self.build.base_model not in ( + TrussTRTLLMModel.DECODER, + TrussTRTLLMModel.ENCODER_TORCH, ): logger.warning( f"Your model architecture {arch} indicates a CausalLM based model. " @@ -889,7 +932,7 @@ def trt_llm_common_validation(config: "TrussConfig"): pass else: raise ValueError( - "TRT-LLM is not supported on CUDA_COMPUTE_75 (T4) and CUDA_COMPUTE_70 (V100) GPUs. \n" + "TRT-LLM and BEI-torch are not supported on CUDA_COMPUTE_75 (T4) and CUDA_COMPUTE_70 (V100) GPUs. \n" "the lowest supported CUDA compute capability is CUDA_COMPUTE_80 (A100) or A10G (CUDA_COMPUTE_86)" ) elif trt_llm_config.build.quantization_type in [ @@ -947,6 +990,7 @@ def trt_llm_validation_v1(config: "TrussConfig") -> "TrussConfig": if trt_llm_config_v1.build.base_model not in [ TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT, + TrussTRTLLMModel.ENCODER_TORCH, ]: current_tags = config.model_metadata.get("tags", []) if ( diff --git a/truss/config.schema.json b/truss/config.schema.json index c9c15bf3f..812a8ed46 100644 --- a/truss/config.schema.json +++ b/truss/config.schema.json @@ -1097,6 +1097,7 @@ "briton_version": null, "bei_version": null, "bei_bert_version": null, + "bei_torch_version": null, "v2_llm_version": null } } @@ -1129,6 +1130,7 @@ "briton_version": null, "bei_version": null, "bei_bert_version": null, + "bei_torch_version": null, "v2_llm_version": null } } @@ -1569,6 +1571,7 @@ "enum": [ "encoder", "encoder_bert", + "encoder_torch", "decoder", "palmyra", "qwen", @@ -1775,6 +1778,18 @@ "default": null, "title": "Bei Bert Version" }, + "bei_torch_version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Bei Torch Version" + }, "v2_llm_version": { "anyOf": [ { diff --git a/truss/contexts/docker_build_setup.py b/truss/contexts/docker_build_setup.py index e98bceb6d..3751f515e 100644 --- a/truss/contexts/docker_build_setup.py +++ b/truss/contexts/docker_build_setup.py @@ -46,6 +46,12 @@ def _fill_trt_llm_versions( ): print(f"Using BEI image: {image_versions.bei_image}") tr.set_base_image(image_versions.bei_image, "/usr/bin/python3") + elif ( + tr.spec.config.trt_llm.build.base_model + == trt_llm_config.TrussTRTLLMModel.ENCODER_TORCH + ): + print(f"Using BEI torch image: {image_versions.bei_torch_image}") + tr.set_base_image(image_versions.bei_torch_image, "/usr/bin/python3") elif ( tr.spec.config.trt_llm.build.base_model == trt_llm_config.TrussTRTLLMModel.ENCODER_BERT diff --git a/truss/contexts/image_builder/serving_image_builder.py b/truss/contexts/image_builder/serving_image_builder.py index e081e342b..2903183ff 100644 --- a/truss/contexts/image_builder/serving_image_builder.py +++ b/truss/contexts/image_builder/serving_image_builder.py @@ -27,6 +27,7 @@ BASE_SERVER_REQUIREMENTS_TXT_FILENAME, BEI_MAX_CONCURRENCY_TARGET_REQUESTS, BEI_REQUIRED_MAX_NUM_TOKENS, + BEI_TORCH_REQUIRED_MAX_NUM_TOKENS, BEI_TRTLLM_CLIENT_BATCH_SIZE, CHAINS_CODE_DIR, CONSTRAINTS_TXT_FILENAME, @@ -610,6 +611,54 @@ def prepare_trtllm_bei_encoder_build_dir(self, build_dir: Path): ) copy_tree_path(DOCKER_SERVER_TEMPLATES_DIR, build_dir, ignore_patterns=[]) + def prepare_bei_encoder_torch_build_dir(self, build_dir: Path): + """prepares the build directory for an ENCODER_TORCH model to launch the BEI torch backend + (text-embeddings-router + vLLM). No engine build; the checkpoint is served directly.""" + config = self._spec.config + assert ( + config.trt_llm + and config.trt_llm.build + and config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER_TORCH + ), ( + "prepare_bei_encoder_torch_build_dir should only be called for ENCODER_TORCH tensorrt-llm model" + ) + assert isinstance(config.trt_llm.root, TRTLLMConfigurationV1), ( + "prepare_bei_encoder_torch_build_dir should only be called for inference_stack v1 tensorrt-llm model" + ) + trt_llm_config: TRTLLMConfigurationV1 = config.trt_llm.root + # Torch backend has no TRT-LLM 32-cap; 256 is a safe upper bound for KV-less + # embedding servers before per-request scheduling overhead dominates. + runtime_max_batch_size = min(trt_llm_config.build.max_batch_size, 256) + runtime_max_batch_tokens = max( + trt_llm_config.build.max_num_tokens, BEI_TORCH_REQUIRED_MAX_NUM_TOKENS + ) + port = 7997 + start_command = " ".join( + [ + "truss-transfer-cli /tmp/bei-model && text-embeddings-router --model-id /tmp/bei-model", + f"--port {port}", + f"--max-batch-requests {runtime_max_batch_size}", + f"--max-batch-tokens {runtime_max_batch_tokens}", + # sentences per JSON payload; capped for request-based autoscaling. + f"--max-client-batch-size {BEI_TRTLLM_CLIENT_BATCH_SIZE}", + # concurrent requests before returning 429. + # https://docs.baseten.co/performance/concurrency#concurrency-target + f"--max-concurrent-requests {BEI_MAX_CONCURRENCY_TARGET_REQUESTS}", + # stricter validation in newer BEI-torch images requires this to + # silently trim overlong inputs instead of 4xx-ing the request. + "--auto-truncate", + ] + ) + self._spec.config.docker_server = DockerServer( + start_command=f"/bin/sh -c '{start_command}'", + server_port=port, + predict_endpoint=trt_llm_config.runtime.webserver_default_route + or "/v1/embeddings", + readiness_endpoint="/health", + liveness_endpoint="/health", + ) + copy_tree_path(DOCKER_SERVER_TEMPLATES_DIR, build_dir, ignore_patterns=[]) + def prepare_trtllm_bert_encoder_build_dir(self, build_dir: Path): """prepares the build directory for a trtllm ENCODER model to launch a Baseten Embeddings Inference (BEI) server""" config = self._spec.config @@ -665,7 +714,11 @@ def prepare_trtllm_decoder_build_dir(self, build_dir: Path): config.trt_llm and config.trt_llm.build and config.trt_llm.build.base_model - not in [TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT] + not in [ + TrussTRTLLMModel.ENCODER, + TrussTRTLLMModel.ENCODER_BERT, + TrussTRTLLMModel.ENCODER_TORCH, + ] ), ( "prepare_trtllm_decoder_build_dir should only be called for decoder tensorrt-llm model" ) @@ -730,6 +783,8 @@ def prepare_image_build_dir( elif config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER_BERT: # Run the specific encoder_bert build self.prepare_trtllm_bert_encoder_build_dir(build_dir=build_dir) + elif config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER_TORCH: + self.prepare_bei_encoder_torch_build_dir(build_dir=build_dir) else: self.prepare_trtllm_decoder_build_dir(build_dir=build_dir) diff --git a/truss/tests/conftest.py b/truss/tests/conftest.py index 46974ef64..94c6e7870 100644 --- a/truss/tests/conftest.py +++ b/truss/tests/conftest.py @@ -816,6 +816,30 @@ def trtllm_config_encoder(default_config) -> Dict[str, Any]: return trtllm_config +@pytest.fixture +def trtllm_config_encoder_torch(default_config) -> Dict[str, Any]: + trtllm_config = default_config + trtllm_config["resources"] = { + "accelerator": Accelerator.H100.value, + "cpu": "1", + "memory": "30Gi", + "use_gpu": True, + "node_count": 1, + } + trtllm_config["trt_llm"] = { + "build": { + "base_model": "encoder_torch", + "checkpoint_repository": { + "source": "HF", + "repo": "nvidia/Nemotron-3-Embed-8B-BF16", + }, + "max_num_tokens": 32768, + }, + "runtime": {}, + } + return trtllm_config + + @pytest.fixture def deprecated_trtllm_config(default_config) -> Dict[str, Any]: trtllm_config = default_config diff --git a/truss/tests/trt_llm/test_trt_llm_config.py b/truss/tests/trt_llm/test_trt_llm_config.py index e005c22cc..6aaf065f8 100644 --- a/truss/tests/trt_llm/test_trt_llm_config.py +++ b/truss/tests/trt_llm/test_trt_llm_config.py @@ -116,6 +116,40 @@ def test_trt_llm_encoder(trtllm_config_encoder): assert config.build.plugin_configuration.use_paged_context_fmha is False +def test_trt_llm_encoder_torch_accepts_no_quant(trtllm_config_encoder_torch): + config = TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"]) + assert config.build.base_model.value == "encoder_torch" + assert config.build.quantization_type.value == "no_quant" + assert config.build.max_num_tokens == 32768 + + +def test_trt_llm_encoder_torch_rejects_build_time_quantization( + trtllm_config_encoder_torch, +): + trtllm_config_encoder_torch["trt_llm"]["build"]["quantization_type"] = "fp8" + with pytest.raises(pydantic.ValidationError, match="encoder_torch"): + TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"]) + + +def test_trt_llm_encoder_torch_rejects_speculator(trtllm_config_encoder_torch): + trtllm_config_encoder_torch["trt_llm"]["build"]["speculator"] = { + "speculative_decoding_mode": "LOOKAHEAD_DECODING", + "lookahead_windows_size": 4, + "lookahead_ngram_size": 3, + "lookahead_verification_set_size": 2, + } + with pytest.raises(pydantic.ValidationError, match="speculator"): + TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"]) + + +def test_trt_llm_encoder_torch_rejects_lora_adapters(trtllm_config_encoder_torch): + trtllm_config_encoder_torch["trt_llm"]["build"]["lora_adapters"] = { + "adapter1": {"source": "HF", "repo": "nvidia/Nemotron-3-Embed-8B-BF16-lora"} + } + with pytest.raises(pydantic.ValidationError, match="lora_adapters"): + TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"]) + + def test_trt_llm_encoder_autoconfig(trtllm_config_encoder): trt_llm_config = TRTLLMConfigurationV1(**trtllm_config_encoder["trt_llm"]) try: diff --git a/truss/trt_llm/config_checks.py b/truss/trt_llm/config_checks.py index f824defc2..5b5f39511 100644 --- a/truss/trt_llm/config_checks.py +++ b/truss/trt_llm/config_checks.py @@ -45,6 +45,7 @@ def add_openai_tag(tr: TrussHandle) -> str: if trt_llm_config.build.base_model in [ TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT, + TrussTRTLLMModel.ENCODER_TORCH, ]: return ("", False) # only briton requires openai-compatible tag, all others don't care about the openai tag