diff --git a/scripts/verify_comment_only_diff.py b/scripts/verify_comment_only_diff.py index 068d3244df..90eafb7f8f 100644 --- a/scripts/verify_comment_only_diff.py +++ b/scripts/verify_comment_only_diff.py @@ -35,6 +35,7 @@ Example: git diff --name-only origin/main..HEAD \\ | xargs python scripts/verify_comment_only_diff.py --base origin/main """ + from __future__ import annotations import argparse @@ -49,7 +50,9 @@ import yaml def _git_show(rev: str, path: str) -> str: return subprocess.check_output( - ["git", "show", f"{rev}:{path}"], text = True, stderr = subprocess.DEVNULL, + ["git", "show", f"{rev}:{path}"], + text = True, + stderr = subprocess.DEVNULL, ) @@ -145,8 +148,7 @@ def _walk_yaml_diff(b: Any, a: Any, prefix: str = "") -> None: elif isinstance(b, list): if len(b) != len(a): print( - f" list len at {prefix or '/'}: " - f"{len(b)} -> {len(a)}", + f" list len at {prefix or '/'}: " f"{len(b)} -> {len(a)}", ) for i, (bi, ai) in enumerate(zip(b, a)): _walk_yaml_diff(bi, ai, f"{prefix}[{i}]") diff --git a/tests/test_gemma4_moe_4bit_swap.py b/tests/test_gemma4_moe_4bit_swap.py index 2093067624..f12b592dc4 100644 --- a/tests/test_gemma4_moe_4bit_swap.py +++ b/tests/test_gemma4_moe_4bit_swap.py @@ -6,6 +6,7 @@ that exercise the swap helper's shape contract, idempotence, and gating behaviour. The full repro (resident VRAM 46 GB -> 14.27 GB, cosine sim 0.994 vs BF16) is documented in the PR description. """ + import importlib import os @@ -54,12 +55,14 @@ def _stub_gemma4_module(): requires_grad = False, ) from transformers.activations import ACT2FN + module.act_fn = ACT2FN[_StubConfig.hidden_activation] return module def test_is_enabled_reads_env_var(): from unsloth.models import gemma4_moe_4bit + old = os.environ.pop("UNSLOTH_GEMMA4_MOE_4BIT", None) try: assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is False @@ -78,6 +81,7 @@ def test_swap_skips_models_without_gemma4_experts(): from unsloth.models.gemma4_moe_4bit import ( swap_gemma4_experts_to_per_expert_linear4bit, ) + model = nn.Sequential(nn.Linear(8, 8), nn.Linear(8, 8)) assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0 @@ -111,6 +115,7 @@ def test_swap_idempotent_on_stub_module_without_cuda(): from unsloth.models.gemma4_moe_4bit import ( swap_gemma4_experts_to_per_expert_linear4bit, ) + if not torch.cuda.is_available(): # CPU-only: bnb's Linear4bit init would fail. Validate the model-walk # path on an empty Sequential to confirm the helper is side-effect-free. diff --git a/unsloth/models/gemma4_moe_4bit.py b/unsloth/models/gemma4_moe_4bit.py index d8453beb98..0863fa477c 100644 --- a/unsloth/models/gemma4_moe_4bit.py +++ b/unsloth/models/gemma4_moe_4bit.py @@ -32,6 +32,7 @@ Gated on UNSLOTH_GEMMA4_MOE_4BIT (default off) and on load_in_4bit=True. Default off until the matching per-expert LoRA path lands; opt in via the env var if you want the VRAM win without QLoRA training. """ + from __future__ import annotations import os @@ -67,7 +68,8 @@ def _per_expert_forward( final_hidden_states = torch.zeros_like(hidden_states) with torch.no_grad(): expert_mask = torch.nn.functional.one_hot( - top_k_index, num_classes = self.num_experts, + top_k_index, + num_classes = self.num_experts, ) expert_mask = expert_mask.permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim = (-1, -2)), 0).nonzero() @@ -81,9 +83,7 @@ def _per_expert_forward( gate_up = self.gate_up_proj_4bit[expert_idx](current_state) gate, up = gate_up.chunk(2, dim = -1) current_hidden_states = self.act_fn(gate) * up - current_hidden_states = self.down_proj_4bit[expert_idx]( - current_hidden_states - ) + current_hidden_states = self.down_proj_4bit[expert_idx](current_hidden_states) current_hidden_states = ( current_hidden_states * top_k_weights[token_idx, top_k_pos, None] ) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index c459a153fe..68cfb8a3fe 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1060,6 +1060,7 @@ class FastBaseModel: is_gemma4_moe_4bit_enabled, swap_gemma4_experts_to_per_expert_linear4bit, ) + if is_gemma4_moe_4bit_enabled(): _swapped = swap_gemma4_experts_to_per_expert_linear4bit( model,