pytensorforge 0.1.0__py3-none-any.whl

src/models/gpt/block.py

@@ -0,0 +1,74 @@
+from src.activations import activation_fns
+from src.models.gpt.attention import CausalSelfAttention
+from src.models.transformers.Dropout import Dropout
+from src.models.transformers.LayerNorm import LayerNorm
+from src.neural.Dense import Dense
+from src.neural.Layer import Layer
+
+
+class GPTBlock(Layer):
+
+    def __init__(self, config, rope=None):
+        super().__init__()
+
+        self.config = config
+
+        self.norm1 = LayerNorm(config.d_model, eps=config.norm_eps)
+        self.attn = CausalSelfAttention(config.d_model, config.n_heads, rope=rope)
+
+        self.norm2 = LayerNorm(config.d_model, eps=config.norm_eps)
+        self.fc1 = Dense(config.ff_dim)
+        self.fc2 = Dense(config.d_model)
+
+        self.dropout = Dropout(config.dropout)
+        self.activation = activation_fns[config.activation]
+
+    def build(self, input_shape):
+        self.norm1.build(input_shape)
+        self.attn.build(input_shape)
+
+        self.norm2.build(input_shape)
+        self.fc1.build(input_shape)
+
+        ff_shape = (*input_shape[:-1], self.fc1.units)
+        self.fc2.build(ff_shape)
+
+        self.dropout.build(input_shape)
+
+        self.built = True
+
+    def call(self, x, mask=None):
+        h = self.norm1.call(x)
+        h = self.attn.call(h, mask=mask)
+        x = x + self.dropout.call(h)
+
+        h = self.norm2.call(x)
+        h = self.fc2.call(self.activation(self.fc1.call(h)))
+        x = x + self.dropout.call(h)
+
+        return x
+
+    def parameters(self):
+        params = []
+        params.extend(self.norm1.parameters())
+        params.extend(self.attn.parameters())
+        params.extend(self.norm2.parameters())
+        params.extend(self.fc1.parameters())
+        params.extend(self.fc2.parameters())
+        return params
+
+    def state_dict(self):
+        return {
+            "norm1": self.norm1.state_dict(),
+            "attn": self.attn.state_dict(),
+            "norm2": self.norm2.state_dict(),
+            "fc1": self.fc1.state_dict(),
+            "fc2": self.fc2.state_dict(),
+        }
+
+    def load_state_dict(self, state):
+        self.norm1.load_state_dict(state["norm1"])
+        self.attn.load_state_dict(state["attn"])
+        self.norm2.load_state_dict(state["norm2"])
+        self.fc1.load_state_dict(state["fc1"])
+        self.fc2.load_state_dict(state["fc2"])

src/models/gpt/config.py

@@ -0,0 +1,103 @@
+import json
+from typing import Optional
+from dataclasses import asdict, dataclass
+
+
+SUPPORTED_ACTIVATIONS = ("gelu", "relu", "tanh", "sigmoid")
+POSITION_ENCODINGS = ("learned", "rope")
+
+
+@dataclass
+class GPTConfig:
+    vocab_size: int
+    context_length: int = 1024
+    d_model: int = 768
+    n_layers: int = 12
+    n_heads: int = 12
+    ff_dim: int = 3072
+    activation: str = "gelu"
+    dropout: float = 0.0
+    norm_eps: float = 1e-5
+    tie_weights: bool = True
+    architecture: str = "gpt"
+    position_encoding: str = "learned"
+    rope_theta: float = 10000.0
+    rope_scaling: str = "none"
+    rope_scaling_factor: float = 1.0
+    trained_context_length: Optional[int] = None
+
+    def __post_init__(self):
+        if self.d_model % self.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+
+        if self.activation not in SUPPORTED_ACTIVATIONS:
+            raise ValueError(
+                f"unsupported activation '{self.activation}'; choose one of {SUPPORTED_ACTIVATIONS}"
+            )
+
+        if self.position_encoding not in POSITION_ENCODINGS:
+            raise ValueError(
+                f"unsupported position_encoding '{self.position_encoding}'; choose one of {POSITION_ENCODINGS}"
+            )
+
+        if self.position_encoding == "rope":
+            if (self.d_model // self.n_heads) % 2 != 0:
+                raise ValueError("rope needs an even head dimension (d_model / n_heads)")
+
+            if not self.rope_theta > 0:
+                raise ValueError("rope_theta must be positive")
+
+        if self.rope_scaling not in ("none", "linear", "ntk"):
+            raise ValueError(f"unknown rope_scaling '{self.rope_scaling}'; choose none, linear or ntk")
+
+        self.rope_scaling_factor = float(self.rope_scaling_factor)
+
+        if not self.rope_scaling_factor >= 1.0:
+            raise ValueError("rope_scaling_factor must be >= 1")
+
+        if (self.rope_scaling != "none" or self.trained_context_length is not None) and not self.uses_rope:
+            raise ValueError("context extension settings require position_encoding: rope")
+
+        if self.trained_context_length is not None and self.trained_context_length > self.context_length:
+            raise ValueError("trained_context_length cannot exceed context_length")
+
+        self.rope_theta = float(self.rope_theta)
+
+    @property
+    def head_dim(self):
+        return self.d_model // self.n_heads
+
+    @property
+    def uses_rope(self):
+        return self.position_encoding == "rope"
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data):
+        known = set(cls.__dataclass_fields__.keys())
+        return cls(**{k: v for k, v in data.items() if k in known})
+
+    @property
+    def base_context_length(self):
+        return self.trained_context_length or self.context_length
+
+    def rotary_tables(self):
+        from src.models.gpt.rope import rotary_tables
+
+        return rotary_tables(self.head_dim, self.context_length, self.rope_theta,
+                             self.rope_scaling, self.rope_scaling_factor)
+
+    @classmethod
+    def normalized(cls, data):
+        return cls.from_dict(data).to_dict()
+
+    def save(self, path):
+        with open(path, "w") as f:
+            json.dump(self.to_dict(), f, indent=2, sort_keys=True)
+
+    @classmethod
+    def load(cls, path):
+        with open(path) as f:
+            return cls.from_dict(json.load(f))

src/models/gpt/context.py

@@ -0,0 +1,44 @@
+from dataclasses import replace
+
+EXTENSION_METHODS = ("extrapolate", "linear", "ntk")
+
+
+def extend_context(config, context_length, method):
+    if method not in EXTENSION_METHODS:
+        raise ValueError(f"unknown context extension method '{method}'; choose one of {EXTENSION_METHODS}")
+
+    if not config.uses_rope:
+        raise ValueError(
+            "context extension needs a rope model; learned positional embeddings have no entries past "
+            f"the trained length of {config.context_length}"
+        )
+
+    base = config.base_context_length
+    context_length = int(context_length)
+
+    if context_length <= base:
+        raise ValueError(f"extended context ({context_length}) must exceed the trained context ({base})")
+
+    factor = context_length / base
+    scaling = "none" if method == "extrapolate" else method
+
+    return replace(
+        config,
+        context_length=context_length,
+        trained_context_length=base,
+        rope_scaling=scaling,
+        rope_scaling_factor=1.0 if scaling == "none" else factor,
+    )
+
+
+def describe_context(config):
+    if config.trained_context_length is None:
+        return {"context_length": config.context_length, "extended": False}
+
+    return {
+        "context_length": config.context_length,
+        "extended": True,
+        "trained_context_length": config.trained_context_length,
+        "method": "extrapolate" if config.rope_scaling == "none" else config.rope_scaling,
+        "factor": config.context_length / config.trained_context_length,
+    }

src/models/gpt/model.py

@@ -0,0 +1,165 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+from src.models.embedding.Embedding import Embedding
+from src.models.gpt.block import GPTBlock
+from src.models.gpt.recompute import checkpoint
+from src.models.transformers.LayerNorm import LayerNorm
+from src.neural.Dense import Dense
+
+
+class GPTModel:
+
+    def __init__(self, config):
+        self.config = config
+
+        self.token_emb = Embedding(config.vocab_size, config.d_model)
+
+        self.pos_emb = None
+        self.rope = None
+
+        if config.uses_rope:
+            self.rope = config.rotary_tables()
+        else:
+            self.pos_emb = Embedding(config.context_length, config.d_model)
+
+        self.blocks = [GPTBlock(config, rope=self.rope) for _ in range(config.n_layers)]
+
+        self.final_norm = LayerNorm(config.d_model, eps=config.norm_eps)
+
+        self.lm_head = None
+        if not config.tie_weights:
+            self.lm_head = Dense(config.vocab_size)
+
+        self.activation_checkpointing = False
+        self.training = True
+        self._built = False
+
+    def build(self):
+        if self._built:
+            return
+
+        self.token_emb.build((None, self.config.context_length))
+
+        if self.pos_emb is not None:
+            self.pos_emb.build((None, self.config.context_length))
+
+        block_shape = (None, self.config.context_length, self.config.d_model)
+
+        for block in self.blocks:
+            block.build(block_shape)
+
+        self.final_norm.build(block_shape)
+
+        if self.lm_head is not None:
+            self.lm_head.build(block_shape)
+
+        self._built = True
+
+    def forward(self, input_ids):
+        if not self._built:
+            self.build()
+
+        batch, seq_len = input_ids.shape
+
+        if seq_len > self.config.context_length:
+            raise ValueError(
+                f"sequence length {seq_len} exceeds context_length "
+                f"{self.config.context_length}"
+            )
+
+        x = self.token_emb.call(input_ids)
+
+        if self.pos_emb is not None:
+            positions = Tensor(
+                np.tile(np.arange(seq_len, dtype=np.float32), (batch, 1)),
+                requires_grad=False,
+            )
+            x = x + self.pos_emb.call(positions)
+
+        for block in self.blocks:
+            if self.activation_checkpointing:
+                x = checkpoint(block.call, x)
+            else:
+                x = block.call(x)
+
+        x = self.final_norm.call(x)
+
+        if self.lm_head is not None:
+            logits = self.lm_head.call(x)
+        else:
+            weight = self.token_emb.weight
+            flat = x.reshape(batch * seq_len, self.config.d_model)
+            logits = flat @ weight.transpose(1, 0)
+            logits = logits.reshape(batch, seq_len, self.config.vocab_size)
+
+        return logits
+
+    def __call__(self, input_ids):
+        return self.forward(input_ids)
+
+    def set_training(self, flag):
+        previous = self.training
+        self.training = bool(flag)
+
+        for block in self.blocks:
+            block.dropout.training = self.training
+
+        return previous
+
+    def parameters(self):
+        if not self._built:
+            self.build()
+
+        params = []
+        params.extend(self.token_emb.parameters())
+
+        if self.pos_emb is not None:
+            params.extend(self.pos_emb.parameters())
+
+        for block in self.blocks:
+            params.extend(block.parameters())
+
+        params.extend(self.final_norm.parameters())
+
+        if self.lm_head is not None:
+            params.extend(self.lm_head.parameters())
+
+        return params
+
+    def num_parameters(self):
+        return sum(p.data.size for p in self.parameters())
+
+    def state_dict(self):
+        state = {
+            "token_emb": self.token_emb.state_dict(),
+            "final_norm": self.final_norm.state_dict(),
+            "blocks": [b.state_dict() for b in self.blocks],
+        }
+
+        if self.pos_emb is not None:
+            state["pos_emb"] = self.pos_emb.state_dict()
+
+        if self.lm_head is not None:
+            state["lm_head"] = self.lm_head.state_dict()
+
+        return state
+
+    def load_state_dict(self, state):
+        if not self._built:
+            self.build()
+
+        if (self.pos_emb is not None) != ("pos_emb" in state):
+            raise ValueError("state does not match the model's position encoding")
+
+        self.token_emb.load_state_dict(state["token_emb"])
+
+        if self.pos_emb is not None:
+            self.pos_emb.load_state_dict(state["pos_emb"])
+        self.final_norm.load_state_dict(state["final_norm"])
+
+        for block, block_state in zip(self.blocks, state["blocks"]):
+            block.load_state_dict(block_state)
+
+        if self.lm_head is not None and "lm_head" in state:
+            self.lm_head.load_state_dict(state["lm_head"])

src/models/gpt/recompute.py

@@ -0,0 +1,35 @@
+import numpy as np
+
+from src.core.Tensor import Tensor, grad_enabled, no_grad
+
+
+def checkpoint(fn, x):
+    if not grad_enabled():
+        return fn(x)
+
+    rng_state = np.random.get_state()
+
+    with no_grad():
+        y_data = fn(Tensor(x.data, requires_grad=False)).data
+
+    out = Tensor(y_data, requires_grad=True, parents=(x,), op="Checkpoint")
+
+    def _backward():
+        resume_state = np.random.get_state()
+        np.random.set_state(rng_state)
+
+        try:
+            x_local = Tensor(x.data, requires_grad=True)
+            y = fn(x_local)
+        finally:
+            np.random.set_state(resume_state)
+
+        y.backward(out.grad, release=True)
+
+        if x.requires_grad:
+            x.grad += x_local.grad
+
+    out._backward = _backward
+
+    return out
+

src/models/gpt/rope.py

@@ -0,0 +1,84 @@
+import threading
+
+import numpy as np
+
+
+ROPE_SCALING = ("none", "linear", "ntk")
+
+
+def scaled_theta(theta, head_dim, scaling, factor):
+    if scaling == "ntk":
+        return float(theta) * float(factor) ** (head_dim / (head_dim - 2))
+
+    return float(theta)
+
+
+class RotaryTables:
+
+    def __init__(self, head_dim, max_positions, theta=10000.0, scaling="none", factor=1.0):
+        if head_dim % 2 != 0:
+            raise ValueError("rotary position encoding needs an even head dimension")
+
+        if scaling not in ROPE_SCALING:
+            raise ValueError(f"unknown rope scaling '{scaling}'; choose one of {ROPE_SCALING}")
+
+        if not factor >= 1.0:
+            raise ValueError("rope scaling factor must be >= 1")
+
+        self.head_dim = head_dim
+        self.half = head_dim // 2
+        self.scaling = scaling
+        self.factor = float(factor)
+        self.theta = scaled_theta(theta, head_dim, scaling, factor)
+        self.max_positions = int(max_positions)
+
+        inv_freq = self.theta ** (-np.arange(0, self.half, dtype=np.float64) / self.half)
+        positions = np.arange(self.max_positions, dtype=np.float64)
+
+        if scaling == "linear":
+            positions = positions / self.factor
+
+        angles = positions[:, None] * inv_freq[None, :]
+
+        self.cos = np.cos(angles).astype(np.float32)
+        self.sin = np.sin(angles).astype(np.float32)
+
+    def tables(self, positions):
+        positions = np.asarray(positions, dtype=np.int64)
+
+        if positions.size and (positions.min() < 0 or positions.max() >= self.max_positions):
+            raise ValueError(f"position outside the rotary table (0..{self.max_positions - 1})")
+
+        return self.cos[positions], self.sin[positions]
+
+    def rotate(self, x, cos, sin, inverse=False):
+        h = self.half
+        x1 = x[..., :h]
+        x2 = x[..., h:]
+        out = np.empty_like(x)
+
+        if inverse:
+            out[..., :h] = x1 * cos + x2 * sin
+            out[..., h:] = x2 * cos - x1 * sin
+        else:
+            out[..., :h] = x1 * cos - x2 * sin
+            out[..., h:] = x2 * cos + x1 * sin
+
+        return out
+
+
+_cache = {}
+_cache_lock = threading.Lock()
+
+
+def rotary_tables(head_dim, max_positions, theta=10000.0, scaling="none", factor=1.0):
+    key = (int(head_dim), int(max_positions), float(theta), scaling, float(factor))
+
+    with _cache_lock:
+        tables = _cache.get(key)
+
+        if tables is None:
+            tables = RotaryTables(head_dim, max_positions, theta, scaling, factor)
+            _cache[key] = tables
+
+        return tables

src/models/regression/Linear.py

@@ -0,0 +1,51 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+# y = XW + B
+class Linear:
+    def __init__(self, weight_shape, bias_shape):
+
+        self.W = Tensor.random(weight_shape)
+
+        self.b = Tensor.random(bias_shape)
+
+        self.dW = Tensor.zeros(weight_shape)
+        self.db = Tensor.zeros(bias_shape)
+
+        self.X = None
+
+    def normalize(self, X):
+        norm = np.linalg.norm(X)
+        return X if norm == 0 else X / norm
+
+    def forward(self, X):
+        self.X = X
+        return np.dot(X, self.W.data) + self.b.data
+
+    def grads(self):
+        return [self.dW, self.db]
+
+    def parameters(self):
+        return [self.W, self.b]
+
+    def backward(self, grad_output):
+        # L = pred - target
+        # u = pred - target
+        # L = u
+        # dL/du = 1
+        # dL/dpred = dL/du * du/dpred
+        # du/dpred = 1 - 0 = 1
+        # dL/dpred = u = pred - target
+        # how do W affect the L?
+        # dL/dW = dL/dpred * dpred/dW
+        # dpred/dW = X
+        # dL/dW = (pred - target) * X
+        self.dW.data += grad_output * self.X.data
+        self.db.data += grad_output
+        return grad_output * self.W.data
+
+    def update(self, learning_rate):
+        self.W.data -= learning_rate * self.dW.data
+        self.b.data -= learning_rate * self.db.data

src/models/regression/Logistic.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+from src.math.sigmoid import sigmoid
+
+
+class Logistic:
+    def __init__(self, weight_shape, bias_shape):
+        self.X = None
+        self.W = Tensor(weight_shape)
+        self.W.random()
+
+        self.b = Tensor(bias_shape)
+        self.b.random()
+
+        self.dW = Tensor(weight_shape)
+        self.db = Tensor(bias_shape)
+
+    def grads(self):
+        return [self.dW, self.db]
+
+    def parameters(self):
+        return [self.W, self.b]
+
+    def forward(self, X):
+        self.X = X
+        return sigmoid(np.dot(X, self.W.data) + self.b.data)
+
+    def backward(self, grad_output):
+        self.dW.data += grad_output * self.X.data
+        self.db.data += grad_output
+        return grad_output * self.W.data
+
+    def update(self, learning_rate):
+        self.W.data -= learning_rate * self.dW.data
+        self.b.data -= learning_rate * self.db.data