From 1ee59c3110aa5f5131f78bdfc7738d174953e131 Mon Sep 17 00:00:00 2001 From: William collins Date: Thu, 12 Mar 2026 08:57:24 +0000 Subject: [PATCH 1/3] feat: Add CGGR (Confidence-Gated Gradient Routing) integration [skip pre-commit.ci] --- unsloth/__init__.py | 44 +---- unsloth/cggr/__init__.py | 70 ++++++++ unsloth/cggr/bridge.py | 346 +++++++++++++++++++++++++++++++++++++++ unsloth/cggr/router.py | 178 ++++++++++++++++++++ 4 files changed, 602 insertions(+), 36 deletions(-) create mode 100644 unsloth/cggr/__init__.py create mode 100644 unsloth/cggr/bridge.py create mode 100644 unsloth/cggr/router.py diff --git a/unsloth/__init__.py b/unsloth/__init__.py index dbacd551c4..aaa4f44a20 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -28,26 +28,16 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules] from .import_fixes import ( fix_message_factory_issue, check_fbgemm_gpu_version, - disable_broken_causal_conv1d, - disable_broken_vllm, - configure_amdgpu_asic_id_table_path, torchvision_compatibility_check, fix_diffusers_warnings, fix_huggingface_hub, ) -# Configure libdrm ids table path early so ROCm can resolve AMD GPU names. -configure_amdgpu_asic_id_table_path() -disable_broken_causal_conv1d() -disable_broken_vllm() fix_message_factory_issue() check_fbgemm_gpu_version() torchvision_compatibility_check() fix_diffusers_warnings() fix_huggingface_hub() -del configure_amdgpu_asic_id_table_path -del disable_broken_causal_conv1d -del disable_broken_vllm del fix_message_factory_issue del check_fbgemm_gpu_version del torchvision_compatibility_check @@ -89,7 +79,7 @@ from importlib.metadata import PackageNotFoundError # Check for unsloth_zoo try: unsloth_zoo_version = importlib_version("unsloth_zoo") - if Version(unsloth_zoo_version) < Version("2026.3.2"): + if Version(unsloth_zoo_version) < Version("2026.1.2"): print( "Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n" "Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`" @@ -135,65 +125,40 @@ from unsloth_zoo.device_type import ( from .import_fixes import ( fix_xformers_performance_issue, fix_vllm_aimv2_issue, - check_vllm_torch_sm100_compatibility, fix_vllm_guided_decoding_params, fix_vllm_pdl_blackwell, - fix_triton_compiled_kernel_missing_attrs, - patch_trunc_normal_precision_issue, ignore_logger_messages, patch_ipykernel_hf_xet, patch_trackio, patch_datasets, patch_enable_input_require_grads, fix_openenv_no_vllm, - patch_openspiel_env_async, fix_executorch, - patch_vllm_for_notebooks, - patch_torchcodec_audio_decoder, - disable_torchcodec_if_broken, - disable_broken_wandb, ) fix_xformers_performance_issue() fix_vllm_aimv2_issue() -# Check vLLM + torch < 2.9.0 + SM100 compatibility BEFORE importing vLLM -check_vllm_torch_sm100_compatibility() fix_vllm_guided_decoding_params() fix_vllm_pdl_blackwell() -fix_triton_compiled_kernel_missing_attrs() -patch_trunc_normal_precision_issue() ignore_logger_messages() patch_ipykernel_hf_xet() patch_trackio() patch_datasets() patch_enable_input_require_grads() fix_openenv_no_vllm() -patch_openspiel_env_async() fix_executorch() -patch_vllm_for_notebooks() -patch_torchcodec_audio_decoder() -disable_torchcodec_if_broken() -disable_broken_wandb() del fix_xformers_performance_issue del fix_vllm_aimv2_issue -del check_vllm_torch_sm100_compatibility del fix_vllm_guided_decoding_params del fix_vllm_pdl_blackwell -del fix_triton_compiled_kernel_missing_attrs -del patch_trunc_normal_precision_issue del ignore_logger_messages del patch_ipykernel_hf_xet del patch_trackio del patch_datasets del patch_enable_input_require_grads del fix_openenv_no_vllm -del patch_openspiel_env_async del fix_executorch -del patch_vllm_for_notebooks -del patch_torchcodec_audio_decoder -del disable_torchcodec_if_broken -del disable_broken_wandb # Torch 2.4 has including_emulation if DEVICE_TYPE == "cuda": @@ -315,6 +280,12 @@ from .chat_templates import * from .tokenizer_utils import * from .trainer import * +# CGGR (Confidence-Gated Gradient Routing) integration +# Optional: requires `pip install cggr` for full functionality +from .cggr import CGGR_AVAILABLE +if CGGR_AVAILABLE: + from .cggr import CGGRUnslothBridge, patch_trainer_for_cggr, create_truncated_router + # Export dataprep utilities for CLI and downstream users from .dataprep.raw_text import RawTextDataLoader, TextPreprocessor from unsloth_zoo.rl_environments import ( @@ -328,3 +299,4 @@ from unsloth_zoo.rl_environments import ( # Patch TRL trainers for backwards compatibility _patch_trl_trainer() + diff --git a/unsloth/cggr/__init__.py b/unsloth/cggr/__init__.py new file mode 100644 index 0000000000..8c30df714b --- /dev/null +++ b/unsloth/cggr/__init__.py @@ -0,0 +1,70 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +CGGR (Confidence-Gated Gradient Routing) integration for Unsloth. + +This module provides selective backpropagation via label masking, enabling +~1.5-2x speedup in backward pass by only computing gradients for "hard" tokens. + +Usage: + from unsloth.cggr import CGGRUnslothBridge + + trainer = SFTTrainer(...) + CGGRUnslothBridge.patch_trainer(trainer, min_tokens_ratio=0.25) + trainer.train() + +Requires: pip install cggr +""" + +import logging + +logger = logging.getLogger(__name__) + +__all__ = [ + "CGGRUnslothBridge", + "patch_trainer_for_cggr", + "create_truncated_router", + "CGGR_AVAILABLE", +] + +# Check if CGGR package is available +try: + import cggr + CGGR_AVAILABLE = True +except ImportError: + CGGR_AVAILABLE = False + logger.debug("CGGR package not installed. Install with: pip install cggr") + +# Conditional imports +if CGGR_AVAILABLE: + from .router import create_truncated_router, TruncatedRouter + from .bridge import CGGRUnslothBridge, patch_trainer_for_cggr +else: + # Provide stub implementations that raise helpful errors + def _cggr_not_available(*args, **kwargs): + raise ImportError( + "CGGR is not installed. Install with: pip install cggr\n" + "For CUDA acceleration: pip install cggr[cuda]" + ) + + create_truncated_router = _cggr_not_available + CGGRUnslothBridge = type("CGGRUnslothBridge", (), { + "patch_trainer": staticmethod(_cggr_not_available), + }) + patch_trainer_for_cggr = _cggr_not_available + + class TruncatedRouter: + def __init__(self, *args, **kwargs): + _cggr_not_available() diff --git a/unsloth/cggr/bridge.py b/unsloth/cggr/bridge.py new file mode 100644 index 0000000000..e429fd953c --- /dev/null +++ b/unsloth/cggr/bridge.py @@ -0,0 +1,346 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +CGGR Bridge for Unsloth integration. + +Provides trainer patching to enable selective backpropagation via label masking. +""" + +import torch +import torch.nn.functional as F +from functools import wraps +from typing import Optional, Dict, Any, Callable +import logging +import warnings + +from .router import TruncatedRouter, create_truncated_router + +logger = logging.getLogger(__name__) + +__all__ = ["CGGRUnslothBridge", "patch_trainer_for_cggr"] + + +class CGGRUnslothBridge: + """ + Bridge class for integrating CGGR with Unsloth trainers. + + Patches the trainer's compute_loss method to apply label masking + before the forward pass, enabling selective gradient computation. + + Example: + >>> from unsloth.cggr import CGGRUnslothBridge + >>> trainer = SFTTrainer(...) + >>> CGGRUnslothBridge.patch_trainer(trainer, min_tokens_ratio=0.25) + >>> trainer.train() + """ + + def __init__( + self, + model: torch.nn.Module, + min_tokens_ratio: float = 0.25, + num_router_layers: int = 2, + warmup_steps: int = 100, + scoring: str = "entropy", + dynamic_threshold: bool = True, + ): + """ + Initialize CGGR bridge. + + Args: + model: The model being trained + min_tokens_ratio: Minimum fraction of tokens to keep gradients for (0.25 = top 25% hardest) + num_router_layers: Number of layers for the truncated router (default: 2) + warmup_steps: Steps before enabling CGGR (train normally first) + scoring: Scoring strategy ('entropy', 'margin', 'loss', 'combined') + dynamic_threshold: Whether to adjust ratio based on batch confidence + """ + self.model = model + self.min_tokens_ratio = min_tokens_ratio + self.num_router_layers = num_router_layers + self.warmup_steps = warmup_steps + self.scoring = scoring + self.dynamic_threshold = dynamic_threshold + + # Create truncated router for difficulty scoring + self.router = create_truncated_router(model, num_layers=num_router_layers) + + # Training state + self.current_step = 0 + self.total_tokens_seen = 0 + self.hard_tokens_seen = 0 + + logger.info( + f"Initialized CGGR Bridge: min_ratio={min_tokens_ratio}, " + f"router_layers={num_router_layers}, warmup={warmup_steps}" + ) + + @torch.inference_mode() + def compute_difficulty_scores( + self, + input_ids: torch.LongTensor, + labels: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Compute difficulty scores for each token using the truncated router. + + Args: + input_ids: Input token IDs [batch, seq_len] + labels: Target labels [batch, seq_len] + attention_mask: Attention mask [batch, seq_len] + + Returns: + difficulty_scores: Per-token difficulty [batch, seq_len] + """ + # Get logits from truncated router (fast forward pass) + logits = self.router(input_ids, attention_mask=attention_mask) + + # Compute difficulty based on scoring strategy + if self.scoring == "entropy": + # High entropy = uncertain = hard + # Use log_softmax for numerical stability (single fused kernel) + log_probs = F.log_softmax(logits, dim=-1) + probs = log_probs.exp() + scores = -torch.sum(probs * log_probs, dim=-1) + elif self.scoring == "margin": + # Small margin between top-2 = hard + # topk is efficient - only partial sort needed + top2 = torch.topk(logits, k=2, dim=-1).values + scores = -(top2[..., 0] - top2[..., 1]) # Negative margin (high = hard) + elif self.scoring == "loss": + # High loss = hard - directly compute per-token loss + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + scores = F.cross_entropy( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1), + reduction="none", + ignore_index=-100, + ).view(shift_labels.shape) + # Pad to match original sequence length + scores = F.pad(scores, (0, 1), value=0) + else: # combined - efficient fused computation + # Compute log_softmax once (fused kernel) + log_probs = F.log_softmax(logits, dim=-1) + probs = log_probs.exp() + + # Entropy from log_probs (reuse computation) + entropy = -torch.sum(probs * log_probs, dim=-1) + + # Margin from topk + top2 = torch.topk(logits, k=2, dim=-1).values + margin = top2[..., 0] - top2[..., 1] + + # Normalize and combine - use in-place operations where possible + entropy_mean = entropy.mean() + entropy_std = entropy.std() + 1e-10 + margin_mean = margin.mean() + margin_std = margin.std() + 1e-10 + + # Combined score: high entropy OR small margin = hard + scores = (entropy - entropy_mean) / entropy_std - (margin - margin_mean) / margin_std + + return scores + + def mask_easy_tokens( + self, + input_ids: torch.LongTensor, + labels: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.LongTensor: + """ + Mask easy tokens in labels with -100 to skip their gradients. + + Args: + input_ids: Input token IDs [batch, seq_len] + labels: Target labels [batch, seq_len] + attention_mask: Attention mask [batch, seq_len] + + Returns: + masked_labels: Labels with easy tokens set to -100 + """ + # During warmup, don't mask anything + if self.current_step < self.warmup_steps: + return labels + + # Clone labels to avoid modifying original + masked_labels = labels.clone() + + # Compute difficulty scores + scores = self.compute_difficulty_scores(input_ids, labels, attention_mask) + + # Get valid (non-ignored) token mask + valid_mask = labels != -100 + + # Compute ratio (avoid CPU sync by keeping on GPU) + ratio = self.min_tokens_ratio + if self.dynamic_threshold and valid_mask.any(): + # More confident batch → keep fewer tokens + valid_scores = scores.masked_select(valid_mask) + score_range = valid_scores.max() - valid_scores.min() + 1e-10 + mean_normalized = (valid_scores.mean() - valid_scores.min()) / score_range + # Lower mean score = more confident = keep fewer tokens + confidence = 1.0 - mean_normalized + ratio = self.min_tokens_ratio + (1.0 - self.min_tokens_ratio) * (1.0 - confidence) * 0.5 + ratio = max(self.min_tokens_ratio, min(ratio.item(), 1.0)) + + # Vectorized masking: compute per-sequence thresholds + batch_size, seq_len = labels.shape + + # Set scores of invalid tokens to -inf so they're never selected as "hard" + scores_for_threshold = scores.clone() + scores_for_threshold.masked_fill_(~valid_mask, float('-inf')) + + # Count valid tokens per sequence + valid_counts = valid_mask.sum(dim=1) # [batch] + + # Compute number to keep per sequence + num_keep = (valid_counts.float() * ratio).long().clamp(min=1) + + # For each sequence, find the threshold score (k-th largest) + # Use topk to find scores we should keep + max_valid = valid_counts.max().item() + if max_valid > 0: + # Sort scores descending to find threshold + sorted_scores, _ = scores_for_threshold.sort(dim=1, descending=True) + + # Get threshold for each sequence (the num_keep-th highest score) + # Clamp indices to valid range + threshold_indices = (num_keep - 1).clamp(min=0, max=seq_len - 1) + thresholds = sorted_scores.gather(1, threshold_indices.unsqueeze(1)).squeeze(1) # [batch] + + # Mask tokens with scores below threshold + below_threshold = scores < thresholds.unsqueeze(1) + mask_tokens = below_threshold & valid_mask + masked_labels.masked_fill_(mask_tokens, -100) + + # Update statistics (use item() only once at end) + total = valid_mask.sum().item() + kept = (masked_labels != -100).sum().item() + self.total_tokens_seen += total + self.hard_tokens_seen += kept + + return masked_labels + + def step(self): + """Called after each training step to update internal state.""" + self.current_step += 1 + + def get_stats(self) -> Dict[str, float]: + """Get CGGR statistics for logging.""" + if self.total_tokens_seen == 0: + return {"cggr/hard_ratio": 0.0, "cggr/step": self.current_step} + return { + "cggr/hard_ratio": self.hard_tokens_seen / self.total_tokens_seen, + "cggr/step": self.current_step, + "cggr/total_tokens": self.total_tokens_seen, + } + + @classmethod + def patch_trainer( + cls, + trainer, + min_tokens_ratio: float = 0.25, + num_router_layers: int = 2, + warmup_steps: int = 100, + scoring: str = "entropy", + dynamic_threshold: bool = True, + ) -> "CGGRUnslothBridge": + """ + Patch a trainer to use CGGR selective backpropagation. + + Args: + trainer: HuggingFace/TRL trainer instance + min_tokens_ratio: Minimum fraction of tokens to keep (0.25 = 25% hardest) + num_router_layers: Layers for difficulty scoring router + warmup_steps: Train normally for this many steps first + scoring: Scoring strategy ('entropy', 'margin', 'loss', 'combined') + dynamic_threshold: Adjust ratio based on batch confidence + + Returns: + CGGRUnslothBridge instance (for accessing stats) + + Example: + >>> bridge = CGGRUnslothBridge.patch_trainer(trainer) + >>> trainer.train() + >>> print(bridge.get_stats()) + """ + # Create bridge instance + bridge = cls( + model=trainer.model, + min_tokens_ratio=min_tokens_ratio, + num_router_layers=num_router_layers, + warmup_steps=warmup_steps, + scoring=scoring, + dynamic_threshold=dynamic_threshold, + ) + + # Store reference on trainer + trainer._cggr_bridge = bridge + + # Patch compute_loss to apply label masking + original_compute_loss = trainer.compute_loss + + @wraps(original_compute_loss) + def cggr_compute_loss(model, inputs, *args, **kwargs): + # Apply CGGR label masking + if "labels" in inputs and inputs["labels"] is not None: + inputs = dict(inputs) # Don't modify original + inputs["labels"] = bridge.mask_easy_tokens( + input_ids=inputs.get("input_ids"), + labels=inputs["labels"], + attention_mask=inputs.get("attention_mask"), + ) + + # Call original compute_loss + outputs = original_compute_loss(model, inputs, *args, **kwargs) + + # Update step counter + bridge.step() + + return outputs + + trainer.compute_loss = cggr_compute_loss + + print(f"🦥 Unsloth + CGGR: Selective backpropagation enabled!") + print(f" → Keeping {min_tokens_ratio*100:.0f}% hardest tokens for gradient computation") + print(f" → Router uses {num_router_layers} layers, warmup={warmup_steps} steps") + + return bridge + + +def patch_trainer_for_cggr( + trainer, + min_tokens_ratio: float = 0.25, + **kwargs, +) -> CGGRUnslothBridge: + """ + Convenience function to patch a trainer for CGGR. + + Equivalent to CGGRUnslothBridge.patch_trainer(). + + Args: + trainer: Trainer instance to patch + min_tokens_ratio: Fraction of tokens to keep (0.25 = 25% hardest) + **kwargs: Additional arguments passed to CGGRUnslothBridge.patch_trainer() + + Returns: + CGGRUnslothBridge instance + """ + return CGGRUnslothBridge.patch_trainer( + trainer, + min_tokens_ratio=min_tokens_ratio, + **kwargs, + ) diff --git a/unsloth/cggr/router.py b/unsloth/cggr/router.py new file mode 100644 index 0000000000..52c8300028 --- /dev/null +++ b/unsloth/cggr/router.py @@ -0,0 +1,178 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Truncated router for CGGR difficulty scoring. + +Creates a lightweight router from the first N layers of a model to score +token difficulty without running the full forward pass. +""" + +import torch +import torch.nn as nn +from typing import Optional, Tuple, Union +import logging + +logger = logging.getLogger(__name__) + +__all__ = ["TruncatedRouter", "create_truncated_router"] + + +class TruncatedRouter(nn.Module): + """ + A truncated version of a language model using only the first N layers. + + Used for fast difficulty scoring in CGGR. Shares weights with the parent + model, so uses zero additional memory. + + Args: + model: The parent HuggingFace model + num_layers: Number of decoder layers to use (default: 2) + """ + + def __init__(self, model: nn.Module, num_layers: int = 2): + super().__init__() + self.num_layers = num_layers + + # Get the base model (handle PEFT wrapping) + base_model = model + if hasattr(model, "base_model"): + base_model = model.base_model + if hasattr(base_model, "model"): + base_model = base_model.model + + # Store reference to model components (shares weights, no copy) + self.embed_tokens = self._get_embed_tokens(base_model) + self.layers = self._get_layers(base_model, num_layers) + self.norm = self._get_norm(base_model) + self.lm_head = self._get_lm_head(model, base_model) + + # Store config for reference + self.config = getattr(base_model, "config", None) + self.dtype = next(model.parameters()).dtype + self.device = next(model.parameters()).device + + def _get_embed_tokens(self, model: nn.Module) -> nn.Module: + """Extract embedding layer from model.""" + if hasattr(model, "embed_tokens"): + return model.embed_tokens + if hasattr(model, "model") and hasattr(model.model, "embed_tokens"): + return model.model.embed_tokens + if hasattr(model, "transformer") and hasattr(model.transformer, "wte"): + return model.transformer.wte # GPT-2 style + raise ValueError(f"Cannot find embedding layer in model: {type(model)}") + + def _get_layers(self, model: nn.Module, num_layers: int) -> nn.ModuleList: + """Extract first N decoder layers.""" + layers = None + if hasattr(model, "layers"): + layers = model.layers + elif hasattr(model, "model") and hasattr(model.model, "layers"): + layers = model.model.layers + elif hasattr(model, "transformer") and hasattr(model.transformer, "h"): + layers = model.transformer.h # GPT-2 style + elif hasattr(model, "encoder") and hasattr(model.encoder, "layer"): + layers = model.encoder.layer # BERT style + + if layers is None: + raise ValueError(f"Cannot find decoder layers in model: {type(model)}") + + # Return reference to first N layers (shares weights) + return nn.ModuleList([layers[i] for i in range(min(num_layers, len(layers)))]) + + def _get_norm(self, model: nn.Module) -> Optional[nn.Module]: + """Extract final normalization layer.""" + if hasattr(model, "norm"): + return model.norm + if hasattr(model, "model") and hasattr(model.model, "norm"): + return model.model.norm + if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"): + return model.transformer.ln_f # GPT-2 style + return None + + def _get_lm_head(self, original_model: nn.Module, base_model: nn.Module) -> nn.Module: + """Extract language model head.""" + if hasattr(original_model, "lm_head"): + return original_model.lm_head + if hasattr(base_model, "lm_head"): + return base_model.lm_head + if hasattr(original_model, "base_model") and hasattr(original_model.base_model, "lm_head"): + return original_model.base_model.lm_head + raise ValueError(f"Cannot find lm_head in model: {type(original_model)}") + + @torch.inference_mode() + def forward( + self, + input_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """ + Forward pass through truncated model to get logits for scoring. + + Args: + input_ids: Input token IDs [batch, seq_len] + attention_mask: Attention mask [batch, seq_len] + + Returns: + logits: Output logits [batch, seq_len, vocab_size] + """ + # Embeddings + hidden_states = self.embed_tokens(input_ids) + + # Pass through truncated layers + for layer in self.layers: + layer_outputs = layer( + hidden_states, + attention_mask=attention_mask, + use_cache=False, + ) + hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs + + # Apply final norm if available + if self.norm is not None: + hidden_states = self.norm(hidden_states) + + # Get logits + logits = self.lm_head(hidden_states) + + return logits + + +def create_truncated_router( + model: nn.Module, + num_layers: int = 2, +) -> TruncatedRouter: + """ + Create a truncated router from a model for CGGR difficulty scoring. + + The router shares weights with the parent model, so uses zero additional + GPU memory. It only runs the first N layers to quickly estimate token + difficulty. + + Args: + model: HuggingFace model (can be PEFT-wrapped) + num_layers: Number of decoder layers to use (default: 2) + + Returns: + TruncatedRouter instance + + Example: + >>> from unsloth import FastLanguageModel + >>> model, tokenizer = FastLanguageModel.from_pretrained(...) + >>> router = create_truncated_router(model, num_layers=2) + """ + router = TruncatedRouter(model, num_layers=num_layers) + logger.info(f"Created truncated router with {num_layers} layers for CGGR scoring") + return router From 5c5d7de2af5dd7ef09907bcf4a78e563de314772 Mon Sep 17 00:00:00 2001 From: William collins Date: Thu, 12 Mar 2026 08:57:35 +0000 Subject: [PATCH 2/3] feat: Add native CGGR support for SFTTrainer [skip pre-commit.ci] --- unsloth/cggr/bridge.py | 49 ++++++++++++++++++++++++------------------ unsloth/cggr/router.py | 19 +++++++++++++++- 2 files changed, 46 insertions(+), 22 deletions(-) diff --git a/unsloth/cggr/bridge.py b/unsloth/cggr/bridge.py index e429fd953c..0701cd8d6e 100644 --- a/unsloth/cggr/bridge.py +++ b/unsloth/cggr/bridge.py @@ -76,10 +76,11 @@ class CGGRUnslothBridge: # Create truncated router for difficulty scoring self.router = create_truncated_router(model, num_layers=num_router_layers) - # Training state + # Training state (keep on device to avoid syncs) self.current_step = 0 - self.total_tokens_seen = 0 - self.hard_tokens_seen = 0 + self.device = next(model.parameters()).device + self.total_tokens_seen = torch.tensor(0, device=self.device, dtype=torch.long) + self.hard_tokens_seen = torch.tensor(0, device=self.device, dtype=torch.long) logger.info( f"Initialized CGGR Bridge: min_ratio={min_tokens_ratio}, " @@ -184,17 +185,22 @@ class CGGRUnslothBridge: # Get valid (non-ignored) token mask valid_mask = labels != -100 - # Compute ratio (avoid CPU sync by keeping on GPU) - ratio = self.min_tokens_ratio - if self.dynamic_threshold and valid_mask.any(): + # Compute ratio (stay on GPU to avoid sync) + ratio = torch.tensor(self.min_tokens_ratio, device=self.device, dtype=scores.dtype) + if self.dynamic_threshold: # More confident batch → keep fewer tokens + # We use a cautious approach: only adjust if we have valid scores + # Use where to avoid sync if possible, but a few ops are fine here valid_scores = scores.masked_select(valid_mask) - score_range = valid_scores.max() - valid_scores.min() + 1e-10 - mean_normalized = (valid_scores.mean() - valid_scores.min()) / score_range - # Lower mean score = more confident = keep fewer tokens - confidence = 1.0 - mean_normalized - ratio = self.min_tokens_ratio + (1.0 - self.min_tokens_ratio) * (1.0 - confidence) * 0.5 - ratio = max(self.min_tokens_ratio, min(ratio.item(), 1.0)) + if valid_scores.numel() > 0: + s_min = valid_scores.min() + s_max = valid_scores.max() + score_range = s_max - s_min + 1e-10 + mean_normalized = (valid_scores.mean() - s_min) / score_range + # Lower mean score = more confident = keep fewer tokens + confidence = 1.0 - mean_normalized + ratio = self.min_tokens_ratio + (1.0 - self.min_tokens_ratio) * (1.0 - confidence) * 0.5 + ratio = ratio.clamp(min=self.min_tokens_ratio, max=1.0) # Vectorized masking: compute per-sequence thresholds batch_size, seq_len = labels.shape @@ -226,11 +232,9 @@ class CGGRUnslothBridge: mask_tokens = below_threshold & valid_mask masked_labels.masked_fill_(mask_tokens, -100) - # Update statistics (use item() only once at end) - total = valid_mask.sum().item() - kept = (masked_labels != -100).sum().item() - self.total_tokens_seen += total - self.hard_tokens_seen += kept + # Update statistics (no .item() here - keeps computation on GPU) + self.total_tokens_seen += valid_mask.sum() + self.hard_tokens_seen += (masked_labels != -100).sum() return masked_labels @@ -239,13 +243,16 @@ class CGGRUnslothBridge: self.current_step += 1 def get_stats(self) -> Dict[str, float]: - """Get CGGR statistics for logging.""" - if self.total_tokens_seen == 0: + """Get CGGR statistics for logging (syncs here).""" + total = self.total_tokens_seen.item() + if total == 0: return {"cggr/hard_ratio": 0.0, "cggr/step": self.current_step} + + hard = self.hard_tokens_seen.item() return { - "cggr/hard_ratio": self.hard_tokens_seen / self.total_tokens_seen, + "cggr/hard_ratio": hard / total, "cggr/step": self.current_step, - "cggr/total_tokens": self.total_tokens_seen, + "cggr/total_tokens": total, } @classmethod diff --git a/unsloth/cggr/router.py b/unsloth/cggr/router.py index 52c8300028..5b644f2127 100644 --- a/unsloth/cggr/router.py +++ b/unsloth/cggr/router.py @@ -116,6 +116,7 @@ class TruncatedRouter(nn.Module): self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, **kwargs, ) -> torch.Tensor: """ @@ -124,6 +125,7 @@ class TruncatedRouter(nn.Module): Args: input_ids: Input token IDs [batch, seq_len] attention_mask: Attention mask [batch, seq_len] + position_ids: Position IDs [batch, seq_len] Returns: logits: Output logits [batch, seq_len, vocab_size] @@ -131,11 +133,26 @@ class TruncatedRouter(nn.Module): # Embeddings hidden_states = self.embed_tokens(input_ids) + # Generate position_ids if not provided (needed for RoPE) + if position_ids is None: + position_ids = torch.arange( + input_ids.size(1), device=input_ids.device + ).unsqueeze(0).expand(input_ids.size(0), -1) + + # Simple expansion for 2D mask to 4D if needed by layers + mask_input = attention_mask + if attention_mask is not None and attention_mask.dim() == 2: + # Convert [batch, seq] to [batch, 1, 1, seq] + mask_input = attention_mask[:, None, None, :] + mask_input = mask_input.to(dtype=hidden_states.dtype) + mask_input = (1.0 - mask_input) * torch.finfo(hidden_states.dtype).min + # Pass through truncated layers for layer in self.layers: layer_outputs = layer( hidden_states, - attention_mask=attention_mask, + attention_mask=mask_input, + position_ids=position_ids, use_cache=False, ) hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs From 458f2a93ef1dcdf4750514925a22532e135f309a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Mar 2026 08:57:42 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks --- unsloth/__init__.py | 2 +- unsloth/cggr/__init__.py | 17 ++-- unsloth/cggr/bridge.py | 185 +++++++++++++++++++++------------------ unsloth/cggr/router.py | 80 +++++++++-------- 4 files changed, 155 insertions(+), 129 deletions(-) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index aaa4f44a20..b0da747cc2 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -283,6 +283,7 @@ from .trainer import * # CGGR (Confidence-Gated Gradient Routing) integration # Optional: requires `pip install cggr` for full functionality from .cggr import CGGR_AVAILABLE + if CGGR_AVAILABLE: from .cggr import CGGRUnslothBridge, patch_trainer_for_cggr, create_truncated_router @@ -299,4 +300,3 @@ from unsloth_zoo.rl_environments import ( # Patch TRL trainers for backwards compatibility _patch_trl_trainer() - diff --git a/unsloth/cggr/__init__.py b/unsloth/cggr/__init__.py index 8c30df714b..0268a29dc7 100644 --- a/unsloth/cggr/__init__.py +++ b/unsloth/cggr/__init__.py @@ -20,7 +20,7 @@ This module provides selective backpropagation via label masking, enabling Usage: from unsloth.cggr import CGGRUnslothBridge - + trainer = SFTTrainer(...) CGGRUnslothBridge.patch_trainer(trainer, min_tokens_ratio=0.25) trainer.train() @@ -42,6 +42,7 @@ __all__ = [ # Check if CGGR package is available try: import cggr + CGGR_AVAILABLE = True except ImportError: CGGR_AVAILABLE = False @@ -58,13 +59,17 @@ else: "CGGR is not installed. Install with: pip install cggr\n" "For CUDA acceleration: pip install cggr[cuda]" ) - + create_truncated_router = _cggr_not_available - CGGRUnslothBridge = type("CGGRUnslothBridge", (), { - "patch_trainer": staticmethod(_cggr_not_available), - }) + CGGRUnslothBridge = type( + "CGGRUnslothBridge", + (), + { + "patch_trainer": staticmethod(_cggr_not_available), + }, + ) patch_trainer_for_cggr = _cggr_not_available - + class TruncatedRouter: def __init__(self, *args, **kwargs): _cggr_not_available() diff --git a/unsloth/cggr/bridge.py b/unsloth/cggr/bridge.py index 0701cd8d6e..4b7a2e8143 100644 --- a/unsloth/cggr/bridge.py +++ b/unsloth/cggr/bridge.py @@ -35,17 +35,17 @@ __all__ = ["CGGRUnslothBridge", "patch_trainer_for_cggr"] class CGGRUnslothBridge: """ Bridge class for integrating CGGR with Unsloth trainers. - + Patches the trainer's compute_loss method to apply label masking before the forward pass, enabling selective gradient computation. - + Example: >>> from unsloth.cggr import CGGRUnslothBridge >>> trainer = SFTTrainer(...) >>> CGGRUnslothBridge.patch_trainer(trainer, min_tokens_ratio=0.25) >>> trainer.train() """ - + def __init__( self, model: torch.nn.Module, @@ -57,7 +57,7 @@ class CGGRUnslothBridge: ): """ Initialize CGGR bridge. - + Args: model: The model being trained min_tokens_ratio: Minimum fraction of tokens to keep gradients for (0.25 = top 25% hardest) @@ -72,21 +72,21 @@ class CGGRUnslothBridge: self.warmup_steps = warmup_steps self.scoring = scoring self.dynamic_threshold = dynamic_threshold - + # Create truncated router for difficulty scoring - self.router = create_truncated_router(model, num_layers=num_router_layers) - + self.router = create_truncated_router(model, num_layers = num_router_layers) + # Training state (keep on device to avoid syncs) self.current_step = 0 self.device = next(model.parameters()).device - self.total_tokens_seen = torch.tensor(0, device=self.device, dtype=torch.long) - self.hard_tokens_seen = torch.tensor(0, device=self.device, dtype=torch.long) - + self.total_tokens_seen = torch.tensor(0, device = self.device, dtype = torch.long) + self.hard_tokens_seen = torch.tensor(0, device = self.device, dtype = torch.long) + logger.info( f"Initialized CGGR Bridge: min_ratio={min_tokens_ratio}, " f"router_layers={num_router_layers}, warmup={warmup_steps}" ) - + @torch.inference_mode() def compute_difficulty_scores( self, @@ -96,29 +96,29 @@ class CGGRUnslothBridge: ) -> torch.Tensor: """ Compute difficulty scores for each token using the truncated router. - + Args: input_ids: Input token IDs [batch, seq_len] labels: Target labels [batch, seq_len] attention_mask: Attention mask [batch, seq_len] - + Returns: difficulty_scores: Per-token difficulty [batch, seq_len] """ # Get logits from truncated router (fast forward pass) - logits = self.router(input_ids, attention_mask=attention_mask) - + logits = self.router(input_ids, attention_mask = attention_mask) + # Compute difficulty based on scoring strategy if self.scoring == "entropy": # High entropy = uncertain = hard # Use log_softmax for numerical stability (single fused kernel) - log_probs = F.log_softmax(logits, dim=-1) + log_probs = F.log_softmax(logits, dim = -1) probs = log_probs.exp() - scores = -torch.sum(probs * log_probs, dim=-1) + scores = -torch.sum(probs * log_probs, dim = -1) elif self.scoring == "margin": # Small margin between top-2 = hard # topk is efficient - only partial sort needed - top2 = torch.topk(logits, k=2, dim=-1).values + top2 = torch.topk(logits, k = 2, dim = -1).values scores = -(top2[..., 0] - top2[..., 1]) # Negative margin (high = hard) elif self.scoring == "loss": # High loss = hard - directly compute per-token loss @@ -127,34 +127,36 @@ class CGGRUnslothBridge: scores = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), - reduction="none", - ignore_index=-100, + reduction = "none", + ignore_index = -100, ).view(shift_labels.shape) # Pad to match original sequence length - scores = F.pad(scores, (0, 1), value=0) + scores = F.pad(scores, (0, 1), value = 0) else: # combined - efficient fused computation # Compute log_softmax once (fused kernel) - log_probs = F.log_softmax(logits, dim=-1) + log_probs = F.log_softmax(logits, dim = -1) probs = log_probs.exp() - + # Entropy from log_probs (reuse computation) - entropy = -torch.sum(probs * log_probs, dim=-1) - + entropy = -torch.sum(probs * log_probs, dim = -1) + # Margin from topk - top2 = torch.topk(logits, k=2, dim=-1).values + top2 = torch.topk(logits, k = 2, dim = -1).values margin = top2[..., 0] - top2[..., 1] - + # Normalize and combine - use in-place operations where possible entropy_mean = entropy.mean() entropy_std = entropy.std() + 1e-10 margin_mean = margin.mean() margin_std = margin.std() + 1e-10 - + # Combined score: high entropy OR small margin = hard - scores = (entropy - entropy_mean) / entropy_std - (margin - margin_mean) / margin_std - + scores = (entropy - entropy_mean) / entropy_std - ( + margin - margin_mean + ) / margin_std + return scores - + def mask_easy_tokens( self, input_ids: torch.LongTensor, @@ -163,30 +165,32 @@ class CGGRUnslothBridge: ) -> torch.LongTensor: """ Mask easy tokens in labels with -100 to skip their gradients. - + Args: input_ids: Input token IDs [batch, seq_len] labels: Target labels [batch, seq_len] attention_mask: Attention mask [batch, seq_len] - + Returns: masked_labels: Labels with easy tokens set to -100 """ # During warmup, don't mask anything if self.current_step < self.warmup_steps: return labels - + # Clone labels to avoid modifying original masked_labels = labels.clone() - + # Compute difficulty scores scores = self.compute_difficulty_scores(input_ids, labels, attention_mask) - + # Get valid (non-ignored) token mask valid_mask = labels != -100 - + # Compute ratio (stay on GPU to avoid sync) - ratio = torch.tensor(self.min_tokens_ratio, device=self.device, dtype=scores.dtype) + ratio = torch.tensor( + self.min_tokens_ratio, device = self.device, dtype = scores.dtype + ) if self.dynamic_threshold: # More confident batch → keep fewer tokens # We use a cautious approach: only adjust if we have valid scores @@ -199,62 +203,67 @@ class CGGRUnslothBridge: mean_normalized = (valid_scores.mean() - s_min) / score_range # Lower mean score = more confident = keep fewer tokens confidence = 1.0 - mean_normalized - ratio = self.min_tokens_ratio + (1.0 - self.min_tokens_ratio) * (1.0 - confidence) * 0.5 - ratio = ratio.clamp(min=self.min_tokens_ratio, max=1.0) - + ratio = ( + self.min_tokens_ratio + + (1.0 - self.min_tokens_ratio) * (1.0 - confidence) * 0.5 + ) + ratio = ratio.clamp(min = self.min_tokens_ratio, max = 1.0) + # Vectorized masking: compute per-sequence thresholds batch_size, seq_len = labels.shape - + # Set scores of invalid tokens to -inf so they're never selected as "hard" scores_for_threshold = scores.clone() - scores_for_threshold.masked_fill_(~valid_mask, float('-inf')) - + scores_for_threshold.masked_fill_(~valid_mask, float("-inf")) + # Count valid tokens per sequence - valid_counts = valid_mask.sum(dim=1) # [batch] - + valid_counts = valid_mask.sum(dim = 1) # [batch] + # Compute number to keep per sequence - num_keep = (valid_counts.float() * ratio).long().clamp(min=1) - + num_keep = (valid_counts.float() * ratio).long().clamp(min = 1) + # For each sequence, find the threshold score (k-th largest) # Use topk to find scores we should keep max_valid = valid_counts.max().item() if max_valid > 0: # Sort scores descending to find threshold - sorted_scores, _ = scores_for_threshold.sort(dim=1, descending=True) - + sorted_scores, _ = scores_for_threshold.sort(dim = 1, descending = True) + # Get threshold for each sequence (the num_keep-th highest score) # Clamp indices to valid range - threshold_indices = (num_keep - 1).clamp(min=0, max=seq_len - 1) - thresholds = sorted_scores.gather(1, threshold_indices.unsqueeze(1)).squeeze(1) # [batch] - + threshold_indices = (num_keep - 1).clamp(min = 0, max = seq_len - 1) + thresholds = sorted_scores.gather( + 1, threshold_indices.unsqueeze(1) + ).squeeze(1) # [batch] + # Mask tokens with scores below threshold below_threshold = scores < thresholds.unsqueeze(1) mask_tokens = below_threshold & valid_mask masked_labels.masked_fill_(mask_tokens, -100) - + # Update statistics (no .item() here - keeps computation on GPU) self.total_tokens_seen += valid_mask.sum() self.hard_tokens_seen += (masked_labels != -100).sum() - + return masked_labels - + def step(self): """Called after each training step to update internal state.""" self.current_step += 1 - + def get_stats(self) -> Dict[str, float]: """Get CGGR statistics for logging (syncs here).""" total = self.total_tokens_seen.item() if total == 0: return {"cggr/hard_ratio": 0.0, "cggr/step": self.current_step} - + hard = self.hard_tokens_seen.item() return { "cggr/hard_ratio": hard / total, "cggr/step": self.current_step, "cggr/total_tokens": total, } - + @classmethod def patch_trainer( cls, @@ -267,7 +276,7 @@ class CGGRUnslothBridge: ) -> "CGGRUnslothBridge": """ Patch a trainer to use CGGR selective backpropagation. - + Args: trainer: HuggingFace/TRL trainer instance min_tokens_ratio: Minimum fraction of tokens to keep (0.25 = 25% hardest) @@ -275,10 +284,10 @@ class CGGRUnslothBridge: warmup_steps: Train normally for this many steps first scoring: Scoring strategy ('entropy', 'margin', 'loss', 'combined') dynamic_threshold: Adjust ratio based on batch confidence - + Returns: CGGRUnslothBridge instance (for accessing stats) - + Example: >>> bridge = CGGRUnslothBridge.patch_trainer(trainer) >>> trainer.train() @@ -286,45 +295,49 @@ class CGGRUnslothBridge: """ # Create bridge instance bridge = cls( - model=trainer.model, - min_tokens_ratio=min_tokens_ratio, - num_router_layers=num_router_layers, - warmup_steps=warmup_steps, - scoring=scoring, - dynamic_threshold=dynamic_threshold, + model = trainer.model, + min_tokens_ratio = min_tokens_ratio, + num_router_layers = num_router_layers, + warmup_steps = warmup_steps, + scoring = scoring, + dynamic_threshold = dynamic_threshold, ) - + # Store reference on trainer trainer._cggr_bridge = bridge - + # Patch compute_loss to apply label masking original_compute_loss = trainer.compute_loss - + @wraps(original_compute_loss) def cggr_compute_loss(model, inputs, *args, **kwargs): # Apply CGGR label masking if "labels" in inputs and inputs["labels"] is not None: inputs = dict(inputs) # Don't modify original inputs["labels"] = bridge.mask_easy_tokens( - input_ids=inputs.get("input_ids"), - labels=inputs["labels"], - attention_mask=inputs.get("attention_mask"), + input_ids = inputs.get("input_ids"), + labels = inputs["labels"], + attention_mask = inputs.get("attention_mask"), ) - + # Call original compute_loss outputs = original_compute_loss(model, inputs, *args, **kwargs) - + # Update step counter bridge.step() - + return outputs - + trainer.compute_loss = cggr_compute_loss - + print(f"🦥 Unsloth + CGGR: Selective backpropagation enabled!") - print(f" → Keeping {min_tokens_ratio*100:.0f}% hardest tokens for gradient computation") - print(f" → Router uses {num_router_layers} layers, warmup={warmup_steps} steps") - + print( + f" → Keeping {min_tokens_ratio*100:.0f}% hardest tokens for gradient computation" + ) + print( + f" → Router uses {num_router_layers} layers, warmup={warmup_steps} steps" + ) + return bridge @@ -335,19 +348,19 @@ def patch_trainer_for_cggr( ) -> CGGRUnslothBridge: """ Convenience function to patch a trainer for CGGR. - + Equivalent to CGGRUnslothBridge.patch_trainer(). - + Args: trainer: Trainer instance to patch min_tokens_ratio: Fraction of tokens to keep (0.25 = 25% hardest) **kwargs: Additional arguments passed to CGGRUnslothBridge.patch_trainer() - + Returns: CGGRUnslothBridge instance """ return CGGRUnslothBridge.patch_trainer( trainer, - min_tokens_ratio=min_tokens_ratio, + min_tokens_ratio = min_tokens_ratio, **kwargs, ) diff --git a/unsloth/cggr/router.py b/unsloth/cggr/router.py index 5b644f2127..58d23dcce9 100644 --- a/unsloth/cggr/router.py +++ b/unsloth/cggr/router.py @@ -32,37 +32,37 @@ __all__ = ["TruncatedRouter", "create_truncated_router"] class TruncatedRouter(nn.Module): """ A truncated version of a language model using only the first N layers. - + Used for fast difficulty scoring in CGGR. Shares weights with the parent model, so uses zero additional memory. - + Args: model: The parent HuggingFace model num_layers: Number of decoder layers to use (default: 2) """ - + def __init__(self, model: nn.Module, num_layers: int = 2): super().__init__() self.num_layers = num_layers - + # Get the base model (handle PEFT wrapping) base_model = model if hasattr(model, "base_model"): base_model = model.base_model if hasattr(base_model, "model"): base_model = base_model.model - + # Store reference to model components (shares weights, no copy) self.embed_tokens = self._get_embed_tokens(base_model) self.layers = self._get_layers(base_model, num_layers) self.norm = self._get_norm(base_model) self.lm_head = self._get_lm_head(model, base_model) - + # Store config for reference self.config = getattr(base_model, "config", None) self.dtype = next(model.parameters()).dtype self.device = next(model.parameters()).device - + def _get_embed_tokens(self, model: nn.Module) -> nn.Module: """Extract embedding layer from model.""" if hasattr(model, "embed_tokens"): @@ -72,7 +72,7 @@ class TruncatedRouter(nn.Module): if hasattr(model, "transformer") and hasattr(model.transformer, "wte"): return model.transformer.wte # GPT-2 style raise ValueError(f"Cannot find embedding layer in model: {type(model)}") - + def _get_layers(self, model: nn.Module, num_layers: int) -> nn.ModuleList: """Extract first N decoder layers.""" layers = None @@ -84,13 +84,13 @@ class TruncatedRouter(nn.Module): layers = model.transformer.h # GPT-2 style elif hasattr(model, "encoder") and hasattr(model.encoder, "layer"): layers = model.encoder.layer # BERT style - + if layers is None: raise ValueError(f"Cannot find decoder layers in model: {type(model)}") - + # Return reference to first N layers (shares weights) return nn.ModuleList([layers[i] for i in range(min(num_layers, len(layers)))]) - + def _get_norm(self, model: nn.Module) -> Optional[nn.Module]: """Extract final normalization layer.""" if hasattr(model, "norm"): @@ -100,17 +100,21 @@ class TruncatedRouter(nn.Module): if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"): return model.transformer.ln_f # GPT-2 style return None - - def _get_lm_head(self, original_model: nn.Module, base_model: nn.Module) -> nn.Module: + + def _get_lm_head( + self, original_model: nn.Module, base_model: nn.Module + ) -> nn.Module: """Extract language model head.""" if hasattr(original_model, "lm_head"): return original_model.lm_head if hasattr(base_model, "lm_head"): return base_model.lm_head - if hasattr(original_model, "base_model") and hasattr(original_model.base_model, "lm_head"): + if hasattr(original_model, "base_model") and hasattr( + original_model.base_model, "lm_head" + ): return original_model.base_model.lm_head raise ValueError(f"Cannot find lm_head in model: {type(original_model)}") - + @torch.inference_mode() def forward( self, @@ -121,49 +125,53 @@ class TruncatedRouter(nn.Module): ) -> torch.Tensor: """ Forward pass through truncated model to get logits for scoring. - + Args: input_ids: Input token IDs [batch, seq_len] attention_mask: Attention mask [batch, seq_len] position_ids: Position IDs [batch, seq_len] - + Returns: logits: Output logits [batch, seq_len, vocab_size] """ # Embeddings hidden_states = self.embed_tokens(input_ids) - + # Generate position_ids if not provided (needed for RoPE) if position_ids is None: - position_ids = torch.arange( - input_ids.size(1), device=input_ids.device - ).unsqueeze(0).expand(input_ids.size(0), -1) - + position_ids = ( + torch.arange(input_ids.size(1), device = input_ids.device) + .unsqueeze(0) + .expand(input_ids.size(0), -1) + ) + # Simple expansion for 2D mask to 4D if needed by layers mask_input = attention_mask if attention_mask is not None and attention_mask.dim() == 2: # Convert [batch, seq] to [batch, 1, 1, seq] mask_input = attention_mask[:, None, None, :] - mask_input = mask_input.to(dtype=hidden_states.dtype) + mask_input = mask_input.to(dtype = hidden_states.dtype) mask_input = (1.0 - mask_input) * torch.finfo(hidden_states.dtype).min - + # Pass through truncated layers for layer in self.layers: layer_outputs = layer( hidden_states, - attention_mask=mask_input, - position_ids=position_ids, - use_cache=False, + attention_mask = mask_input, + position_ids = position_ids, + use_cache = False, ) - hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs - + hidden_states = ( + layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs + ) + # Apply final norm if available if self.norm is not None: hidden_states = self.norm(hidden_states) - + # Get logits logits = self.lm_head(hidden_states) - + return logits @@ -173,23 +181,23 @@ def create_truncated_router( ) -> TruncatedRouter: """ Create a truncated router from a model for CGGR difficulty scoring. - + The router shares weights with the parent model, so uses zero additional GPU memory. It only runs the first N layers to quickly estimate token difficulty. - + Args: model: HuggingFace model (can be PEFT-wrapped) num_layers: Number of decoder layers to use (default: 2) - + Returns: TruncatedRouter instance - + Example: >>> from unsloth import FastLanguageModel >>> model, tokenizer = FastLanguageModel.from_pretrained(...) >>> router = create_truncated_router(model, num_layers=2) """ - router = TruncatedRouter(model, num_layers=num_layers) + router = TruncatedRouter(model, num_layers = num_layers) logger.info(f"Created truncated router with {num_layers} layers for CGGR scoring") return router