Decreasing warmup for use on spot instances

[mattkrick]

Feature request

I'd like to deploy TEI (CPU v1.8 serving Qwen3 0.6B quantized) on a spot instance. To do so, the container would have to start up in 14 seconds to maintain QoS.
Poking around, I saw MAX_WARMUP_SEQUENCE_LENGTH, but that appears to only be used for Intel HPU deployments. I already set my max tokens to a modest value (MAX_BATCH_TOKENS=1028).
Is there anything else I can do to drive down warmup time, or anything I can cache to speed it up for a subsequent startup?

Motivation

Reduce costs by moving TEI replicas to spot instances

Your contribution

Happy to write a PR if there's a suitable path forward, or update the docs if a strategy already exists

[br0kej]

Hey @mattkrick. I am investigating a similar problem. Did you get anywhere with this?

[alvarobartt]

Hey @mattkrick indeed setting --max-batch-tokens to a lower value than the default which is 16384 should do the work, but that means that only up to 1024 tokens will fit within a batch, which you need to be careful as that might be lower than the maximum model's sequence length, meaning that input sequences will be truncated to at most 1024 tokens.

So say that the warmup is also a check to ensure that your instance can handle up to 16384 tokens by default, which is basically 512 tokens per sequence with a default client batch size (on Candle) of 32.

In this case, as you'd like to serve Qwen3-Embedding-0.6B on CPU and you mention quantized I guess you're referring to ONNX, specifically https://huggingface.co/onnx-community/Qwen3-Embedding-0.6B-ONNX; and given that its max_position_embeddings i.e., maximum input sequence length, is 32768, that means that a "full sequence" would be at most 32768 tokens, and if larger truncated to it. Also note that given that the default batch size of ORT is 8, that'd mean to benefit from batching + maximum context length you'd need to be able to handle at most 8 x 32768 = 262144 tokens. So really if your use-case won't use sequences / documents that large, lowering the value for the --max-batch-tokens should do the work, anyway we're investigating on how to reduce that to a proper sensible default as per #792 (and other related PRs).

Sorry for the long response, but hope that helps clarify things a bit 🙏🏻

P.S. Let me know if reducing --max-batch-tokens 1024 did the work (+ nit that you mentioned MAX_BATCH_TOKENS=1028 but it should rather be a power of 8 instead for ensuring data alignment).

[mattkrick]

Thanks @alvarobartt for the thoughtful response. Unfortunately, setting --max-batch-tokens 1024 (apologize, 1028 was a typo) still results in startups of about 2 minutes.

The model I've been using is the INT8, but all of them in the Qwen family seem to take awhile. https://huggingface.co/Svenni551/Qwen3-Embedding-0.6B-ONNX-INT8

Just an idea, but if we had warmup read/write a hash of the TEI version + flags + model + hardware architecture, then we could connect that to a persisted volume. Before warmup, it could read the hash & know that warmup has passed on a similar deployment infra previously & skip it. Long cold start + fast restarts!