Fix soft-prompt optimization diverging to NaN in half precision

[matanbt]

The bug

SoftPromptOptimizer optimizes the tensor returned by the model's embedding layer, so the parameter inherits the model's dtype. Adam keeps its state in that same dtype, and in fp16 both grad ** 2 and the default eps=1e-8 flush to zero:

>>> torch.tensor(1e-8, dtype=torch.float16).item()      # eps
0.0
>>> torch.tensor(1e-4, dtype=torch.float16).pow(2).item()  # grad ** 2
0.0

So the very first step evaluates m_hat / (sqrt(0) + 0), the prompt becomes ±inf, and every subsequent loss is NaN:

# gte-modernbert-base, which loads in fp16
losses=[0.0089569091796875, nan, nan, nan]
best_trigger_emb=tensor([[-inf, inf, inf, ..., inf, -inf, -inf], ...])

This affects any model loaded in fp16/bf16. fp32 models are unaffected — on all-MiniLM-L6-v2 (fp32) the same code descends normally, which is why it went unnoticed.

The fix

Keep a float32 master copy of the soft prompt, as mixed-precision training does, and cast to the model dtype only for the forward/backward. The result's embedding is cast back, so the returned dtype is unchanged.

Why it wasn't caught

tests/conftest.py loads every model in float32 by design ("Both run on CPU in float32 so the suite is portable"), so no existing test exercises the half-precision path. The added test builds an fp16 encoder explicitly.

Verified both directions:

• unpatched optimizer + new test → fails with AssertionError: [-0.4306640625, nan, nan]
• patched optimizer + new test → passes

Full suite: 8 passed (excluding test_api_integrations.py, which needs credentials).

🤖 Generated with Claude Code