pytensorforge 0.1.0__py3-none-any.whl

src/training/losses.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+from src.core.Tensor import Tensor
+
+
+def masked_cross_entropy(logits, targets, mask):
+    data = logits.data
+    targets = np.asarray(targets, dtype=np.int64).reshape(-1)
+    weights = np.asarray(mask, dtype=np.float32).reshape(-1)
+
+    if data.ndim != 2 or data.shape[0] != targets.shape[0] or weights.shape[0] != targets.shape[0]:
+        raise ValueError("masked_cross_entropy expects logits (N, V), targets (N,) and mask (N,)")
+
+    count = float(weights.sum())
+    denom = max(count, 1.0)
+
+    shifted = data - data.max(axis=1, keepdims=True)
+    log_norm = np.log(np.exp(shifted).sum(axis=1))
+    rows = np.arange(targets.shape[0])
+    nll = log_norm - shifted[rows, targets]
+
+    out = Tensor(np.float32(float(np.dot(nll, weights)) / denom), requires_grad=logits.requires_grad,
+                 parents=(logits,), op="MaskedCrossEntropy")
+    out.token_count = count
+
+    def _backward():
+        if not logits.requires_grad:
+            return
+
+        active = np.nonzero(weights)[0]
+
+        if active.size == 0:
+            return
+
+        probs = np.exp(shifted[active] - log_norm[active, None])
+        probs[np.arange(active.size), targets[active]] -= 1.0
+        probs *= (weights[active] * (float(out.grad) / denom))[:, None]
+        logits.grad[active] += probs
+
+    out._backward = _backward
+
+    return out

src/training/precision.py

@@ -0,0 +1,141 @@
+import math
+
+import numpy as np
+
+from src.core.Tensor import set_matmul_policy
+
+PRECISIONS = ("fp32", "bf16", "fp16")
+LOSS_SCALING = ("auto", "dynamic", "none")
+
+
+def round_bf16(x):
+    x = np.ascontiguousarray(x, dtype=np.float32)
+    bits = x.view(np.uint32)
+    finite = (bits & np.uint32(0x7F800000)) != np.uint32(0x7F800000)
+    lsb = (bits >> np.uint32(16)) & np.uint32(1)
+    rounded = (bits + np.uint32(0x7FFF) + lsb) & np.uint32(0xFFFF0000)
+    is_nan = ~finite & ((bits & np.uint32(0x007FFFFF)) != np.uint32(0))
+    special = (bits & np.uint32(0xFFFF0000)) | np.where(is_nan, np.uint32(0x00400000), np.uint32(0))
+    return np.where(finite, rounded, special).astype(np.uint32).view(np.float32)
+
+
+def round_fp16(x):
+    with np.errstate(over="ignore"):
+        return np.asarray(x, dtype=np.float32).astype(np.float16).astype(np.float32)
+
+
+ROUNDERS = {"bf16": round_bf16, "fp16": round_fp16}
+
+
+class autocast:
+
+    def __init__(self, precision):
+        if precision not in PRECISIONS:
+            raise ValueError(f"unknown precision '{precision}'; choose one of {PRECISIONS}")
+
+        self.precision = precision
+        self._previous = None
+
+    def __enter__(self):
+        self._previous = set_matmul_policy(ROUNDERS.get(self.precision))
+        return self
+
+    def __exit__(self, *exc):
+        set_matmul_policy(self._previous)
+        return False
+
+
+class GradScaler:
+
+    def __init__(
+        self,
+        enabled=True,
+        init_scale=65536.0,
+        growth_factor=2.0,
+        backoff_factor=0.5,
+        growth_interval=2000,
+        min_scale=1.0,
+        max_scale=2.0 ** 24,
+    ):
+        if init_scale <= 0 or growth_factor <= 1 or not 0 < backoff_factor < 1 or growth_interval < 1:
+            raise ValueError("invalid loss scaler settings")
+
+        self.enabled = bool(enabled)
+        self.scale = float(init_scale) if enabled else 1.0
+        self.growth_factor = float(growth_factor)
+        self.backoff_factor = float(backoff_factor)
+        self.growth_interval = int(growth_interval)
+        self.min_scale = float(min_scale)
+        self.max_scale = float(max_scale)
+        self.good_steps = 0
+        self.skipped_steps = 0
+
+    def loss_multiplier(self):
+        return self.scale if self.enabled else 1.0
+
+    def grads_finite(self, params):
+        for p in params:
+            if p.requires_grad and p._grad is not None and not np.isfinite(p._grad).all():
+                return False
+
+        return True
+
+    def unscale_and_check(self, params):
+        if not self.enabled:
+            return True
+
+        if not self.grads_finite(params):
+            return False
+
+        inv = 1.0 / self.scale
+
+        for p in params:
+            if p.requires_grad and p._grad is not None:
+                p._grad *= inv
+
+        return True
+
+    def update(self, found_overflow):
+        if not self.enabled:
+            return
+
+        if found_overflow:
+            self.skipped_steps += 1
+            self.good_steps = 0
+            self.scale = max(self.min_scale, self.scale * self.backoff_factor)
+            return
+
+        self.good_steps += 1
+
+        if self.good_steps >= self.growth_interval:
+            self.good_steps = 0
+            self.scale = min(self.max_scale, self.scale * self.growth_factor)
+
+    def state_dict(self):
+        return {
+            "enabled": self.enabled,
+            "scale": self.scale,
+            "good_steps": self.good_steps,
+            "skipped_steps": self.skipped_steps,
+        }
+
+    def load_state_dict(self, state):
+        if bool(state.get("enabled", False)) != self.enabled:
+            raise ValueError("checkpoint loss-scaling mode does not match this run")
+
+        self.scale = float(state["scale"])
+        self.good_steps = int(state["good_steps"])
+        self.skipped_steps = int(state["skipped_steps"])
+
+        if not math.isfinite(self.scale) or self.scale <= 0:
+            raise ValueError("checkpoint carries an invalid loss scale")
+
+
+def resolve_loss_scaling(precision, loss_scaling):
+    if loss_scaling not in LOSS_SCALING:
+        raise ValueError(f"unknown loss_scaling '{loss_scaling}'; choose one of {LOSS_SCALING}")
+
+    if loss_scaling == "auto":
+        return precision == "fp16"
+
+    return loss_scaling == "dynamic"

src/training/profiler.py

@@ -0,0 +1,38 @@
+import time
+from collections import defaultdict
+
+
+class Profiler:
+
+    def __init__(self):
+        self.totals = defaultdict(float)
+        self.window = defaultdict(float)
+
+    def add(self, name, seconds):
+        self.totals[name] += seconds
+        self.window[name] += seconds
+
+    def section(self, name):
+        return _Section(self, name)
+
+    def take_window(self):
+        snapshot = dict(self.window)
+        self.window.clear()
+        return snapshot
+
+    def summary(self):
+        total = sum(self.totals.values()) or 1.0
+        return {k: {"seconds": v, "share": v / total} for k, v in sorted(self.totals.items())}
+
+
+class _Section:
+
+    def __init__(self, profiler, name):
+        self.profiler = profiler
+        self.name = name
+
+    def __enter__(self):
+        self.start = time.perf_counter()
+
+    def __exit__(self, *exc):
+        self.profiler.add(self.name, time.perf_counter() - self.start)

src/training/scheduler.py

@@ -0,0 +1,50 @@
+import math
+
+
+class LRScheduler:
+
+    def __init__(
+        self,
+        base_lr,
+        warmup_steps=0,
+        decay="cosine",
+        total_steps=None,
+        min_lr=0.0,
+    ):
+        if decay not in ("cosine", "linear", "constant"):
+            raise ValueError(f"unknown decay '{decay}'")
+
+        self.base_lr = base_lr
+        self.warmup_steps = warmup_steps
+        self.decay = decay
+        self.total_steps = total_steps
+        self.min_lr = min_lr
+        self.step_count = 0
+
+    def get_lr(self, step):
+        if self.warmup_steps and step < self.warmup_steps:
+            return self.base_lr * (step + 1) / self.warmup_steps
+
+        if self.decay == "constant" or not self.total_steps:
+            return self.base_lr
+
+        span = max(1, self.total_steps - self.warmup_steps)
+        progress = (step - self.warmup_steps) / span
+        progress = min(max(progress, 0.0), 1.0)
+
+        if self.decay == "linear":
+            return self.base_lr + (self.min_lr - self.base_lr) * progress
+
+        cosine = 0.5 * (1 + math.cos(math.pi * progress))
+        return self.min_lr + (self.base_lr - self.min_lr) * cosine
+
+    def step(self):
+        lr = self.get_lr(self.step_count)
+        self.step_count += 1
+        return lr
+
+    def state_dict(self):
+        return {"step_count": self.step_count}
+
+    def load_state_dict(self, state):
+        self.step_count = state["step_count"]