diff --git a/unsloth/__init__.py b/unsloth/__init__.py index dbacd551c4..b0da747cc2 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,13 @@ 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 ( diff --git a/unsloth/cggr/__init__.py b/unsloth/cggr/__init__.py new file mode 100644 index 0000000000..0268a29dc7 --- /dev/null +++ b/unsloth/cggr/__init__.py @@ -0,0 +1,75 @@ +# 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..4b7a2e8143 --- /dev/null +++ b/unsloth/cggr/bridge.py @@ -0,0 +1,366 @@ +# 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 (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) + + 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 (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) + 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 + + # 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 (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, + 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..58d23dcce9 --- /dev/null +++ b/unsloth/cggr/router.py @@ -0,0 +1,203 @@ +# 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, + position_ids: Optional[torch.LongTensor] = 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] + 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) + ) + + # 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 = mask_input, + position_ids = position_ids, + 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