Add base_model=encoder_torch for BEI-torch backend by bohanyao-baseten · Pull Request #2592 · basetenlabs/truss

bohanyao-baseten · 2026-08-04T04:37:12Z

Summary

Make BEI-torch a first-class truss primitive by adding base_model: encoder_torch as a new value under trt_llm.build, alongside the existing encoder and encoder_bert. Users write:

trt_llm:
  build:
    base_model: encoder_torch
    checkpoint_repository:
      repo: nvidia/Nemotron-3-Embed-8B-BF16
      source: HF
    max_num_tokens: 32768
    quantization_type: no_quant

and the image builder auto-synthesizes a docker_server running text-embeddings-router + vLLM directly on the checkpoint — no TRT engine build step.

Unblocks Harvey / Turbopuffer (Nemotron-3-Embed-8B/1B) and any future embedding model whose architecture the TRT-LLM encoder path can't compile (LlamaBidirectional, Gemma2Embedding, jina-v3, etc.).

Changes

truss/base/trt_llm_config.py — new TrussTRTLLMModel.ENCODER_TORCH. _bei_specfic_migration rejects build-time quantization (torch backend has no engine build; users point at a pre-quantized checkpoint instead) and speculator (no draft/target loop for embedding/reranker). bei_torch_image on ImageVersions, bei_torch_version on VersionsOverrides. webserver_default_route auto-detection, the ForCausalLM warning, and the chunked-context skip all extended to include encoder_torch.
truss/config.schema.json — enum + bei_torch_version field.
truss/trt_llm/config_checks.py — encoder_torch joins the "skip openai-tag rewrite" branch alongside the other BEI variants.
truss/contexts/image_builder/serving_image_builder.py — new prepare_bei_encoder_torch_build_dir. Caps --max-batch-requests at 256 (torch backend has no TRT-LLM 32-cap), floors --max-batch-tokens at 32768, enables --auto-truncate (newer BEI-torch images require it).
truss/contexts/docker_build_setup.py — dispatch to the new bei_torch_image slot.
truss/tests/{conftest.py,trt_llm/test_trt_llm_config.py} — new fixture + three tests (accept no_quant, reject fp8, reject speculator).

Open questions for reviewers

bei_torch_image default tag — ImageVersions.bei_torch_image defaults to baseten/bei:torch-dev-placeholder; the backend-injected BEI_TORCH_IMAGE_URI constance (paired baseten PR) is the real source of truth. Needs a real semver tag before merge. Ask @michaelfeil.
--pooling in the start command — the plan proposed --pooling mean explicit. I dropped it because text-embeddings-router --pooling overrides 1_Pooling/config.json, and forcing mean would break last-token-pooling checkpoints (Qwen3-Embedding). Happy to add back if that's wrong.
weights: vs trt_llm.checkpoint_repository — kept the simpler path of reusing checkpoint_repository, per the plan's leaning (open question Update authors so that we can be contacted #4).

Test plan

uv run pytest truss/tests/trt_llm/test_trt_llm_config.py truss/tests/util/test_config_checks.py — 17 passed
uv run ruff check ... — clean
Push a Nemotron-3-Embed-1B truss with this branch, confirm MODEL_READY and /v1/embeddings returns valid embeddings (blocked on the paired basetenlabs/baseten PR to wire BEI_TORCH_IMAGE_URI through)

🤖 Generated with Claude Code

Make BEI-torch a first-class truss primitive alongside `encoder` and `encoder_bert`. Users write `base_model: encoder_torch` and the image builder emits a `docker_server` running text-embeddings-router + vLLM directly; no TRT engine build step. Targets embedding / reranker / classification models the TRT-LLM `encoder` path can't compile — LlamaBidirectional (Nemotron-3-Embed), Gemma2Embedding, jina-v3, etc. - `TrussTRTLLMModel.ENCODER_TORCH` on the enum + JSON schema. - `_bei_specfic_migration` rejects build-time quantization (torch backend has no engine build; users must point at a pre-quantized checkpoint instead). - `prepare_bei_encoder_torch_build_dir` emits the docker_server with 256-cap on --max-batch-requests, 32768-min on --max-batch-tokens, and --auto-truncate on by default. - docker_build_setup routes ENCODER_TORCH to a new `bei_torch_image` slot on ImageVersions (backend inserts the resolved tag). - `bei_torch_version` on VersionsOverrides for user pinning. - Chunked-context / route-autodetect / ForCausalLM-warning branches extended so encoder_torch is treated like the other BEI variants. See plan: https://app.notion.com/p/3ae91d24727381d5902de1d02944ead4 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

- `trt_llm_validation_v1` was still requiring an `openai-compatible` / `openai-non-compatible` model_metadata tag for encoder_torch, since the base_model wasn't in the exempt list. Added to the exemption next to ENCODER / ENCODER_BERT. - `_bei_specfic_migration` now rejects `speculator:` on encoder_torch with a helpful message (BEI-torch is embedding/reranker serving; no draft/target loop applies). - Restructured the encoder_torch / encoder branches from two `if`s to `if / elif` so accidental fall-through can't happen if the base_model set widens later. - Reworded the T4/V100 rejection message to name BEI-torch alongside TRT-LLM (encoder_torch shares the SM89+ floor with the vLLM runtime). - Added `test_trt_llm_encoder_torch_rejects_speculator`. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

CLAassistant · 2026-08-04T04:50:04Z

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

Three hand-written templates for the new top-level `bei_torch:` config block (BEI-torch backend). Lives in a new `11-embeddings-reranker-classification-torch/` directory mirroring the existing `-tensorrt/` layout. - BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16 (H100) - BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16 (L4) - BEI-torch-nvidia-llama-nemotron-rerank-1b-v2 (L4) Each README documents the three-JSON-patch workaround for the stock NVIDIA repos (transformers-incompatible tokenizer_class, extra_special_tokens list, unregistered `llama_bidirec` model_type) until BEI-torch lands the serde alias and tokenizer-config fallback. Not yet wired into `generate_templates.py` — follow-up work to add a `BEI_TORCH` `Solution` class before tier-2/3. Paired with: - basetenlabs/truss#2592 - basetenlabs/baseten#25084 Plan: https://app.notion.com/p/3ae91d24727381d5902de1d02944ead4 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Three hand-written templates for `base_model: encoder_torch` (BEI-torch backend). Lives in a new `11-embeddings-reranker- classification-torch/` directory mirroring the existing `-tensorrt/` layout. - BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16 (H100) - BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16 (L4) - BEI-torch-nvidia-llama-nemotron-rerank-1b-v2 (L4) Each README documents the three-JSON-patch workaround for the stock NVIDIA repos (transformers-incompatible tokenizer_class, extra_special_tokens list, unregistered `llama_bidirec` model_type) until BEI-torch lands the serde alias and tokenizer-config fallback. Not yet wired into `generate_templates.py` — follow-up work to add a `BEI_TORCH` `Solution` class before tier-2/3. Paired with: - basetenlabs/truss#2592 - basetenlabs/baseten#25084 Plan: https://app.notion.com/p/3ae91d24727381d5902de1d02944ead4 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…NUM_TOKENS - `_bei_specfic_migration` now rejects `lora_adapters` on encoder_torch. LoRA adapters are a TRT-LLM build-time concept; vLLM serves the checkpoint as-is, so adapters would be silently ignored. - Extracted the `32768` magic number in `prepare_bei_encoder_torch_build_dir` to a named constant `BEI_TORCH_REQUIRED_MAX_NUM_TOKENS`, matching the `BEI_REQUIRED_MAX_NUM_TOKENS` precedent. Comment now explains why the user's value is clamped upward (vLLM rejects --max-batch-tokens below this; long-context embedders need it). - Added `test_trt_llm_encoder_torch_rejects_lora_adapters`. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

bohanyao-baseten · 2026-08-10T17:54:44Z

@copilot resolve the merge conflicts in this pull request

…hout default Co-authored-by: bohanyao-baseten <290184426+bohanyao-baseten@users.noreply.github.com>

Copilot · 2026-08-10T17:56:43Z

@copilot resolve the merge conflicts in this pull request

Resolved the merge conflict in truss/base/trt_llm_config.py (commit 9410a3b8). The conflict was that main removed the default value from beibert_image while our branch had added the new bei_torch_image field alongside it. Both changes are now incorporated.

michaelfeil · 2026-08-18T17:28:33Z

 BEI_TRTLLM_CLIENT_BATCH_SIZE = 128
 BEI_MAX_CONCURRENCY_TARGET_REQUESTS = 2048
 BEI_REQUIRED_MAX_NUM_TOKENS = 16384
+# BEI-torch (vLLM) rejects --max-batch-tokens below this; long-context embedders


mhh, please remove this comment vLLM... internals

michaelfeil · 2026-08-18T17:28:58Z

 BEI_MAX_CONCURRENCY_TARGET_REQUESTS = 2048
 BEI_REQUIRED_MAX_NUM_TOKENS = 16384
+# BEI-torch (vLLM) rejects --max-batch-tokens below this; long-context embedders
+# (Nemotron-3-Embed at 32k) need the headroom. User max_num_tokens is clamped


not sure if max_num_tokens is clamped? WDYM? We are not clamping down, right?

michaelfeil · 2026-08-18T17:34:55Z

-            if (
-                "ForCausalLM" in arch
-                and self.build.base_model != TrussTRTLLMModel.DECODER
+            if "ForCausalLM" in arch and self.build.base_model not in (


puh, this is a hard one? Is that actually true? I forgot, but don't users need to convert models or do you wanna do it on startup?

- Remove comment referencing vLLM internals in constants.py - Remove comment about encoder_torch bidirectional override in trt_llm_config.py - Remove comment about vLLM max-batch-tokens in serving_image_builder.py The behavior is unchanged; these were just explanatory comments that referenced implementation details reviewers preferred to keep out.

bohanyao-baseten · 2026-08-18T19:50:46Z

Addressed all three review comments in fcfa5e2:

constants.py:30 — Removed the vLLM internals comment. The constant is just a floor value now.
constants.py:31 — You're right, the comment was misleading. The code does max(user_value, BEI_TORCH_REQUIRED_MAX_NUM_TOKENS) which is a floor (upgrade), not a clamp-down. Removed.
trt_llm_config.py:718 — For encoder_torch, BEI-torch (text-embeddings-router + vLLM) handles the causal→bidirectional attention override at runtime, so no user-side conversion is needed. The warning was overly cautious. Removed.

All tests pass, ruff clean.

bohanyao-baseten mentioned this pull request Aug 4, 2026

Add tier-1 BEI-torch example templates basetenlabs/truss-examples#584

Open

2 tasks

bohanyao-baseten force-pushed the encoder-torch branch from 22d758d to 12ebb78 Compare August 4, 2026 21:56

bohanyao-baseten changed the title ~~Add base_model=encoder_torch for BEI-torch backend~~ Add top-level bei_torch: config block for BEI-torch backend Aug 4, 2026

bohanyao-baseten force-pushed the encoder-torch branch from 6c28ee2 to 22d758d Compare August 4, 2026 22:19

bohanyao-baseten changed the title ~~Add top-level bei_torch: config block for BEI-torch backend~~ Add base_model=encoder_torch for BEI-torch backend Aug 4, 2026

bohanyao-baseten and others added 2 commits August 4, 2026 15:23

Merge branch 'main' into encoder-torch

6c084bc

bohanyao-baseten marked this pull request as ready for review August 5, 2026 21:34

Copilot started work on behalf of bohanyao-baseten August 10, 2026 17:55 View session

Resolve merge conflict: keep bei_torch_image, adopt beibert_image wit…

9410a3b

…hout default Co-authored-by: bohanyao-baseten <290184426+bohanyao-baseten@users.noreply.github.com>

Copilot finished work on behalf of bohanyao-baseten August 10, 2026 17:57

bohanyao-baseten added 3 commits August 11, 2026 08:27

Merge branch 'main' into encoder-torch

5404fbc

Merge branch 'main' into encoder-torch

6c4ee86

Merge branch 'main' into encoder-torch

f000092

michaelfeil reviewed Aug 18, 2026

View reviewed changes

Navigation Menu

Uh oh!

Add base_model=encoder_torch for BEI-torch backend - #2592

Add base_model=encoder_torch for BEI-torch backend#2592
bohanyao-baseten wants to merge 9 commits into
mainfrom
encoder-torch

bohanyao-baseten commented Aug 4, 2026 •

edited

Loading

Uh oh!

CLAassistant commented Aug 4, 2026

Uh oh!

bohanyao-baseten commented Aug 10, 2026

Uh oh!

Copilot AI commented Aug 10, 2026

Uh oh!

michaelfeil Aug 18, 2026

Uh oh!

michaelfeil Aug 18, 2026

Uh oh!

michaelfeil Aug 18, 2026

Uh oh!

bohanyao-baseten commented Aug 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

bohanyao-baseten commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Open questions for reviewers

Test plan

Uh oh!

CLAassistant commented Aug 4, 2026

Uh oh!

bohanyao-baseten commented Aug 10, 2026

Uh oh!

Copilot AI commented Aug 10, 2026

Uh oh!

michaelfeil Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

michaelfeil Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

michaelfeil Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

bohanyao-baseten commented Aug 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bohanyao-baseten commented Aug 4, 2026 •

edited

Loading