From 29a0218103c01a643f2f1ae156eb99617f7f5ba0 Mon Sep 17 00:00:00 2001 From: MatanBT Date: Sun, 2 Aug 2026 15:56:49 +0300 Subject: [PATCH 1/2] Fix soft-prompt optimization diverging to NaN in half precision `SoftPromptOptimizer` optimizes the trigger embeddings 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 -- so the first step evaluates `m_hat / (sqrt(0) + 0)`, the prompt becomes +-inf, and every subsequent loss is NaN. Repro (gte-modernbert-base, which loads fp16): losses=[0.0089569091796875, nan, nan, nan] 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. fp32 models were never affected, which is why this went unnoticed -- the test suite loads every model in float32 by design. The added test therefore builds an fp16 encoder explicitly; it fails on the unpatched optimizer with `[-0.4306640625, nan, nan]` and passes with the fix. Co-Authored-By: Claude Opus 5 --- tests/test_optimizer_encoder.py | 33 +++++++++++++++++++++++++++++++ tropt/optimizer/soft_optimizer.py | 20 ++++++++++++++++--- 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/tests/test_optimizer_encoder.py b/tests/test_optimizer_encoder.py index 5fd03dc..d123ec4 100644 --- a/tests/test_optimizer_encoder.py +++ b/tests/test_optimizer_encoder.py @@ -2,8 +2,13 @@ import math +from tropt.common import Targets from tropt.loss import SimilarityLoss +from tropt.model.huggingface import EncoderHFModel from tropt.optimizer import GASLITEOptimizer +from tropt.optimizer.soft_optimizer import SoftPromptOptimizer + +from tests.conftest import ENCODER_NAME def test_gaslite_optimizer_end_to_end( @@ -29,3 +34,31 @@ def test_gaslite_optimizer_end_to_end( assert math.isfinite(result.best_loss) assert result.losses is not None and len(result.losses) >= 1 assert len(result.losses) <= 2 + + +def test_soft_prompt_stays_finite_in_half_precision(encoder_templates): + """Soft-prompt optimization must not diverge when the model is fp16. + + Adam keeps its state in the parameter's dtype, and in fp16 both `grad ** 2` + and the default `eps=1e-8` flush to zero -- so the first step divides by + zero, the prompt becomes +-inf and every later loss is NaN. The rest of this + suite loads models in float32, where the bug is invisible. + """ + fp16_encoder = EncoderHFModel( + model_name=ENCODER_NAME, device="cpu", dtype="float16", + ) + targets = Targets( + target_vectors=fp16_encoder(["This product is excellent."]).detach(), + ) + + result = SoftPromptOptimizer( + model=fp16_encoder, loss=SimilarityLoss(), num_steps=3, seed=0, + ).optimize_trigger( + templates=encoder_templates, + initial_trigger="! ! ! ! !", + targets=targets, + ) + + assert result.losses is not None and len(result.losses) == 3 + assert all(math.isfinite(loss) for loss in result.losses), result.losses + assert math.isfinite(result.best_loss) diff --git a/tropt/optimizer/soft_optimizer.py b/tropt/optimizer/soft_optimizer.py index 6872609..96a252d 100644 --- a/tropt/optimizer/soft_optimizer.py +++ b/tropt/optimizer/soft_optimizer.py @@ -72,6 +72,17 @@ def optimize_trigger( trigger_embeds = self.model._embedding_layer(trigger_ids.unsqueeze(0)) # (1, trigger_seq_len, embd_dim) + # Optimize a float32 master copy of the soft prompt, casting to the + # model's dtype only for the forward/backward. + # + # Adam's update underflows in half precision: with fp16 parameters both + # `grad ** 2` and the default `eps=1e-8` flush to 0, so the very first + # step computes `m_hat / (0 + 0)` and the prompt becomes +-inf -- every + # subsequent loss is NaN. This bites any model loaded in fp16/bf16 + # (fp32 models were unaffected, which is why it went unnoticed). + model_dtype = trigger_embeds.dtype + trigger_embeds = trigger_embeds.float() + # Initialize the optimizer on the trigger embeddings optimizer = self.GDOptimizer([trigger_embeds], lr=self.learning_rate) @@ -83,19 +94,22 @@ def optimize_trigger( # Compute gradients w.r.t. trigger embeddings trigger_grad, curr_loss = self.model.compute_grad_from_embeds( loss_func=self.loss_func, - candidate_trigger_embeds=trigger_embeds, + candidate_trigger_embeds=trigger_embeds.to(model_dtype), normalize_grads=False, return_loss=True, ) # grad: (1, trigger_seq_len, embed_dim); loss: (1,) curr_loss = curr_loss.item() # Set gradient on trigger embeddings - trigger_embeds.grad = trigger_grad + trigger_embeds.grad = trigger_grad.float() # Adam step optimizer.step() - best.update(loss=curr_loss, trigger_emb=trigger_embeds.detach().squeeze(0)) + best.update( + loss=curr_loss, + trigger_emb=trigger_embeds.detach().squeeze(0).to(model_dtype), + ) self.log(loss=curr_loss, lr=optimizer.param_groups[0]["lr"], grad_norm=trigger_grad.norm().item()) result = best.to_result() From 5d96be893f21071dc6c721a34de27a08fbf3f376 Mon Sep 17 00:00:00 2001 From: MatanBT Date: Mon, 3 Aug 2026 14:52:20 +0300 Subject: [PATCH 2/2] Fixes --- tests/test_optimizer_encoder.py | 33 ------------------------------- tropt/optimizer/gbda_optimizer.py | 14 +++++++------ tropt/optimizer/pez_optimizer.py | 8 +++++--- tropt/optimizer/soft_optimizer.py | 11 +---------- 4 files changed, 14 insertions(+), 52 deletions(-) diff --git a/tests/test_optimizer_encoder.py b/tests/test_optimizer_encoder.py index d123ec4..5fd03dc 100644 --- a/tests/test_optimizer_encoder.py +++ b/tests/test_optimizer_encoder.py @@ -2,13 +2,8 @@ import math -from tropt.common import Targets from tropt.loss import SimilarityLoss -from tropt.model.huggingface import EncoderHFModel from tropt.optimizer import GASLITEOptimizer -from tropt.optimizer.soft_optimizer import SoftPromptOptimizer - -from tests.conftest import ENCODER_NAME def test_gaslite_optimizer_end_to_end( @@ -34,31 +29,3 @@ def test_gaslite_optimizer_end_to_end( assert math.isfinite(result.best_loss) assert result.losses is not None and len(result.losses) >= 1 assert len(result.losses) <= 2 - - -def test_soft_prompt_stays_finite_in_half_precision(encoder_templates): - """Soft-prompt optimization must not diverge when the model is fp16. - - Adam keeps its state in the parameter's dtype, and in fp16 both `grad ** 2` - and the default `eps=1e-8` flush to zero -- so the first step divides by - zero, the prompt becomes +-inf and every later loss is NaN. The rest of this - suite loads models in float32, where the bug is invisible. - """ - fp16_encoder = EncoderHFModel( - model_name=ENCODER_NAME, device="cpu", dtype="float16", - ) - targets = Targets( - target_vectors=fp16_encoder(["This product is excellent."]).detach(), - ) - - result = SoftPromptOptimizer( - model=fp16_encoder, loss=SimilarityLoss(), num_steps=3, seed=0, - ).optimize_trigger( - templates=encoder_templates, - initial_trigger="! ! ! ! !", - targets=targets, - ) - - assert result.losses is not None and len(result.losses) == 3 - assert all(math.isfinite(loss) for loss in result.losses), result.losses - assert math.isfinite(result.best_loss) diff --git a/tropt/optimizer/gbda_optimizer.py b/tropt/optimizer/gbda_optimizer.py index 9ed064c..62a35f0 100644 --- a/tropt/optimizer/gbda_optimizer.py +++ b/tropt/optimizer/gbda_optimizer.py @@ -146,17 +146,19 @@ def optimize_trigger( trigger_ids: Int[Tensor, "trigger_seq_len"] = tokenizer.encode_trigger(initial_trigger).to(self.model.device) vocab_size = self.model.vocab_size device = self.model.device + model_dtype = self.model.dtype trigger_seq_len = trigger_ids.shape[0] - # Initialize logit matrix theta (here called `trigger_probs`) + # Initialize logit matrix theta (here called `trigger_probs`). + # stored in high precision for the optimizer if self.init_mode == "random": trigger_probs: Float[Tensor, "seq_len vocab_size"] = ( - torch.randn(trigger_seq_len, vocab_size, device=device, dtype=self.model.dtype) + torch.randn(trigger_seq_len, vocab_size, device=device, dtype=torch.float32) * self.init_noise_scale ) else: # "from_trigger" trigger_probs: Float[Tensor, "seq_len vocab_size"] = torch.zeros( - trigger_seq_len, vocab_size, device=device, dtype=self.model.dtype, + trigger_seq_len, vocab_size, device=device, dtype=torch.float32, ) for i in range(trigger_seq_len): trigger_probs[i, trigger_ids[i]] = self.initial_coeff @@ -181,7 +183,7 @@ def optimize_trigger( # (we repeat `trigger_probs_samples` so the gradient computation will draw multiple samples (w/ gumbel-softmax) from the same (optimized) distribution.) trigger_probs_samples = trigger_probs.unsqueeze(0).repeat(self.n_grad_samples, 1, 1) trigger_grad = self.model.compute_grad_from_tokens( - candidate_trigger_probs=trigger_probs_samples, # (n_grad_samples, trigger_seq_len, vocab_size) + candidate_trigger_probs=trigger_probs_samples.to(model_dtype), # (n_grad_samples, trigger_seq_len, vocab_size) loss_func=self.loss_func, do_gumbel_softmax=True, gumbel_softmax_temp=temperature, @@ -191,7 +193,7 @@ def optimize_trigger( avg_grad = trigger_grad.mean(dim=0) # -> (trigger_seq_len, vocab_size) # take grad step: - trigger_probs.grad = avg_grad + trigger_probs.grad = avg_grad.float() if self.grad_clip_norm is not None: torch.nn.utils.clip_grad_norm_([trigger_probs], self.grad_clip_norm) optimizer.step() @@ -239,7 +241,7 @@ def optimize_trigger( best_trigger_ids=final_best_ids, losses=best.losses, trigger_strs=best.trigger_strs, - best_trigger_probs=trigger_probs.detach(), + best_trigger_probs=trigger_probs.detach().to(model_dtype), ) return result diff --git a/tropt/optimizer/pez_optimizer.py b/tropt/optimizer/pez_optimizer.py index 1ecfec9..a3b24cf 100644 --- a/tropt/optimizer/pez_optimizer.py +++ b/tropt/optimizer/pez_optimizer.py @@ -82,6 +82,8 @@ def optimize_trigger( # Initialize continuous embeddings from the initial trigger tokens trigger_embeds = self.model._embedding_layer(trigger_ids.unsqueeze(0)) # (1, trigger_seq_len, embed_dim) + model_dtype = trigger_embeds.dtype + trigger_embeds = trigger_embeds.float() # stored in high precision for the optimizer # Initialize optimizer on continuous embeddings optimizer = self.GDOptimizer( @@ -95,7 +97,7 @@ def optimize_trigger( # Forward projection: project continuous embeddings to nearest vocab tokens projected_ids, projected_embeds = self._project_to_vocab( - trigger_embeds.squeeze(0), embedding_matrix + trigger_embeds.squeeze(0).to(model_dtype), embedding_matrix ) # Compute gradient w.r.t. the projected embeddings. @@ -110,7 +112,7 @@ def optimize_trigger( curr_loss = curr_loss.item() # Set gradient on the continuous embeddings and step - trigger_embeds.grad = trigger_grad + trigger_embeds.grad = trigger_grad.float() optimizer.step() # Decode current discrete trigger for tracking @@ -122,7 +124,7 @@ def optimize_trigger( # Final projection and evaluation on discrete tokens final_ids, _ = self._project_to_vocab( - trigger_embeds.squeeze(0), embedding_matrix + trigger_embeds.squeeze(0).to(model_dtype), embedding_matrix ) final_trigger_str = tokenizer.decode_trigger(final_ids) final_loss = self.model.compute_loss_from_tokens( diff --git a/tropt/optimizer/soft_optimizer.py b/tropt/optimizer/soft_optimizer.py index 96a252d..6bfe788 100644 --- a/tropt/optimizer/soft_optimizer.py +++ b/tropt/optimizer/soft_optimizer.py @@ -71,17 +71,8 @@ def optimize_trigger( trigger_ids = tokenizer.encode_trigger(initial_trigger).to(self.model.device) trigger_embeds = self.model._embedding_layer(trigger_ids.unsqueeze(0)) # (1, trigger_seq_len, embd_dim) - - # Optimize a float32 master copy of the soft prompt, casting to the - # model's dtype only for the forward/backward. - # - # Adam's update underflows in half precision: with fp16 parameters both - # `grad ** 2` and the default `eps=1e-8` flush to 0, so the very first - # step computes `m_hat / (0 + 0)` and the prompt becomes +-inf -- every - # subsequent loss is NaN. This bites any model loaded in fp16/bf16 - # (fp32 models were unaffected, which is why it went unnoticed). model_dtype = trigger_embeds.dtype - trigger_embeds = trigger_embeds.float() + trigger_embeds = trigger_embeds.float() # stored in high precision for the optimizer # Initialize the optimizer on the trigger embeddings optimizer = self.GDOptimizer([trigger_embeds], lr=self.learning_rate)