Memory spikes on long context lengths

[imomayiz]

When testing Qwen3-0.6B on context lengths > 8k, I get very high memory spikes , with bsz=1
For eg. a 16k context sample does not fit in 40gb VRAM.
Did anyone have this problem?

[afalf]

Could you please provide your inference code?

[imomayiz]

Hello @afalf,
I use SentenceTransformer with prompt_name = "document"

class Embedder:
      def __init__(self, model_name_or_path: str):
          if torch.cuda.is_available():
              torch.cuda.empty_cache()
          
          self.model = SentenceTransformer(
            embedding_model,
            model_kwargs={"device_map": "cuda:0"})
  
      def embed(self, texts: List[str], prompt_name: str) -> List[List[float]]:
          document_embeddings = self.model.encode(texts, prompt_name=prompt_name)
          return document_embeddings.tolist()

[DehaoDong]

Encountered same problem with long input length. And I notice that there's a discussion on huggingface repo (qwen3 embedding 0.6b) stating that they encountered memory spike when inputting mis-spelled words.

I tested qwen3 embedding 4b and 8b as well, this problem happens on all qwen3 embedding models as far as I know.