Upload via push_to_hf.py

.gitattributes

@@ -34,3 +34,40 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 arc2u-pretrain/identifiers.json filter=lfs diff=lfs merge=lfs -text
+step_110598 filter=lfs diff=lfs merge=lfs -text
+step_129031 filter=lfs diff=lfs merge=lfs -text
+step_147464 filter=lfs diff=lfs merge=lfs -text
+step_165897 filter=lfs diff=lfs merge=lfs -text
+step_18433 filter=lfs diff=lfs merge=lfs -text
+step_184330 filter=lfs diff=lfs merge=lfs -text
+step_202763 filter=lfs diff=lfs merge=lfs -text
+step_221196 filter=lfs diff=lfs merge=lfs -text
+step_239629 filter=lfs diff=lfs merge=lfs -text
+step_258062 filter=lfs diff=lfs merge=lfs -text
+step_276495 filter=lfs diff=lfs merge=lfs -text
+step_294928 filter=lfs diff=lfs merge=lfs -text
+step_313361 filter=lfs diff=lfs merge=lfs -text
+step_331794 filter=lfs diff=lfs merge=lfs -text
+step_350227 filter=lfs diff=lfs merge=lfs -text
+step_36866 filter=lfs diff=lfs merge=lfs -text
+step_368660 filter=lfs diff=lfs merge=lfs -text
+step_387093 filter=lfs diff=lfs merge=lfs -text
+step_405526 filter=lfs diff=lfs merge=lfs -text
+step_423959 filter=lfs diff=lfs merge=lfs -text
+step_442392 filter=lfs diff=lfs merge=lfs -text
+step_460825 filter=lfs diff=lfs merge=lfs -text
+step_479258 filter=lfs diff=lfs merge=lfs -text
+step_497691 filter=lfs diff=lfs merge=lfs -text
+step_516124 filter=lfs diff=lfs merge=lfs -text
+step_534557 filter=lfs diff=lfs merge=lfs -text
+step_55299 filter=lfs diff=lfs merge=lfs -text
+step_552990 filter=lfs diff=lfs merge=lfs -text
+step_571423 filter=lfs diff=lfs merge=lfs -text
+step_589856 filter=lfs diff=lfs merge=lfs -text
+step_608289 filter=lfs diff=lfs merge=lfs -text
+step_626722 filter=lfs diff=lfs merge=lfs -text
+step_645155 filter=lfs diff=lfs merge=lfs -text
+step_663588 filter=lfs diff=lfs merge=lfs -text
+step_682021 filter=lfs diff=lfs merge=lfs -text
+step_73732 filter=lfs diff=lfs merge=lfs -text
+step_92165 filter=lfs diff=lfs merge=lfs -text

all_config.yaml

@@ -0,0 +1,69 @@
+arch:
+  H_cycles: 3
+  H_layers: 0
+  L_cycles: 4
+  L_layers: 2
+  expansion: 4
+  forward_dtype: bfloat16
+  grid_noise_fraction: 0.0
+  grid_noise_prob: 0.0
+  grid_token_dropout: 0.0
+  halt_exploration_prob: 0.1
+  halt_max_steps: 16
+  halt_max_steps_eval: null
+  hidden_size: 512
+  loss:
+    loss_type: stablemax_cross_entropy
+    name: losses@ACTLossHead
+  mlp_t: false
+  name: recursive_reasoning.trm@TinyRecursiveReasoningModel_ACTV1
+  no_ACT_continue: true
+  num_heads: 8
+  pos_encodings: rope
+  puzzle_emb_dropout: 0.0
+  puzzle_emb_len: 1
+  puzzle_emb_ndim: 512
+beta1: 0.9
+beta2: 0.95
+checkpoint_every_eval: true
+checkpoint_every_n_steps: null
+checkpoint_path: checkpoints/Arc2concept-aug-1000-ACT-torch/pretrain_arc2_rearc_200k
+data_paths:
+- data/rearc-pretrain
+- data/arc2u-pretrain
+data_paths_test:
+- data/arc2u-pretrain
+dataloader_num_workers: 1
+dataloader_persistent_workers: true
+dataloader_pin_memory: true
+dataloader_prefetch_factor: 8
+ema: true
+ema_rate: 0.999
+entity: trelis
+epochs: 100000
+eval_global_batch_size: null
+eval_interval: 2500
+eval_max_augmentations: null
+eval_save_outputs: []
+evaluators:
+- name: arc@ARC
+freeze_weights: false
+freeze_weights_epochs: null
+global_batch_size: 768
+grid_noise_fraction: null
+grid_noise_prob: null
+halt_max_steps_eval: null
+load_checkpoint: checkpoints/Arc2concept-aug-1000-ACT-torch/pretrain_arc2_rearc_100k/step_737320
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+max_eval_examples_per_puzzle: null
+max_examples_per_puzzle: 5
+min_eval_interval: 0
+project_name: Arc2concept-aug-1000-ACT-torch
+puzzle_emb_lr: 0.01
+puzzle_emb_reinit_strategy: mean
+puzzle_emb_weight_decay: 0.1
+run_name: pretrain_arc2_rearc_200k
+seed: 0
+weight_decay: 0.1

losses.py

@@ -0,0 +1,103 @@
+from typing import Any, Tuple, Dict, Sequence, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+import math
+
+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, valid_mask=None):
+    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)
+
+    if valid_mask is None:
+        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],
+        # 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
+        new_carry, outputs = self.model(**model_kwargs)
+        labels = new_carry.current_data["labels"]
+
+        with torch.no_grad():
+            # Preds
+            outputs["preds"] = torch.argmax(outputs["logits"], dim=-1)
+
+            # Correctness
+            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)
+            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
+
+        lm_loss = (self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID, valid_mask=mask) / 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); Alexia: This fits Q-learning, but seems totally unecessary
+        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()
+        # Filter outputs for return
+        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}
+
+        return new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all()
+

step_110598

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

step_129031

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

step_147464

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

step_165897

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

step_18433

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

step_184330

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

step_202763

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

step_221196

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

step_239629

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

step_258062

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