Full finetuning of large language models without large memory requirements
from .sigma import *
from .optim import *
from .trainer import *
from .embed import *
from .lr import CosineDecayWithWarmup

__init__.py
euclaise/SlimTrainer
import torch
from torch import nn
from torch.nn import functional as F

class QuantEmbedding(nn.Module):
    def __init__(self, old_emb: nn.Embedding):
        super().__init__()

        old_weight = old_emb.weight.data

        self.embedding_dim = old_emb.embedding_dim
        self.num_embeddings = old_emb.num_embeddings

        mins = old_weight.min(dim=1).values.view(-1, 1)
        maxs = old_weight.max(dim=1).values.view(-1, 1)

        self.scales = nn.Parameter((maxs - mins) / 255.0, requires_grad=False)
        self.means = nn.Parameter(old_weight.mean(dim=1).view(-1, 1), requires_grad=False)

        self.scales[self.scales == 0] = 1e-8

        self.weight = nn.Parameter(((old_weight - self.means) / self.scales).round().clamp(min=-128, max=127).to(torch.int8), requires_grad=False)

        self.requires_grad = False

    def forward(self, idx):
        quantized_emb = self.weight[idx].float()

        return (self.scales[idx] * (quantized_emb + self.means[idx])).bfloat16()

    def dequantize(self):
        weight = (self.weight.to("cpu").float() + self.means.cpu()) * self.scales.cpu()
        res = torch.nn.Embedding(_weight = weight, embedding_dim = self.embedding_dim, num_embeddings=self.num_embeddings, dtype=torch.float32, device="cpu")
        return res.to(self.weight.device).bfloat16()

def quantembs(module):
    for name, submodule in module.named_children():
        if isinstance(submodule, nn.Embedding):
            setattr(module, name, QuantEmbedding(submodule))
            del submodule
        else:
            quantembs(submodule)

def dequantembs(module):
    for name, submodule in module.named_children():
        if isinstance(submodule, QuantEmbedding):
            setattr(module, name, submodule.dequantize())
        else:
            dequantembs(submodule)

embed.py
euclaise/SlimTrainer
import math

class LRScheduler():
    def __init__(self, optimizer, total_steps, warmup_steps, max_lr=1e-3):
        pass
    def step(self):
        self.steps += 1

    def get_lr(self):
        pass

class CosineDecayWithWarmup(LRScheduler):
    def __init__(self, optimizer, total_steps, warmup_steps, max_lr=1e-3):
        self.total_steps = total_steps
        self.warmup_steps = warmup_steps
        self.max_lr = max_lr
        self.steps = 1

    def step(self):
        self.steps += 1

    def get_lr(self):
        if self.steps < self.warmup_steps:
            lr = self.max_lr * (self.steps / self.warmup_steps)
        else:
            remaining_steps = self.steps - self.warmup_steps
            total_cosine_steps = self.total_steps - self.warmup_steps
            cos_inner = (math.pi * remaining_steps) / total_cosine_steps
            lr = self.max_lr * (1 + math.cos(cos_inner)) / 2

        return lr

lr.py
euclaise/SlimTrainer
import torch
import math
from dataclasses import dataclass, field
from typing import List, Optional
from collections import defaultdict

# Somewhat based on https://gist.github.com/albanD/18c240bd2e09f9d93f5c4a0c9ccda39e and LOMO

@dataclass
class OverlapOptimizer:
    model: torch.nn.Module
    lr: Optional[float] = None
    decay: Optional[float] = 0.0
    _acc_grads: Optional[List] = field(default_factory=lambda: [])

    def init(self):
        for p in self.model.parameters():
            if p.requires_grad:
                self.prepare(p)
                self.hook(p)

    def step(self, loss, lr):
        pass

    def hook(self, p):
        pass

@dataclass
class OverlapLion(OverlapOptimizer):
    beta1: float = 0.9
    beta2: float = 0.99

    def prepare(self, p):
        return

    def step(self, loss, lr):
        self.lr = lr
        loss.backward()

    def prepare(self, p):
        p._m = torch.zeros_like(p)

    def hook(self, p):
        ag = p.view_as(p).grad_fn.next_functions[0][0]
        p._acc_grads = [ag]

        @torch.no_grad()
        def gf(*_):
            p.data.mul_(1 - self.lr * self.decay)

            update = p._m.clone().mul_(self.beta1).add_(p.grad, alpha=1 - self.beta1).sign_()
            p.add_(update, alpha=-self.lr)

            p._m.mul_(self.beta2).add_(g, alpha=1 - self.beta2)

            p.grad = None

        ag.register_hook(gf)

@dataclass
class OverlapSGD(OverlapOptimizer):
    sign: bool = False

    def prepare(self, p):
        return

    def step(self, loss, lr):
        self.lr = lr
        loss.backward()

    def hook(self, p):
        ag = p.view_as(p).grad_fn.next_functions[0][0]
        p._acc_grads = [ag]

        @torch.no_grad()
        def gf(*_):
            if self.sign:
                p.add_(p.grad.sign(), alpha=-self.lr)
            else:
                p.add_(p.grad, alpha=-self.lr)
            p.grad = None

        ag.register_hook(gf)

@dataclass
class Adalite(OverlapOptimizer):
    use_lr: bool = True
    eps1: float = 1e-30
    eps2: float = 1e-3
    beta2_decay: float = 0.8
    lars: bool = True
    use_rms: bool = False # RMS normalization from Adafactor - seems to take up memory but doesn't help performance
    _t: int = 0

    def step(self, loss, lr=None):
        self._t += 1
        self.lr = lr
        loss.backward()

    def prepare(self, p):
        if len(p.shape) == 2:
            n = p.shape[0]
            m = p.shape[1]

            p._r = torch.zeros(n, device=p.device, dtype=p.dtype)
            p._c = torch.zeros(m, device=p.device, dtype=p.dtype)

        else:
            p._v = torch.zeros_like(p)

    def _rms(self, x):
        return x.square().mean().sqrt()

    def hook(self, p):
        ag = p.view_as(p).grad_fn.next_functions[0][0]
        p._acc_grads = [ag]

        @torch.no_grad()
        def gf(*_):
            if self.use_lr:
                alpha = self.lr
            else:
                rho = min(1e-2, 1 / math.sqrt(self._t))
                alpha = max(self._rms(p), self.eps2)*rho

            g = p.grad

            beta2t = 1.0 - math.pow(self._t, -self.beta2_decay)
            u = g.square() + self.eps2

            if len(p.shape) == 2:
                p._r.mul_(beta2t).add_(u.mean(dim=1), alpha=1-beta2t)
                p._c.mul_(beta2t).add_(u.mean(dim=0), alpha=1-beta2t)

                r_factor = (p._r / p._r.mean(dim=-1)).rsqrt_().unsqueeze(-1)
                c_factor = p._c.unsqueeze(-2).rsqrt()
                m = r_factor @ c_factor
                m.mul_(g)
            else:
                p._v.mul_(beta2t).add_(u, alpha=1-beta2t)
                m = p._v.rsqrt() * g


            if self.use_rms:
                m.div_(max(1.0, self._rms(m)))

            p_norm = p.norm()
            g_norm = g.norm()

            if p_norm != 0 and g_norm != 0:
                m.mul_(p_norm / g_norm)

            p.add_(m, alpha=-alpha)
            p.grad = None

        ag.register_hook(gf)

optim.py
euclaise/SlimTrainer
import torch
from torch.nn import functional as F
from torch import nn

class SigmaLinear(nn.Module):
    def __init__(self, layer: nn.Linear, num_steps=1, num_steps_start=5):
        super().__init__()
        self.b = nn.Parameter(layer.bias.data) if layer.bias != None else None
        self.W = nn.Parameter(layer.weight.T.data)
        self.num_steps = num_steps

        d, c = self.W.shape

        to = {
            'device': self.W.device,
            'dtype': self.W.dtype
        }
        
        self.u = F.normalize(torch.normal(0, 1, size=(d,)), dim=0).to(**to)
        self.v = F.normalize(torch.normal(0, 1, size=(c,)), dim=0).to(**to)

        self.gamma = nn.Parameter(self._get_sigma(num_steps_start))

    def _get_sigma(self, steps):
        with torch.no_grad():
            for i in range(steps):
                self.u = F.normalize(self.W @ self.v, dim=0)
                self.v = F.normalize(self.W.T @ self.u, dim=0)
        return torch.einsum('d,dc,c->', self.u, self.W, self.v)

    def to_linear(self):
        l = nn.Linear(self.W.shape[1], self.W.shape[0], bias=(self.b is not None), device=self.W.device, dtype=self.W.dtype)
        l.weight.data = self.gamma / self._get_sigma(self.num_steps) * self.W.T
        if self.b is not None:
            l.bias.data = self.b

        return l

    def forward(self, x):
        x = x * self.gamma / self._get_sigma(self.num_steps)
        if self.b is not None:
            return x @ self.W + self.b
        else:
            return x @ self.W

def sigmafy(module):
    for name, submodule in module.named_children():
        if isinstance(submodule, nn.Linear) and "embed" not in name and "lm_head" not in name:
            setattr(module, name, SigmaLinear(submodule))
            del submodule
        else:
            sigmafy(submodule)

def desigmafy(module):
    for name, submodule in module.named_children():
        if isinstance(submodule, SigmaLinear):
            setattr(module, name, submodule.to_linear())
        else:
            desigmafy(submodule)


sigma.py
euclaise/SlimTrainer
from dataclasses import dataclass
import math
from datasets import Dataset
from typing import Callable, Mapping
from transformers import PreTrainedModel
from tqdm import trange
from tqdm.contrib import tenumerate
from torch.utils.data import DataLoader
import torch
from typing import Optional
import wandb

from .optim import OverlapOptimizer
from .lr import LRScheduler

@dataclass
class SlimTrainer():
    model: PreTrainedModel
    optim: OverlapOptimizer
    train_data: Dataset
    epochs: int
    data_collater: Callable
    batch_size: int
    scheduler: LRScheduler
    wandb_entity: Optional[str]
    wandb_project: Optional[str]
    wandb_name: Optional[str]
    report_steps: int = 8
    neft: bool = False
    freeze_embeds: bool = True

    def train(self):
        first = True
        if self.wandb_entity is not None:
            wandb.init(entity=self.wandb_entity, project=self.wandb_project, name=self.wandb_name)

        loader = DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True, collate_fn=self.data_collater)

        total_batches = len(loader)
        self.optim.init()

        if self.freeze_embeds:
            self.model.get_input_embeddings().requires_grad = False
            self.model.get_output_embeddings().requires_grad = False

        if self.neft:
            embedding_layer = self.model.get_input_embeddings()

        for epoch in trange(self.epochs, desc="Epoch"):
            if hasattr(self.scheduler, "epoch_init"):
                self.scheduler.epoch_init()

            loss_avg = 0
            for batch_idx, batch in tenumerate(loader, desc="Batch"):
                if self.neft:
                    embeds = embedding_layer(batch['input_ids'].cuda())
                    noise = (torch.rand_like(embeds) - 0.5) * 10/math.sqrt(512 * embeds.shape[-1])
                    loss = self.model(
                        inputs_embeds=embeds + noise,
                        labels=batch['labels'].cuda()
                    ).loss
                else:
                    loss = self.model(
                        input_ids=batch['input_ids'].cuda(),
                        labels=batch['labels'].cuda()
                    ).loss

                self.optim.step(loss, self.scheduler.get_lr()) # Backwards pass is mixed with optimization pass
                self.scheduler.step()
                loss_avg += loss.detach().item()


                if (batch_idx + 1) % self.report_steps == 0:
                    if self.wandb_entity is not None:
                        wandb.log({'loss': loss_avg / self.report_steps})
                        wandb.log({'epoch': epoch + batch_idx / total_batches})
                        wandb.log({'step': epoch*len(loader) + batch_idx})
                        wandb.log({'learning_rate': self.scheduler.get_lr()})
                    loss_avg = 0

            if hasattr(self.scheduler, "epoch_end"):
                self.scheduler.epoch_end()

trainer.py
euclaise/SlimTrainer