diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index ecfeec381..70c96fb84 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -345,12 +345,23 @@ class CodePredictorBaseModel(nn.Module): inputs_embeds: torch.Tensor, position_ids: torch.Tensor, ) -> torch.Tensor: + # Run the transformer body in float32 when the model is in fp16. + # fp16 lacks the dynamic range for stable attention scores and + # SiLU-gated MLP intermediates, producing NaN on GPUs without + # native bf16 support (Turing, Volta). The RMSNorm and RoPE + # layers already upcast internally; this extends the same + # treatment to attention and MLP. + input_dtype = inputs_embeds.dtype + use_fp32 = input_dtype == torch.float16 + if use_fp32: + inputs_embeds = inputs_embeds.float() hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - for layer in self.layers: - hidden_states = layer(hidden_states, position_embeddings) - hidden_states = self.norm(hidden_states) - return hidden_states + with torch.amp.autocast(inputs_embeds.device.type, enabled=use_fp32, dtype=torch.float32): + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings) + hidden_states = self.norm(hidden_states) + return hidden_states.to(input_dtype) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters(remove_duplicate=False)) @@ -686,12 +697,12 @@ class CodePredictorWrapper(nn.Module): logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) if s_top_p < 1.0: sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True) - sorted_probs = F.softmax(sorted_logits, dim=-1) + sorted_probs = F.softmax(sorted_logits, dim=-1, dtype=torch.float32) cumulative_probs = sorted_probs.cumsum(dim=-1) remove_mask = (cumulative_probs - sorted_probs) >= s_top_p sorted_logits[remove_mask] = float("-inf") logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) - probs = F.softmax(logits, dim=-1) + probs = F.softmax(logits, dim=-1, dtype=torch.float32) code = torch.multinomial(probs, num_samples=1) else: # "per_call" mode: temperature-scaled + top-k @@ -700,7 +711,7 @@ class CodePredictorWrapper(nn.Module): if top_k > 0: topk_vals, _ = scaled.topk(top_k, dim=-1) scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) - probs = F.softmax(scaled, dim=-1) + probs = F.softmax(scaled, dim=-1, dtype=torch.float32) code = torch.multinomial(probs, num_samples=1) else: code = logits.argmax(dim=-1, keepdim=True)