[Qwen 3.5][gemma4] Qwen35 and Gemma 4 fast inference (clean-base mirror of #588) #10
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -759,8 +759,8 @@ def grpo_accumulated_loss( | |
| unwrapped_model = trainer.accelerator.unwrap_model(trainer.model, keep_fp32_wrapper = False) | ||
|
|
||
| for module in unwrapped_model.modules(): | ||
| if hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "io_same_device"): | ||
| module._hf_hook.io_same_device = False | ||
| if hasattr(module, "rope_deltas"): | ||
| module.rope_deltas = None | ||
| pass | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,7 +17,7 @@ | |
| import torch | ||
| import os | ||
| from .common import TEMPORARY_PATCHES | ||
| from .utils import raise_error, patch_function | ||
|
|
||
|
|
||
| # ============================================================================ | ||
|
|
@@ -618,3 +618,56 @@ def forward(self, hidden_states, position_embeddings, attention_mask=None, **kwa | |
| Gemma4AudioAttention.forward = forward | ||
| pass | ||
| TEMPORARY_PATCHES.append(patch_Gemma4AudioAttention) | ||
|
|
||
|
|
||
| # Gemma-4 float16 MLP overflow fix. | ||
| # | ||
| # `down_proj(act_fn(gate_proj(x)) * up_proj(x))` overflows fp16 at layers.0 | ||
| # (E2B) / layers.1 (E4B): the product + fp16 matmul accumulator saturate to | ||
| # +-inf, poison the residual stream, and generation samples NaN logits that | ||
| # trip the CUDA categorical sampler on GRPO step ~2. | ||
| # | ||
| # Fix: fp32 gate*up, clamp to a safe bound, fp16 cast, nan_to_num on the | ||
| # down_proj output. Gated on gate output dtype so bf16/fp32 users see no | ||
| # change and no env flag is required. RMSNorm / Attention / Embedding | ||
| # patches are unnecessary (verified by bisection - identical loss/kl/grad | ||
| # trajectories). | ||
|
|
||
|
|
||
| def patch_Gemma4TextMLP(): | ||
| """fp16 overflow clamp for Gemma4TextMLP. | ||
|
|
||
| Does gate*up in fp32, clamps to a safe fp16 bound, then nan_to_nums | ||
| the down_proj output. Self-gated on gate dtype - no-op on bf16/fp32. | ||
| """ | ||
| try: | ||
| import transformers.models.gemma4.modeling_gemma4 as mod | ||
| except ImportError: | ||
| return | ||
| try: | ||
| Gemma4TextMLP = mod.Gemma4TextMLP | ||
| except AttributeError as e: | ||
| return raise_error("Gemma4TextMLP.forward", e) | ||
|
|
||
| # Largest value representable in both fp16 and bf16 (65536 rounds to | ||
| # fp16 inf). | ||
| _SAFE_FP16 = 65280.0 | ||
|
|
||
| def forward(self, x): | ||
| gate = self.gate_proj(x) | ||
| # Check matmul output dtype so autocast / PEFT fp16 casts are caught. | ||
| if gate.dtype != torch.float16: | ||
| return self.down_proj(self.act_fn(gate) * self.up_proj(x)) | ||
| product = self.act_fn(gate.float()) * self.up_proj(x).float() | ||
| product = torch.clamp(product, min=-_SAFE_FP16, max=_SAFE_FP16) | ||
| out = self.down_proj(product.to(gate.dtype)) | ||
| # Zero overflows so the residual identity path survives. | ||
| return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0) | ||
| try: | ||
| patch_function( | ||
| Gemma4TextMLP, "forward", forward, fullgraph=False, | ||
| ) | ||
| except Exception as e: | ||
| return raise_error("Gemma4TextMLP.forward", e) | ||
| pass | ||
| TEMPORARY_PATCHES.append(patch_Gemma4TextMLP) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1167,6 +1167,13 @@ def _is_fused_module(name: str) -> bool: | |
| if hasattr(layer, "layer_scalar"): | ||
| state_dict[f"{vllm_text_model_prefix}.layers.{kk}.layer_scalar"] = layer.layer_scalar.data | ||
| quant_state_dict[f"{vllm_text_model_prefix}.layers.{kk}.layer_scalar"] = layer.layer_scalar.data | ||
| for per_layer_name in ("per_layer_input_gate", "per_layer_projection"): | ||
| per_layer_module = getattr(layer, per_layer_name, None) | ||
| if per_layer_module is not None and hasattr(per_layer_module, "weight"): | ||
| get_state_dict( | ||
| f"{vllm_text_model_prefix}.layers.{kk}.{per_layer_name}", | ||
| 0, state_dict, per_layer_module, | ||
| ) | ||
| pass | ||
|
|
||
| if len(skipped_layernorms) != 0: | ||
|
|
@@ -1216,6 +1223,8 @@ def assert_same_state_dict(old_state_dict, new_state_dict): | |
| def _normalize_state_dict_tensor(value): | ||
| if isinstance(value, torch.nn.Parameter): | ||
| value = value.detach() | ||
| if not isinstance(value, torch.Tensor): | ||
| return value | ||
| if value.is_sparse: | ||
| value = value.to_dense() | ||
| return value.contiguous() | ||
|
Comment on lines
+1223
to
+1230
There was a problem hiding this comment. The def _normalize_state_dict_tensor(value):
if isinstance(value, torch.nn.Parameter):
value = value.detach()
return value.to_dense().contiguous() if value.is_sparse else value.contiguous() |
||
|
|
@@ -1396,8 +1405,14 @@ def _override_to(self, *args, **kwargs): | |
| if layer_name in quant_state_dict: | ||
| # for attributes of type nn.Parameter, there's no .weight | ||
| layer_name_br = re.sub(r"\.([\d]{1,})\.", r"[\1].", layer_name) | ||
| value = _unwrap_tensor(weight) | ||
| parent_expr, attr_name = layer_name_br.rsplit(".", 1) | ||
|
There was a problem hiding this comment. If |
||
| parent_module = eval(f"new_model.{parent_expr}") | ||
| if attr_name in getattr(parent_module, "_buffers", {}): | ||
| parent_module._buffers[attr_name] = value | ||
| else: | ||
| layer = torch.nn.Parameter(value, requires_grad = False) | ||
| exec(f"new_model.{layer_name_br} = layer") | ||
| continue | ||
| elif fp8_weight_scale is not None: | ||
| if fp8_weight_scale.ndim == 1: | ||
|
|
@@ -1448,9 +1463,23 @@ def _override_to(self, *args, **kwargs): | |
| layer.weight = torch.nn.Parameter(_unwrap_tensor(weight), requires_grad = False) | ||
| layer.bias = bias | ||
| else: | ||
| # LayerNorms (including vision norms) and depthwise Conv1d | ||
| weight_param = torch.nn.Parameter(_unwrap_tensor(weight), requires_grad=False) | ||
| layer_name_br = re.sub(r"\.([\d]{1,})\.", r"[\1].", layer_name) | ||
| if layer_name.endswith(".conv1d"): | ||
| target = eval(f"new_model.{layer_name_br}") | ||
| w = _unwrap_tensor(weight) | ||
| out_channels = w.shape[0] | ||
| kernel_size = w.shape[-1] | ||
| target.out_channels = out_channels | ||
| target.in_channels = out_channels | ||
| target.groups = out_channels | ||
| target.kernel_size = (kernel_size,) | ||
| target.padding = (kernel_size - 1,) | ||
| target.weight = weight_param | ||
| if bias is not None: | ||
| target.bias = bias | ||
| continue | ||
|
Comment on lines
+1469
to
+1482
There was a problem hiding this comment. The use of |
||
| # Set weight | ||
| exec(f"new_model.{layer_name_br}.weight = None") | ||
| exec(f"new_model.{layer_name_br}.weight = weight_param") | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Similar to the
patched_supports_lorafunction, usingisinstancehere is more robust than comparing class names as strings.