Quantize a single tensor obtained from a float32 model

[Boltzmachine]

I have a model consisting of two parts: the first is an encoder of float32, and the second is a quantized LLM.
If I load the LLM in bfloat16, I can do encoder(x).bfloat() and feed it into the LLM. But for LLM in 8bit, I cannot find a corresponding way to convert the output of the encoder.

[gyanu2507]

You do not quantize the encoder output to INT8 to feed an 8-bit LLM. The 8-bit part is the weights of the Linear layers. Activations stay fp16/bf16 (with the outlier split inside those layers).

So the same cast you already use is the right one:

hidden = encoder(x)
hidden = hidden.to(torch.float16)  # or bfloat16, matching bnb_4bit_compute_dtype / the 8-bit compute dtype
logits = llm_model(inputs_embeds=hidden)

There is no public encoder(x).int8() helper, and you do not want one. Linear8bitLt will quantize activations internally where it needs to.

Watch dtypes: if the encoder is fp32 and the LLM compute dtype is fp16, a silent overflow on the first projection is common. Cast once, and keep the sequence/hidden size the LLM was built for. inputs_embeds skips the token embedding table, which is what you want if the encoder already produced that space.