Delete shrek-tiny

shrek-tiny/all_config.yaml

@@ -1,36 +0,0 @@
-arch:
-  H_cycles: 2
-  H_layers: 2
-  L_cycles: 2
-  L_layers: 2
-  expansion: 4
-  halt_exploration_prob: 0.1
-  halt_max_steps: 16
-  hidden_size: 512
-  loss:
-    loss_type: stablemax_cross_entropy
-    name: losses@ACTLossHead
-  name: hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1
-  num_heads: 8
-  pos_encodings: rope
-  puzzle_emb_ndim: 512
-beta1: 0.9
-beta2: 0.95
-checkpoint_every_eval: true
-checkpoint_path: checkpoints/HRM_Sudoku_Comparison/SHREK_Tiny_Sudoku
-data_path: ../../dataset/data/sudoku-extreme-1k-aug-1000-hint
-ema: true
-ema_rate: 0.999
-epochs: 40000
-eval_interval: 1000
-eval_save_outputs: []
-global_batch_size: 768
-lr: 0.0001
-lr_min_ratio: 1.0
-lr_warmup_steps: 2000
-project_name: HRM_Sudoku_Comparison
-puzzle_emb_lr: 0.0001
-puzzle_emb_weight_decay: 1.0
-run_name: SHREK_Tiny_Sudoku
-seed: 0
-weight_decay: 1.0

shrek-tiny/hrm_act_v1.py

@@ -1,408 +0,0 @@
-from typing import Tuple, List, Dict, Optional
-from dataclasses import dataclass
-import math
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-from pydantic import BaseModel
-
-from models.common import trunc_normal_init_
-from models.layers import rms_norm, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear
-from models.sparse_embedding import CastedSparseEmbedding
-from models.hrm.error_singals import get_error_signal  # SHREK: error signal module
-
-
-@dataclass
-class HierarchicalReasoningModel_ACTV1InnerCarry:
-    z_H: torch.Tensor
-    z_L: torch.Tensor
-    # SHREK: prev_pred stores last step's argmax predictions for flip rate computation.
-    # zeros = fresh start (first step after init or reset gives flip_rate ≈ 1.0)
-    prev_pred: torch.Tensor        # (B, seq_len) int32
-    # SHREK: Q-values cached in carry — no longer used for Q-targets (see Bug 4),
-    # but kept because removing them causes torch.compile regression.
-    prev_q_halt: torch.Tensor      # (B,) float32
-    prev_q_continue: torch.Tensor  # (B,) float32
-
-
-@dataclass
-class HierarchicalReasoningModel_ACTV1Carry:
-    inner_carry: HierarchicalReasoningModel_ACTV1InnerCarry
-    
-    steps: torch.Tensor
-    halted: torch.Tensor
-    
-    current_data: Dict[str, torch.Tensor]
-
-
-class HierarchicalReasoningModel_ACTV1Config(BaseModel):
-    batch_size: int
-    seq_len: int
-    puzzle_emb_ndim: int = 0
-    num_puzzle_identifiers: int
-    vocab_size: int
-
-    H_cycles: int
-    L_cycles: int
-
-    H_layers: int
-    L_layers: int
-
-    # Transformer config
-    hidden_size: int
-    expansion: float
-    num_heads: int
-    pos_encodings: str
-
-    rms_norm_eps: float = 1e-5
-    rope_theta: float = 10000.0
-    
-    # Halting Q-learning config
-    halt_max_steps: int
-    halt_exploration_prob: float
-
-    # SHREK: error injection warmup — ramps alpha from 0 to alpha_max over warmup steps.
-    # Prevents small models from collapsing before the error estimator is accurate.
-    alpha_max: float = 0.01
-    alpha_warmup_steps: int = 5000
-
-    forward_dtype: str = "bfloat16"
-
-
-class HierarchicalReasoningModel_ACTV1Block(nn.Module):
-    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
-        super().__init__()
-
-        self.self_attn = Attention(
-            hidden_size=config.hidden_size,
-            head_dim=config.hidden_size // config.num_heads,
-            num_heads=config.num_heads,
-            num_key_value_heads=config.num_heads,
-            causal=False
-        )
-        self.mlp = SwiGLU(
-            hidden_size=config.hidden_size,
-            expansion=config.expansion,
-        )
-        self.norm_eps = config.rms_norm_eps
-
-    def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> torch.Tensor:
-        # Post Norm
-        # Self Attention
-        hidden_states = rms_norm(hidden_states + self.self_attn(cos_sin=cos_sin, hidden_states=hidden_states), variance_epsilon=self.norm_eps)
-        # Fully Connected
-        hidden_states = rms_norm(hidden_states + self.mlp(hidden_states), variance_epsilon=self.norm_eps)
-        return hidden_states
-
-
-class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module):
-    def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]):
-        super().__init__()
-
-        self.layers = torch.nn.ModuleList(layers)
-
-    def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor:
-        # Input injection (add)
-        hidden_states = hidden_states + input_injection
-        # Layers
-        for layer in self.layers:
-            hidden_states = layer(hidden_states=hidden_states, **kwargs)
-
-        return hidden_states
-
-
-class HierarchicalReasoningModel_ACTV1_Inner(nn.Module):
-    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
-        super().__init__()
-        self.config = config
-        self.forward_dtype = getattr(torch, self.config.forward_dtype)
-
-        # I/O
-        self.embed_scale  = math.sqrt(self.config.hidden_size)
-        embed_init_std = 1.0 / self.embed_scale
-
-        self.embed_tokens = CastedEmbedding(self.config.vocab_size, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
-        self.lm_head      = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False)
-        # Q-head: same as original HRM — reads CLS token (position 0) only
-        self.q_head       = CastedLinear(self.config.hidden_size, 2, bias=True)
-
-        # SHREK: error_encoder maps the scalar error score -> hidden_size vector for injection into z_H
-        # alpha follows a linear warmup schedule (0 → alpha_max over warmup steps).
-        # This lets the error estimator train before its signal affects z_H.
-        self.error_encoder  = nn.Linear(1, self.config.hidden_size)
-        # SHREK: step counter for alpha warmup (not a learned parameter)
-        self.register_buffer('_alpha_step', torch.tensor(0, dtype=torch.long))
-        # SHREK: error_estimator reads z_H and predicts how wrong the model is.
-        # trained via auxiliary loss in pretrain.py using the real lm_loss as target.
-        # catches "stuck but wrong" — a model confidently on the wrong answer.
-        # flip rate catches oscillation; estimator catches confident-but-wrong.
-        self.error_estimator = nn.Linear(self.config.hidden_size, 1)
-
-        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)  # ceil div
-        if self.config.puzzle_emb_ndim > 0:
-            # Zero init puzzle embeddings
-            self.puzzle_emb = CastedSparseEmbedding(self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim,
-                                                    batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype)
-
-        # LM Blocks
-        if self.config.pos_encodings == "rope":
-            self.rotary_emb = RotaryEmbedding(dim=self.config.hidden_size // self.config.num_heads,
-                                              max_position_embeddings=self.config.seq_len + self.puzzle_emb_len,
-                                              base=self.config.rope_theta)
-        elif self.config.pos_encodings == "learned":
-            self.embed_pos = CastedEmbedding(self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
-        else:
-            raise NotImplementedError()
-
-        # Reasoning Layers
-        self.H_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.H_layers)])
-        self.L_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.L_layers)])
-        
-        # Initial states
-        self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
-        self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
-
-        # Q head special init
-        # Init Q to (almost) zero for faster learning during bootstrapping
-        with torch.no_grad():
-            self.q_head.weight.zero_()
-            self.q_head.bias.fill_(-5)  # type: ignore
-
-    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor):
-        # Token embedding
-        embedding = self.embed_tokens(input.to(torch.int32))
-
-        # Puzzle embeddings
-        if self.config.puzzle_emb_ndim > 0:
-            puzzle_embedding = self.puzzle_emb(puzzle_identifiers)
-            
-            pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]
-            if pad_count > 0:
-                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))
-
-            embedding = torch.cat((puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding), dim=-2)
-
-        # Position embeddings
-        if self.config.pos_encodings == "learned":
-            # scale by 1/sqrt(2) to maintain forward variance
-            embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype))
-
-        # Scale
-        return self.embed_scale * embedding
-
-    def empty_carry(self, batch_size: int):
-        return HierarchicalReasoningModel_ACTV1InnerCarry(
-            z_H=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
-            z_L=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
-            # SHREK: zeros = no previous prediction — first step gives flip_rate ≈ 1.0
-            # device=H_init.device ensures prev_pred is on CUDA, matching z_H and logits
-            prev_pred=torch.zeros(batch_size, self.config.seq_len, dtype=torch.int32, device=self.H_init.device),
-            prev_q_halt=torch.full((batch_size,), -5.0, device=self.H_init.device),
-            prev_q_continue=torch.full((batch_size,), -5.0, device=self.H_init.device),
-        )
-        
-    def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry):
-        # SHREK: zero out prev_pred for reset sequences so they start fresh.
-        # a reset sequence is one that just halted — it will solve a new puzzle next.
-        # zeroing prev_pred means first step gives flip_rate ≈ 1.0 (maximum uncertainty).
-        new_prev_pred = torch.where(
-            reset_flag.view(-1, 1),
-            torch.zeros_like(carry.prev_pred),
-            carry.prev_pred
-        )
-        new_prev_q_halt     = torch.where(reset_flag, torch.full_like(carry.prev_q_halt,     -5.0), carry.prev_q_halt)
-        new_prev_q_continue = torch.where(reset_flag, torch.full_like(carry.prev_q_continue, -5.0), carry.prev_q_continue)
-
-        return HierarchicalReasoningModel_ACTV1InnerCarry(
-            z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H),
-            z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L),
-            prev_pred=new_prev_pred,
-            prev_q_halt=new_prev_q_halt,
-            prev_q_continue=new_prev_q_continue,
-        )
-
-
-    # SHREK: removed task_type parameter — error signal is now universal (no task rules needed)
-    def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor], require_trace=False):
-    #  -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
-        seq_info = dict(
-            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None,
-        )
-
-        # Input encoding
-        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"])
-
-        z_H_trace = []
-
-        # Forward iterations
-        with torch.no_grad():
-            z_H, z_L = carry.z_H, carry.z_L
-
-            for _H_step in range(self.config.H_cycles):
-                for _L_step in range(self.config.L_cycles):
-                    if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)):
-                        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
-
-                if not (_H_step == self.config.H_cycles - 1):
-                    z_H = self.H_level(z_H, z_L, **seq_info)
-                    if require_trace:
-                        z_H_trace.append(z_H.detach().cpu().clone())
-
-        assert not z_H.requires_grad and not z_L.requires_grad
-
-        # 1-step grad
-        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
-        z_H = self.H_level(z_H, z_L, **seq_info)
-
-        if require_trace:
-            z_H_trace.append(z_H.detach().cpu().clone())
-
-        # LM Outputs — decode z_H into token predictions
-        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]   # (B, seq_len, vocab_size)
-
-        # SHREK Component 1: Combined Error Signal
-        # Signal A — flip rate: what fraction of tokens changed from last step?
-        #   catches oscillation (model changing its mind between wrong options)
-        flip_err, current_pred = get_error_signal(output, carry.prev_pred)     # (B,), (B, seq_len)
-
-        # Signal B — learned estimator: reads z_H and predicts how wrong the model is.
-        #   catches stuck-but-wrong (model confidently on wrong answer without oscillating)
-        #   average over content positions (skip puzzle embedding prefix positions)
-        #   CRITICAL: detach z_H_mean so aux_loss only trains error_estimator weights,
-        #   NOT the main reasoning layers. Without detach, aux_loss sends parasitic
-        #   gradients through z_H → H_level/L_level that conflict with lm_loss.
-        z_H_mean    = z_H[:, self.puzzle_emb_len:].mean(dim=1).detach()        # (B, hidden_size)
-        learned_err = torch.sigmoid(self.error_estimator(z_H_mean.float()))    # (B, 1)
-        learned_err = learned_err.squeeze(-1)                                  # (B,)
-
-        # SHREK: combined error = 50/50 blend of both signals
-        # flip_err works immediately from step 1 (no learning needed)
-        # learned_err becomes accurate over training and takes over as the stronger signal
-        error = 0.5 * flip_err + 0.5 * learned_err                            # (B,)
-
-        # Q head: same as original HRM — reads CLS token (position 0) before error injection.
-        # This keeps Q-learning stable by matching HRM's gradient structure exactly.
-        q_logits  = self.q_head(z_H[:, 0].to(torch.float32)).to(torch.float32) # (B, 2)
-
-        # SHREK: inject combined error into z_H (AFTER Q-head, only affects carry for next step)
-        # error_encoder maps scalar -> hidden_size vector
-        # alpha follows linear warmup: 0 → alpha_max over warmup steps
-        # scaled by 1/sqrt(hidden_size) so injection is proportional to model size
-        error_emb = self.error_encoder(error.unsqueeze(-1))                    # (B, hidden_size)
-        # SHREK: compute alpha from warmup schedule (not learned)
-        # During warmup, alpha ramps linearly from 0 to alpha_max.
-        # After warmup, alpha stays at alpha_max.
-        # Uses torch.clamp instead of Python min() to stay compatible with torch.compile.
-        with torch.no_grad():
-            if self.training:
-                self._alpha_step += 1
-            alpha = self.config.alpha_max * torch.clamp(self._alpha_step / self.config.alpha_warmup_steps, max=1.0)
-        scale = math.sqrt(self.config.hidden_size)
-        z_H = z_H + (alpha * error_emb.unsqueeze(1) / scale).to(z_H.dtype)   # (B, seq_len, hidden_size)
-
-        # New carry: store error-injected z_H so next ACT step starts from it
-        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(
-            z_H=z_H.detach(),
-            z_L=z_L.detach(),
-            # SHREK: store current predictions — next step compares against these for flip rate
-            prev_pred=current_pred.detach(),
-            # SHREK: Q-values written to carry for torch.compile compatibility (not read for Q-targets)
-            prev_q_halt=q_logits[..., 0].detach(),
-            prev_q_continue=q_logits[..., 1].detach(),
-        )
-
-        # SHREK: also return learned_err so pretrain.py can compute auxiliary loss
-        # auxiliary loss trains the estimator: predicted error should match real lm_loss
-        if require_trace:
-            return z_H_trace, new_carry, output, (q_logits[..., 0], q_logits[..., 1]), learned_err
-        else:
-            return new_carry, output, (q_logits[..., 0], q_logits[..., 1]), learned_err
-
-
-class HierarchicalReasoningModel_ACTV1(nn.Module):
-    """ACT wrapper."""
-
-    def __init__(self, config_dict: dict):
-        super().__init__()
-        self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict)
-        self.inner = HierarchicalReasoningModel_ACTV1_Inner(self.config)
-
-    @property
-    def puzzle_emb(self):
-        return self.inner.puzzle_emb
-
-    def initial_carry(self, batch: Dict[str, torch.Tensor]):
-        batch_size = batch["inputs"].shape[0]
-
-        return HierarchicalReasoningModel_ACTV1Carry(
-            inner_carry=self.inner.empty_carry(batch_size),  # Empty is expected, it will be reseted in first pass as all sequences are halted.
-            
-            steps=torch.zeros((batch_size, ), dtype=torch.int32),
-            halted=torch.ones((batch_size, ), dtype=torch.bool),  # Default to halted
-            
-            current_data={k: torch.empty_like(v) for k, v in batch.items()}
-        )
-        
-    # SHREK: removed task_type parameter — error signal is universal, no task rules needed
-    def forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor], require_trace=False):
-        # -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor], torch.Tensor]:
-        # Update data, carry (removing halted sequences)
-        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
-
-        new_steps = torch.where(carry.halted, 0, carry.steps)
-
-        new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()}
-
-        # SHREK: run inner forward — unpack learned_err for auxiliary loss in pretrain.py
-        if require_trace:
-            z_H_trace, new_inner_carry, logits, (q_halt_logits, q_continue_logits), learned_err = self.inner(new_inner_carry, new_current_data, require_trace=require_trace)
-        else:
-            new_inner_carry, logits, (q_halt_logits, q_continue_logits), learned_err = self.inner(new_inner_carry, new_current_data)
-
-        outputs = {
-            "logits": logits,
-            "q_halt_logits": q_halt_logits,
-            "q_continue_logits": q_continue_logits,
-            "learned_err": learned_err,   # SHREK: for auxiliary loss in pretrain.py
-        }
-        
-        with torch.no_grad():
-            # Step
-            new_steps = new_steps + 1
-            is_last_step = new_steps >= self.config.halt_max_steps
-            
-            halted = is_last_step
-
-            # if training, and ACT is enabled
-            if self.training and (self.config.halt_max_steps > 1):
-                # Halt signal
-                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
-                halted = halted | (q_halt_logits > q_continue_logits)
-
-                # Exploration
-                min_halt_steps = (torch.rand_like(q_halt_logits) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
-
-                halted = halted & (new_steps >= min_halt_steps)
-
-                # SHREK: Q-target via double forward pass (same as original HRM).
-                # Run inner() a second time to get NEXT step's Q-values (step T+1).
-                # Save/restore _alpha_step so the second call doesn't double-count
-                # the warmup — without this, alpha reaches full strength at step ~2500
-                # instead of 5000, destabilizing early training.
-                # NOTE: prev_q fields in carry are not read here — they must stay
-                # in the dataclass for torch.compile compatibility.
-                saved_alpha_step = self.inner._alpha_step.clone()
-                next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-2]
-                self.inner._alpha_step.copy_(saved_alpha_step)
-
-                outputs["target_q_continue"] = torch.sigmoid(
-                    torch.where(is_last_step,
-                        next_q_halt_logits,
-                        torch.maximum(next_q_halt_logits, next_q_continue_logits))
-                )
-
-        if require_trace:
-            return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs, new_steps, (q_halt_logits > q_continue_logits), z_H_trace
-        else:
-            return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs, new_steps, (q_halt_logits > q_continue_logits)

shrek-tiny/losses.py

@@ -1,130 +0,0 @@
-from typing import Any, Tuple, Dict, Sequence, Optional
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-IGNORE_LABEL_ID = -100
-
-
-def s(x, epsilon=1e-30):
-    return torch.where(
-        x<0,
-        1/(1-x+ epsilon),
-        x + 1
-    )
-
-
-def log_stablemax(x, dim=-1):
-    s_x = s(x)
-    return torch.log(s_x/torch.sum(s_x, dim=dim, keepdim=True))
-
-
-def stablemax_cross_entropy(logits, labels, ignore_index: int = -100):
-    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)
-
-    valid_mask = labels != ignore_index
-    transformed_labels = torch.where(valid_mask, labels, 0)
-    prediction_logprobs = torch.gather(logprobs, index=transformed_labels.to(torch.long).unsqueeze(-1), dim=-1).squeeze(-1)
-
-    return -torch.where(valid_mask, prediction_logprobs, 0)
-
-
-def softmax_cross_entropy(logits, labels, ignore_index: int = -100):
-    # Cast logits to f32
-    # Flatten logits
-    return F.cross_entropy(logits.to(torch.float32).view(-1, logits.shape[-1]), labels.to(torch.long).view(-1), ignore_index=ignore_index, reduction="none").view(labels.shape)
-
-
-class ACTLossHead(nn.Module):
-    def __init__(self, model: nn.Module, loss_type: str):
-        super().__init__()
-        self.model = model
-        self.loss_fn = globals()[loss_type]
-        
-    def initial_carry(self, *args, **kwargs):
-        return self.model.initial_carry(*args, **kwargs)  # type: ignore
-
-    def forward(
-        self,
-        return_keys: Sequence[str],
-        require_trace=False,
-        # Model args
-        **model_kwargs,
-    ) -> Tuple[Any, torch.Tensor, Dict[str, torch.Tensor], Optional[Dict[str, torch.Tensor]], torch.Tensor]:
-        # Model logits
-        # B x SeqLen x D
-        if require_trace:
-            new_carry, outputs, steps, act_halt, z_H_trace = self.model(**model_kwargs, require_trace=require_trace)
-        else:
-            new_carry, outputs, steps, act_halt = self.model(**model_kwargs)
-        labels = new_carry.current_data["labels"]
-        alpha = 1.0
-        ds_mask = (alpha ** (16-steps.detach())).unsqueeze(dim=1)
-        # print(ds_mask.shape)
-
-        # Correctness
-        with torch.no_grad():
-            mask = labels != IGNORE_LABEL_ID
-            loss_counts = mask.sum(-1)
-            loss_divisor = loss_counts.clamp_min(1).unsqueeze(-1)  # Avoid NaNs in division
-
-            is_correct = mask & (torch.argmax(outputs["logits"], dim=-1) == labels)
-            seq_is_correct = is_correct.sum(-1) == loss_counts
-            
-            # # Metrics (halted)
-            # valid_metrics = new_carry.halted & (loss_counts > 0)
-            valid_metrics = (loss_counts > 0)
-
-            metrics = {
-                "count": valid_metrics.sum(),
-                
-                "accuracy":       torch.where(valid_metrics, (is_correct.to(torch.float32) / loss_divisor).sum(-1), 0).sum(),
-                "exact_accuracy": (valid_metrics & seq_is_correct).sum(),
-
-                "q_halt_accuracy": (valid_metrics & ((outputs["q_halt_logits"] >= 0) == seq_is_correct)).sum(),
-                "steps":          torch.where(valid_metrics, new_carry.steps, 0).sum(),
-            }
-
-        # Losses
-        # FIXME: Assuming the batch is always full
-        loss = self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID)
-        # print(loss.shape)
-        lm_loss = (ds_mask * loss / loss_divisor).sum()
-        q_halt_loss = F.binary_cross_entropy_with_logits(outputs["q_halt_logits"], seq_is_correct.to(outputs["q_halt_logits"].dtype), reduction="sum")
-
-        metrics.update({
-            "lm_loss": lm_loss.detach(),
-            "q_halt_loss": q_halt_loss.detach(),
-        })
-
-        # Q continue (bootstrapping target loss)
-        q_continue_loss = 0
-        if "target_q_continue" in outputs:
-            q_continue_loss = F.binary_cross_entropy_with_logits(outputs["q_continue_logits"], outputs["target_q_continue"], reduction="sum")
-
-            metrics["q_continue_loss"] = q_continue_loss.detach()
-
-        # SHREK: auxiliary loss — trains error_estimator to predict per-sample lm_loss.
-        # Computed here (not in pretrain.py) so learned_err still has gradients.
-        # per_sample_lm_loss is detached so only error_estimator weights get updated by aux_loss.
-        aux_loss = 0
-        if "learned_err" in outputs:
-            per_sample_lm_loss = (ds_mask * loss / loss_divisor).sum(dim=1).detach()  # (B,)
-            # SHREK: min-max normalize target to 0-1 so it matches sigmoid output range.
-            # When all losses are equal (early training), target = 0 (no distinguishable error).
-            # When losses vary, target spreads across 0-1 giving the estimator a useful signal.
-            lm_min = per_sample_lm_loss.min()
-            lm_max = per_sample_lm_loss.max()
-            target = (per_sample_lm_loss - lm_min) / (lm_max - lm_min + 1e-6)       # (B,) in [0, 1]
-            aux_loss = F.mse_loss(outputs["learned_err"], target, reduction="sum")
-            metrics["aux_loss"] = aux_loss.detach()
-
-        # Filter outputs for return
-        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}
-
-        if require_trace:
-            return z_H_trace, new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss) + 0.1 * aux_loss, metrics, detached_outputs, new_carry.halted.all(), new_carry.halted & act_halt
-        else:
-            return new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss) + 0.1 * aux_loss, metrics, detached_outputs, new_carry.halted.all(), new_carry.halted & act_halt

shrek-tiny/step_10416

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1f521e134348438f583491076366e9914e462af5cc2bb6163384d8c0ab911150
-size 54598517

shrek-tiny/step_11718

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:02a532206b6ffc960f1bb9f5528c7183fd81737f8ef5c88bcde8e658fe74d854
-size 54598517

shrek-tiny/step_1302

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ce1e025463aada89d20b6dd3a1f542e6979e78380664ba2a13fe97ee72398a99
-size 54598483

shrek-tiny/step_13020

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:799c12a2e32c2e9d7c50268d681e542a2de8d2c4eb1f1e4acf115b0c09180203
-size 54598517

shrek-tiny/step_14322

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1ef40dd222d3dbab83829031faf493fa025634ad339cb020d10d7361bc85133d
-size 54598517

shrek-tiny/step_15624

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7d1235e34a8b4c29bef805dfc58d933e47987e81636b7773b5ece22138c7d2cf
-size 54598517

shrek-tiny/step_16926

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:21efbca43b3499f82b875e3ba7935db426d260a957f58fb1deeb4ce200610b91
-size 54598517

shrek-tiny/step_18228

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:10ad075dbc85077383adf0fe4379631c7b06ef8d3b7f361d1ecf99937c7cc284
-size 54598517

shrek-tiny/step_19530

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9095fc8b4de55edcda4a776d18b1a91d71513ab12562856910ed170e76cefcb2
-size 54598517

shrek-tiny/step_20832

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f26e72b6a4c04cc42b3e136fce9aff806b2460e0419e3b70fc7a568575de6b8b
-size 54598517

shrek-tiny/step_22134

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:298ddccbe55ec8e8a8487c69596dea40f657262290ca7b1df4a02ad1b3455539
-size 54598517

shrek-tiny/step_23436

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e154d33d1682f65b2f42a57e4f2501e1e204ecd98d8792cc2142180340bddfa4
-size 54598517

shrek-tiny/step_24738

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5c9aa05424c2596ae35af7fabe34236ba59205c3c0629ca49eb100af90f65c9f
-size 54598517

shrek-tiny/step_2604

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8ed874d5ae52fc19167e29f37cff72c5709100ab9c8c657adf1fc45379a22abb
-size 54598483

shrek-tiny/step_26040

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:710fce62342c8db7fcd7098b2c41fe9c0b815edd5cbbc9440cab5817b6e2367c
-size 54598517

shrek-tiny/step_27342

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cf364c4d3654225412638f6110dc23274fd642050b140a1e7b5f35a2e83f95b8
-size 54598517

shrek-tiny/step_28644

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2c2727f1f7cee574e6668ad5f6159f380b5288b3eccc7b498f20df79d3ca2464
-size 54598517

shrek-tiny/step_29946

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e6558147ed93f8e9b065291c7c433152e199c2abff502c8162318434c9120131
-size 54598517

shrek-tiny/step_31248

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:66ae477da49037765ec8f256e7e29b02a14fc1468284ec0dd0393ab718bf4308
-size 54598517

shrek-tiny/step_32550

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:208ce2338a610d623aa633e7d0432873afa6214ac25d1590568b4f8e84847d79
-size 54598517

shrek-tiny/step_33852

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4af79b0c7756e82dbb6c1f332b3512778cc0349ee80651e0aa93f80dc8ca9a7d
-size 54598517

shrek-tiny/step_35154

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8af708d3e7a5aea9a6fac5fb82270435f404e6fc3e99ee54e6c9c5a0ea251d25
-size 54598517

shrek-tiny/step_36456

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:850ec865d80577580e98d0ced4c40f276a45d8e895a9d21e1492787153b825a9
-size 54598517

shrek-tiny/step_37758

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8749bd18d7b192c11aa9258147fe27ce140330307c5364171513e3fef06ce815
-size 54598517

shrek-tiny/step_3906

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2cfe4986310699fee8b4c7aac140de82ea37c48b076e1aaa1aa8e027d1f0ef4f
-size 54598483

shrek-tiny/step_39060

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3a067f3a96e024dc649799c164486c3996c490a47217297ff2f31532c22a7a28
-size 54598517

shrek-tiny/step_40362

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a60abecd938f8d59d53e5b46c919f9fa0b465c67e5acc2491a545d344ea4706f
-size 54598517

shrek-tiny/step_41664

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:57111671199da5fbdf860b29814aa2566db67d43cc0854e87c89b7fef25074c8
-size 54598517

shrek-tiny/step_42966

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3f8d01a695e277c7e3cb58ce38fe5fe227c1f9619b45d8e45b9cd90bffa2b2d0
-size 54598517

shrek-tiny/step_44268

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b840f9100ec94f566ce6bbd6958bf8d5460a5142d63a365a09e1a04fc543b04b
-size 54598517

shrek-tiny/step_45570

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3626f3ab51661725c17ec1b17d81db438f7bbcfff0361c0b8d1e8c7b7970db9d
-size 54598517

shrek-tiny/step_46872

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8c412c892a1f174f665c3be0ba30378633950c190ab0a656efde1a8587930dfc
-size 54598517

shrek-tiny/step_48174

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0e6ec2a15dfa10e086327bc1678f1bfa8dcd60c9aec22c39abc37d21ef97e1aa
-size 54598517

shrek-tiny/step_49476

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8e30617f346f39194e04d1831bd88ae6808b76bc4520d87518a426bdfb8f6824
-size 54598517

shrek-tiny/step_50778

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0efb27ceb2cb2df99a87ec0da2db9b7333b8120a7730c3eae51067437cc6844d
-size 54598517

shrek-tiny/step_5208

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:344ee11d7b0e226dfc90c9b889aed3b19b807226d674be0dfb9589b101ae1e52
-size 54598483

shrek-tiny/step_52080

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a4ff48f5373ed2b36a0ae1c8ca96813779180c839674976b01334d59b1c70643
-size 54598517

shrek-tiny/step_6510

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:68ca57695c3ec6ad31c599b44e722a5bf5da7a738d550b67bcb3842982187c7a
-size 54598483

shrek-tiny/step_7812

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8f0273ae881fd2b62660ebd03f75e877b9fc03c99f661234d36ad9de711abe13
-size 54598483

shrek-tiny/step_9114

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1a0e6fa1f81c4d918854e80d1f44ccd482dc4cdb3b90412dd9ad96cc61660356
-size 54598483