Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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"}'
```
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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}'
```
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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"
}'
```
Original file line number Diff line number Diff line change
@@ -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
27 changes: 27 additions & 0 deletions 11-embeddings-reranker-classification-torch/README.md
Original file line number Diff line number Diff line change
@@ -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.
Loading