[pre-commit.ci] auto fixes from pre-commit.com hooks

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 (

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()

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,
+    )

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