Upload 5 files

algorithms/backprop.py

@@ -0,0 +1,24 @@
+import torch
+from .base_optimizer import BaseOptimizer
+
+class Backpropagation(BaseOptimizer):
+    def __init__(self, model, config):
+        super().__init__(model, config)
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=float(config.get('learning_rate', 5e-5)),
+            weight_decay=float(config.get('weight_decay', 0.01))
+        )
+
+    def set_accelerator(self, accelerator):
+        super().set_accelerator(accelerator)
+        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
+
+    def step(self, inputs, labels):
+        self.model.train()
+        outputs = self.model(**inputs, labels=labels)
+        loss = outputs.loss
+        self.accelerator.backward(loss)
+        self.optimizer.step()
+        self.optimizer.zero_grad(set_to_none=True)
+        return float(loss.item())

algorithms/base_optimizer.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+
+class BaseOptimizer(ABC):
+    """Abstract base class for all training algorithms."""
+    def __init__(self, model, config):
+        self.model = model
+        self.config = config
+        self.accelerator = None
+
+    def set_accelerator(self, accelerator):
+        self.accelerator = accelerator
+
+    @abstractmethod
+    def step(self, inputs, labels):
+        """Performs a single training step; must return a Python float loss."""
+        raise NotImplementedError

algorithms/feedback_alignment.py

@@ -0,0 +1,54 @@
+import torch
+from torch import nn
+from .base_optimizer import BaseOptimizer
+
+class FeedbackAlignment(BaseOptimizer):
+    """Direct Feedback Alignment on the custom ResearchTransformer."""
+    def __init__(self, model, config):
+        super().__init__(model, config)
+        self.lr = float(config.get('learning_rate', 1e-4))
+        self.optimizer = torch.optim.AdamW([p for b in model.h for p in b.parameters()] + list(model.ln_f.parameters()) + list(model.lm_head.parameters()), lr=self.lr)
+        self.ce = nn.CrossEntropyLoss()
+        self.feedback = nn.ModuleList([
+            nn.Linear(model.config.n_embd, model.config.n_embd, bias=False) for _ in model.h
+        ])
+        for fb in self.feedback:
+            for p in fb.parameters():
+                p.requires_grad_(False)
+
+    def set_accelerator(self, accelerator):
+        super().set_accelerator(accelerator)
+        self.model, self.optimizer, self.feedback = self.accelerator.prepare(self.model, self.optimizer, self.feedback)
+
+    def step(self, inputs, labels):
+        self.model.train()
+        self.optimizer.zero_grad(set_to_none=True)
+
+        input_ids = inputs['input_ids']
+        device = input_ids.device
+        B, T = input_ids.shape
+        pos = torch.arange(0, T, device=device).unsqueeze(0)
+        x = self.model.wte(input_ids) + self.model.wpe(pos)
+        x = self.model.drop(x)
+
+        block_outs = []
+        for block in self.model.h:
+            x = block(x.detach())
+            block_outs.append(x)
+
+        x_final = self.model.ln_f(block_outs[-1])
+        logits = self.model.lm_head(x_final)
+
+        B, T, V = logits.shape
+        loss = self.ce(logits[:, :-1, :].contiguous().view(-1, V), labels[:, 1:].contiguous().view(-1))
+
+        grad_final, = torch.autograd.grad(loss, block_outs[-1], retain_graph=True)
+
+        for i in reversed(range(len(block_outs))):
+            pseudo_err = self.feedback[i](grad_final.detach())
+            block_outs[i].backward(pseudo_err, retain_graph=True)
+
+        self.accelerator.clip_grad_norm_(self.model.parameters(), 1.0)
+        self.optimizer.step()
+
+        return float(loss.item())

algorithms/forward_forward.py

@@ -0,0 +1,63 @@
+import torch
+from torch import nn
+from .base_optimizer import BaseOptimizer
+
+class ForwardForward(BaseOptimizer):
+    def __init__(self, model, config):
+        super().__init__(model, config)
+        self.threshold = float(config.get('threshold', 2.0))
+        self.lr = float(config.get('learning_rate', 1e-4))
+        self.optimizers = [torch.optim.AdamW(b.parameters(), lr=self.lr) for b in model.h]
+        self.sigmoid = nn.Sigmoid()
+
+    def set_accelerator(self, accelerator):
+        super().set_accelerator(accelerator)
+        prepared = self.accelerator.prepare(self.model, *self.optimizers)
+        self.model = prepared[0]
+        self.optimizers = list(prepared[1:])
+
+    def step(self, inputs, labels):
+        self.model.train()
+        total = 0.0
+
+        input_ids = inputs['input_ids']
+        device = input_ids.device
+        B, T = input_ids.shape
+        pos = torch.arange(0, T, device=device).unsqueeze(0)
+        x = self.model.wte(input_ids) + self.model.wpe(pos)
+        x = self.model.drop(x)
+
+        neg_ids = input_ids.clone()
+        for b in range(B):
+            idx = torch.randperm(T, device=device)
+            neg_ids[b] = neg_ids[b, idx]
+        x_neg = self.model.wte(neg_ids) + self.model.wpe(pos)
+        x_neg = self.model.drop(x_neg)
+
+        for i, block in enumerate(self.model.h):
+            opt = self.optimizers[i]
+            opt.zero_grad(set_to_none=True)
+
+            xp = x.detach()
+            xn = x_neg.detach()
+
+            op = block(xp)
+            on = block(xn)
+
+            gp = (op.pow(2).mean(dim=-1))
+            gn = (on.pow(2).mean(dim=-1))
+
+            loss = torch.log1p(torch.exp(-(gp - self.threshold))).mean() + torch.log1p(torch.exp(gn - self.threshold)).mean()
+            self.accelerator.backward(loss)
+            opt.step()
+            total += float(loss.item())
+
+            x = op.detach()
+            x_neg = on.detach()
+
+        with torch.no_grad():
+            logits = self.model.lm_head(self.model.ln_f(x))
+            V = logits.size(-1)
+            ce = nn.CrossEntropyLoss()
+            proxy = ce(logits[:, :-1, :].contiguous().view(-1, V), labels[:, 1:].contiguous().view(-1)).item()
+        return total / max(1, len(self.model.h))

algorithms/synthetic_gradients.py

@@ -0,0 +1,62 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .base_optimizer import BaseOptimizer
+
+class GradientSynthesizer(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size)
+        )
+    def forward(self, x):
+        return self.net(x.detach())
+
+class SyntheticGradients(BaseOptimizer):
+    def __init__(self, model, config):
+        super().__init__(model, config)
+        self.hidden = model.config.n_embd
+        self.main_lr = float(config.get('main_learning_rate', 1e-5))
+        self.synth_lr = float(config.get('synth_learning_rate', 1e-4))
+        self.model_opt = torch.optim.AdamW(model.parameters(), lr=self.main_lr)
+        self.synths = nn.ModuleList([GradientSynthesizer(self.hidden) for _ in model.h])
+        self.synth_opts = [torch.optim.AdamW(s.parameters(), lr=self.synth_lr) for s in self.synths]
+        self.ce = nn.CrossEntropyLoss()
+
+    def set_accelerator(self, accelerator):
+        super().set_accelerator(accelerator)
+        packs = [self.model, self.model_opt, self.synths] + self.synth_opts
+        prepped = self.accelerator.prepare(*packs)
+        self.model, self.model_opt, self.synths, *self.synth_opts = prepped
+
+    def step(self, inputs, labels):
+        self.model.train()
+        for opt in self.synth_opts:
+            opt.zero_grad(set_to_none=True)
+        self.model_opt.zero_grad(set_to_none=True)
+
+        logits, block_outs = self.model(inputs['input_ids'], return_activations=True)
+        for i, out in enumerate(block_outs):
+            pred_grad = self.synths[i](out)
+            out.backward(pred_grad, retain_graph=True)
+        self.accelerator.clip_grad_norm_(self.model.parameters(), 1.0)
+        self.model_opt.step()
+        self.model_opt.zero_grad(set_to_none=True)
+
+        logits, block_outs = self.model(inputs['input_ids'], return_activations=True)
+        B, T, V = logits.shape
+        task_loss = self.ce(logits[:, :-1, :].contiguous().view(-1, V), labels[:, 1:].contiguous().view(-1))
+        self.accelerator.backward(task_loss, retain_graph=True)
+
+        for i, out in enumerate(block_outs):
+            if out.grad is not None:
+                true_grad = out.grad.detach()
+                pred_grad = self.synths[i](out.detach())
+                synth_loss = F.mse_loss(pred_grad, true_grad)
+                self.accelerator.backward(synth_loss)
+                self.synth_opts[i].step()
+                self.synth_opts[i].zero_grad(set_to_none=True)
+
+        return float(task_loss.item())