[feat] Added debug tools.

nanovllm/debug/adapters/torch_adapter.py

@@ -0,0 +1,119 @@
+"""Torch reference model adapter for breakpoint alignment."""
+
+from typing import Generator, Set, Optional
+import torch
+
+from ..breakpoints import Breakpoint, BreakpointType
+from .base import SteppableModel
+
+
+class TorchSteppable(SteppableModel):
+    """
+    Steppable adapter for the torch reference Qwen3 implementation.
+
+    Wraps tests/modeling_qwen3.py Qwen3ForCausalLM model.
+    """
+
+    def __init__(
+        self,
+        model,
+        enabled_breakpoints: Optional[Set[BreakpointType]] = None,
+    ):
+        """
+        Args:
+            model: Qwen3ForCausalLM from tests/modeling_qwen3.py
+            enabled_breakpoints: Set of breakpoint types to yield at
+        """
+        super().__init__(enabled_breakpoints)
+        self.model = model
+        self.model.eval()
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.model.model.layers)
+
+    def step(
+        self,
+        input_ids: torch.Tensor,
+        positions: Optional[torch.Tensor] = None,
+        is_prefill: bool = True,
+    ) -> Generator[Breakpoint, None, torch.Tensor]:
+        """
+        Generator that manually steps through the torch model.
+
+        The torch model uses [batch, seq_len, hidden_size] shapes.
+        """
+        # Ensure batch dimension
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)
+
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+
+        # Generate position IDs if not provided
+        if positions is None:
+            positions = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
+        elif positions.dim() == 1:
+            positions = positions.unsqueeze(0)
+
+        with torch.no_grad():
+            # === EMBEDDING ===
+            hidden_states = self.model.model.embed_tokens(input_ids)
+
+            if self.is_enabled(BreakpointType.EMBEDDING):
+                yield Breakpoint(
+                    bp_type=BreakpointType.EMBEDDING,
+                    layer_idx=None,
+                    tensor=hidden_states.detach().clone(),
+                    name="Embedding",
+                )
+
+            # Create causal attention mask
+            causal_mask = torch.triu(
+                torch.full((seq_len, seq_len), float("-inf"), device=device),
+                diagonal=1,
+            )
+            attention_mask = causal_mask.unsqueeze(0).unsqueeze(0)
+
+            # === DECODER LAYERS ===
+            for layer_idx, layer in enumerate(self.model.model.layers):
+                hidden_states, _, _ = layer(
+                    hidden_states=hidden_states,
+                    position_ids=positions,
+                    attention_mask=attention_mask,
+                    past_key_value=None,
+                    use_cache=False,
+                    output_qkv=False,
+                )
+
+                if self.is_enabled(BreakpointType.LAYER_OUTPUT):
+                    yield Breakpoint(
+                        bp_type=BreakpointType.LAYER_OUTPUT,
+                        layer_idx=layer_idx,
+                        tensor=hidden_states.detach().clone(),
+                        name=f"Layer {layer_idx}",
+                    )
+
+            # === FINAL NORM ===
+            hidden_states = self.model.model.norm(hidden_states)
+
+            if self.is_enabled(BreakpointType.FINAL_NORM):
+                yield Breakpoint(
+                    bp_type=BreakpointType.FINAL_NORM,
+                    layer_idx=None,
+                    tensor=hidden_states.detach().clone(),
+                    name="Final Norm",
+                )
+
+            # === LM HEAD ===
+            logits = self.model.lm_head(hidden_states)
+
+            if self.is_enabled(BreakpointType.LM_HEAD):
+                yield Breakpoint(
+                    bp_type=BreakpointType.LM_HEAD,
+                    layer_idx=None,
+                    tensor=logits.detach().clone(),
+                    name="LM Head",
+                )
+
+            return logits