From 828fb0355bb0e54e374b50fa406b39674d88ad53 Mon Sep 17 00:00:00 2001 From: Bohan Yao Date: Mon, 3 Aug 2026 21:38:25 -0700 Subject: [PATCH] Add tier-1 BEI-torch example templates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../README.md | 34 +++++++++++++ .../config.yaml | 23 +++++++++ .../README.md | 50 +++++++++++++++++++ .../config.yaml | 23 +++++++++ .../README.md | 39 +++++++++++++++ .../config.yaml | 25 ++++++++++ .../README.md | 27 ++++++++++ 7 files changed, 221 insertions(+) create mode 100644 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/README.md create mode 100644 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/config.yaml create mode 100644 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md create mode 100644 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/config.yaml create mode 100644 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/README.md create mode 100644 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/config.yaml create mode 100644 11-embeddings-reranker-classification-torch/README.md diff --git a/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/README.md b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/README.md new file mode 100644 index 00000000..7590ab10 --- /dev/null +++ b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/README.md @@ -0,0 +1,34 @@ +# BEI-torch with nvidia/Nemotron-3-Embed-1B-BF16 + +Deployment for [Nemotron-3-Embed-1B-BF16](https://huggingface.co/nvidia/Nemotron-3-Embed-1B-BF16) +on the BEI-torch backend. Native embedding dim 2048; MRL heads at 512 / 1024 / 2048. + +Fits on an L4 (24 GB) at bf16. + +## Prerequisites + +1. [Baseten account](https://app.baseten.co/signup) and [API key](https://app.baseten.co/settings/account/api_keys). +2. Latest truss: `pip install --upgrade truss` +3. Set `hf_access_token` as a [Baseten secret](https://app.baseten.co/settings/secrets). + +## Known issue: stock `nvidia/Nemotron-3-Embed-1B-BF16` fails to load + +Same three JSON patches apply as for the 8B model — see +[`../BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md`](../BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md). +Point `checkpoint_repository.repo` at a fork with those patches applied. + +## Deploy + +```sh +git clone https://github.com/basetenlabs/truss-examples.git +cd 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16 +truss push --publish +``` + +## Call + +```bash +curl -X POST https://model-xxxxxx.api.baseten.co/environments/production/sync/v1/embeddings \ + -H "Authorization: Api-Key $BASETEN_API_KEY" \ + -d '{"input": "text string", "model": "model"}' +``` diff --git a/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/config.yaml b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/config.yaml new file mode 100644 index 00000000..2a5a95b3 --- /dev/null +++ b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16/config.yaml @@ -0,0 +1,23 @@ +model_metadata: + example_model_input: + encoding_format: float + input: text string + model: model +model_name: BEI-torch-nvidia-Nemotron-3-Embed-1B-bf16-truss-example +python_version: py39 +resources: + accelerator: L4 + cpu: '1' + memory: 15Gi + use_gpu: true +trt_llm: + build: + base_model: encoder_torch + checkpoint_repository: + repo: nvidia/Nemotron-3-Embed-1B-BF16 + revision: main + source: HF + max_num_tokens: 32768 + quantization_type: no_quant + runtime: + webserver_default_route: /v1/embeddings diff --git a/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md new file mode 100644 index 00000000..5e49b48b --- /dev/null +++ b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md @@ -0,0 +1,50 @@ +# BEI-torch with nvidia/Nemotron-3-Embed-8B-BF16 + +Deployment for [Nemotron-3-Embed-8B-BF16](https://huggingface.co/nvidia/Nemotron-3-Embed-8B-BF16) +on the BEI-torch backend (text-embeddings-router + vLLM). Native embedding dim 4096; MRL heads +at 512 / 1024 / 2048 / 4096. + +## Prerequisites + +1. [Baseten account](https://app.baseten.co/signup) and [API key](https://app.baseten.co/settings/account/api_keys). +2. Latest truss: `pip install --upgrade truss` +3. Set `hf_access_token` as a [Baseten secret](https://app.baseten.co/settings/secrets). + +## Known issue: stock `nvidia/Nemotron-3-Embed-8B-BF16` fails to load + +Three JSON files in the stock repo trip HF `transformers`; fork the repo and apply these +patches, then set `checkpoint_repository.repo` to your fork: + +1. In `tokenizer_config.json`: remove the `tokenizer_class: "TokenizersBackend"` key. +2. In `tokenizer_config.json`: remove `extra_special_tokens` (or convert the list to a dict). +3. In `config.json`: change `model_type: "llama_bidirec"` → `"llama"` and set + `architectures: ["LlamaBidirectionalModel"]`. + +The BEI-torch arch registry maps `LlamaBidirectionalModel` to a bidirectional-attention +Llama variant; `model_type: "llama"` only exists so the transformers config parser accepts +the file. This workaround will be unnecessary once the BEI-torch `llama_bidirec` serde +alias and tokenizer-config fallback land upstream. + +## Deploy + +```sh +git clone https://github.com/basetenlabs/truss-examples.git +cd 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16 +truss push --publish +``` + +## Call + +```bash +curl -X POST https://model-xxxxxx.api.baseten.co/environments/production/sync/v1/embeddings \ + -H "Authorization: Api-Key $BASETEN_API_KEY" \ + -d '{"input": "text string", "model": "model"}' +``` + +Matryoshka embeddings (native 4096; also supports 512, 1024, 2048): + +```bash +curl -X POST https://model-xxxxxx.api.baseten.co/environments/production/sync/v1/embeddings \ + -H "Authorization: Api-Key $BASETEN_API_KEY" \ + -d '{"input": "text string", "model": "model", "dimensions": 1024}' +``` diff --git a/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/config.yaml b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/config.yaml new file mode 100644 index 00000000..6997944f --- /dev/null +++ b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/config.yaml @@ -0,0 +1,23 @@ +model_metadata: + example_model_input: + encoding_format: float + input: text string + model: model +model_name: BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16-truss-example +python_version: py39 +resources: + accelerator: H100 + cpu: '1' + memory: 30Gi + use_gpu: true +trt_llm: + build: + base_model: encoder_torch + checkpoint_repository: + repo: nvidia/Nemotron-3-Embed-8B-BF16 + revision: main + source: HF + max_num_tokens: 32768 + quantization_type: no_quant + runtime: + webserver_default_route: /v1/embeddings diff --git a/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/README.md b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/README.md new file mode 100644 index 00000000..67bc1330 --- /dev/null +++ b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/README.md @@ -0,0 +1,39 @@ +# BEI-torch with nvidia/llama-nemotron-rerank-1b-v2 + +Deployment for [llama-nemotron-rerank-1b-v2](https://huggingface.co/nvidia/llama-nemotron-rerank-1b-v2) +on the BEI-torch backend. Reranker analog to Nemotron-3-Embed; scores a `(query, documents[])` +pair via `/rerank`. + +Fits on an L4 (24 GB) at bf16. + +## Prerequisites + +1. [Baseten account](https://app.baseten.co/signup) and [API key](https://app.baseten.co/settings/account/api_keys). +2. Latest truss: `pip install --upgrade truss` +3. Set `hf_access_token` as a [Baseten secret](https://app.baseten.co/settings/secrets). + +## Known issue: stock `nvidia/llama-nemotron-rerank-1b-v2` fails to load + +Same three JSON patches apply as for the Nemotron-3-Embed models — see +[`../BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md`](../BEI-torch-nvidia-Nemotron-3-Embed-8B-bf16/README.md). +Point `checkpoint_repository.repo` at a fork with those patches applied. + +## Deploy + +```sh +git clone https://github.com/basetenlabs/truss-examples.git +cd 11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2 +truss push --publish +``` + +## Call + +```bash +curl -X POST https://model-xxxxxx.api.baseten.co/environments/production/sync/rerank \ + -H "Authorization: Api-Key $BASETEN_API_KEY" \ + -d '{ + "query": "what is baseten?", + "documents": ["Baseten is a model deployment platform.", "Bananas are yellow."], + "model": "model" + }' +``` diff --git a/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/config.yaml b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/config.yaml new file mode 100644 index 00000000..d5c2af03 --- /dev/null +++ b/11-embeddings-reranker-classification-torch/BEI-torch-nvidia-llama-nemotron-rerank-1b-v2/config.yaml @@ -0,0 +1,25 @@ +model_metadata: + example_model_input: + query: what is baseten? + documents: + - Baseten is a model deployment platform. + - Bananas are yellow. + model: model +model_name: BEI-torch-nvidia-llama-nemotron-rerank-1b-v2-truss-example +python_version: py39 +resources: + accelerator: L4 + cpu: '1' + memory: 15Gi + use_gpu: true +trt_llm: + build: + base_model: encoder_torch + checkpoint_repository: + repo: nvidia/llama-nemotron-rerank-1b-v2 + revision: main + source: HF + max_num_tokens: 32768 + quantization_type: no_quant + runtime: + webserver_default_route: /rerank diff --git a/11-embeddings-reranker-classification-torch/README.md b/11-embeddings-reranker-classification-torch/README.md new file mode 100644 index 00000000..626e4ecf --- /dev/null +++ b/11-embeddings-reranker-classification-torch/README.md @@ -0,0 +1,27 @@ +# BEI-torch templates + +Reference truss configs for the BEI-torch backend (`base_model: encoder_torch`) — the torch/vLLM +runtime used for embedding, reranker, and classification models whose architecture the +TRT-LLM `encoder`/`encoder_bert` paths cannot compile (LlamaBidirectional / Ministral3-Embed, +Gemma2Embedding, jina-v3, etc.). + +Templates in this directory are hand-written pending integration with +`../11-embeddings-reranker-classification-tensorrt/templating/generate_templates.py`. +When picking `encoder_torch` vs `encoder`/`encoder_bert`, use: + +- **`encoder_bert`** — BERT-family (BERT, ModernBERT, XLM-R, jina-v2). Fastest cold-start. +- **`encoder`** — causal-arch models with FP8 support on H100/B200 (Qwen3-Embedding, SFR-Embedding). +- **`encoder_torch`** — everything else. LlamaBidirectional (Nemotron-3-Embed), Gemma2Embedding, + jina-v3, custom architectures BEI-torch supports but TRT-LLM does not. + +## Known issue with stock NVIDIA Nemotron-3-Embed repos + +The stock `nvidia/Nemotron-3-Embed-*` repos ship JSON that HF `transformers` rejects on load: + +1. `tokenizer_config.json` has `tokenizer_class: "TokenizersBackend"` (not a real class). +2. `tokenizer_config.json` has `extra_special_tokens` as a list; transformers expects a dict. +3. `config.json` has `model_type: "llama_bidirec"` which is not registered with transformers. + +Until the BEI-torch backend lands the `llama_bidirec` serde alias and a tokenizer-config +fallback, point `checkpoint_repository.repo` at a fork with those three fields patched. +See the per-template README for the exact patch set.