tensorless 0.9.3__tar.gz → 0.9.4__tar.gz

{tensorless-0.9.3/tensorless.egg-info → tensorless-0.9.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorless
-Version: 0.9.3
+Version: 0.9.4
 Summary: ML with maximum automation and minimum setup.
 Author: Tensorless Contributors
 License: MIT

{tensorless-0.9.3 → tensorless-0.9.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorless"
-version = "0.9.3"
+version = "0.9.4"
 description = "ML with maximum automation and minimum setup."
 readme = "README.md"
 requires-python = ">=3.9"

{tensorless-0.9.3 → tensorless-0.9.4}/tensorless/backends/jax_backend.py

@@ -39,7 +39,97 @@ def is_available() -> bool:
         return False
 
 
-class JaxTinyTransformer(TinyTransformer):
+def supports_fused_training(model) -> bool:
+    """Whether `model` can use the device-resident fused AdamW/Adam train
+    step (see `_JaxFusedAdamMixin`) instead of round-tripping full
+    parameter/gradient dicts through host NumPy every step."""
+    return isinstance(model, _JaxFusedAdamMixin)
+
+
+def _adam_update_tree(jnp, params, m, v, grads, step_count, lr, weight_decay, decoupled):
+    """Fused AdamW/Adam parameter update over a pytree dict of device arrays.
+
+    Runs entirely inside `jax.jit` as part of the fused train step below, so
+    params/moments never have to leave the accelerator to be updated.
+    """
+    beta1_correction = 1.0 - 0.9 ** step_count
+    beta2_correction = 1.0 - 0.999 ** step_count
+    new_params, new_m, new_v = {}, {}, {}
+    for name, p in params.items():
+        g = grads[name]
+        if decoupled:
+            p = p * (1.0 - lr * weight_decay)
+        else:
+            g = g + weight_decay * p
+        m_i = 0.9 * m[name] + 0.1 * g
+        v_i = 0.999 * v[name] + 0.001 * g * g
+        m_hat = m_i / beta1_correction
+        v_hat = v_i / beta2_correction
+        new_params[name] = p - lr * m_hat / (jnp.sqrt(v_hat) + 1e-8)
+        new_m[name] = m_i
+        new_v[name] = v_i
+    return new_params, new_m, new_v
+
+
+class _JaxFusedAdamMixin:
+    """Keeps parameters + Adam moment estimates resident on the GPU/TPU
+    across many training steps, instead of re-uploading the full parameter
+    dict and downloading the full gradient dict through host NumPy on
+    *every single step*.
+
+    Compiling the forward pass with `jax.jit` (see `_get_forward_jit`) fixes
+    re-tracing overhead, but each `loss_and_backward()` call still does two
+    full host<->device transfers of the whole model (params up, grads down),
+    and the NumPy `Adam`/`AdamW` optimizer in `engine.py` then runs the
+    update on the CPU. For small/medium models that transfer, not compute,
+    dominates the step time and is why GPU utilization can stay low even
+    after JIT is fixed. This mixin fuses forward+backward+optimizer-update
+    into one jitted function so state only crosses the host/device boundary
+    at explicit sync points (checkpointing, validation, end of training) --
+    see `training/trainer.py`'s use of `init_fused_state`/`sync_fused_to_host`.
+    """
+
+    _fused_params = None
+    _fused_m = None
+    _fused_v = None
+    _fused_step = 0
+    _fused_train_step_key = None
+
+    def fused_training_supported(self) -> bool:
+        return True
+
+    def init_fused_state(self, weight_decay: float, decoupled: bool, step: int = 0):
+        _, jnp = _jax()
+        params = self._jax_params()
+        self._fused_params = params
+        self._fused_m = {name: jnp.zeros_like(value) for name, value in params.items()}
+        self._fused_v = {name: jnp.zeros_like(value) for name, value in params.items()}
+        self._fused_step = int(step)
+        self._fused_train_step_key = None
+
+    def load_fused_optimizer_moments(self, m_by_name, v_by_name):
+        """Restore Adam moment estimates (as NumPy, from a resumed checkpoint)."""
+        _, jnp = _jax()
+        self._fused_m = {name: jnp.asarray(value, dtype=jnp.float32) for name, value in m_by_name.items()}
+        self._fused_v = {name: jnp.asarray(value, dtype=jnp.float32) for name, value in v_by_name.items()}
+
+    def sync_fused_to_host(self):
+        """Write device-resident params/moments back to host NumPy.
+
+        Called only at checkpoint/validation/end-of-training boundaries, so
+        this cost is paid a handful of times per epoch rather than every step.
+        """
+        if self._fused_params is None:
+            return {}, {}
+        for name, parameter in self.named_parameters():
+            parameter.data[...] = np.asarray(self._fused_params[name], dtype=np.float32)
+        m_by_name = {name: np.asarray(value, dtype=np.float32) for name, value in self._fused_m.items()}
+        v_by_name = {name: np.asarray(value, dtype=np.float32) for name, value in self._fused_v.items()}
+        return m_by_name, v_by_name
+
+
+
+class JaxTinyTransformer(_JaxFusedAdamMixin, TinyTransformer):
     """The native transformer parameter layout with JAX math and gradients."""
 
     def __init__(self, *args, **kwargs):
@@ -188,16 +278,98 @@ class JaxTinyTransformer(TinyTransformer):
             parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32)
         return float(loss) / loss_scale
 
-    def generate(self, input_ids, max_new_tokens, temperature=.8, top_k=40, eos_id: Optional[int] = None):
+    def _get_fused_train_step(self, weight_decay: float, decoupled: bool):
+        cache_key = (weight_decay, decoupled)
+        if self._fused_train_step_key != cache_key:
+            jax, jnp = _jax()
+            loss_scale = 128.0 if self.precision == "fp16" else 1.0
+
+            def loss_fn(params, ids, target, mask):
+                forward = jax.checkpoint(self._forward_jax) if self.gradient_checkpointing else self._forward_jax
+                logits = forward(params, ids, mask)
+                flat_logits = logits.reshape(-1, logits.shape[-1]).astype(jnp.float32)
+                flat_target = target.reshape(-1)
+                log_probs = jax.nn.log_softmax(flat_logits, axis=-1)
+                losses = -jnp.take_along_axis(log_probs, flat_target[:, None], axis=1).squeeze(1)
+                if self.task == "text-generation":
+                    valid = flat_target != self.pad_id
+                    loss = jnp.sum(jnp.where(valid, losses, 0.0)) / jnp.maximum(valid.sum(), 1)
+                else:
+                    loss = jnp.mean(losses)
+                return loss * loss_scale
+
+            grad_fn = jax.value_and_grad(loss_fn)
+
+            def step(params, m, v, step_count, ids, target, mask, lr, grad_clip):
+                loss, grads = grad_fn(params, ids, target, mask)
+                grads = {name: g / loss_scale for name, g in grads.items()}
+                grad_norm = jnp.sqrt(sum(jnp.sum(g.astype(jnp.float32) ** 2) for g in grads.values()))
+                clip_scale = jnp.where(grad_clip > 0, jnp.minimum(1.0, grad_clip / (grad_norm + 1e-12)), 1.0)
+                grads = {name: g * clip_scale for name, g in grads.items()}
+                new_params, new_m, new_v = _adam_update_tree(
+                    jnp, params, m, v, grads, step_count, lr, weight_decay, decoupled
+                )
+                return new_params, new_m, new_v, loss / loss_scale, grad_norm
+
+            self._fused_train_step_fn = jax.jit(step)
+            self._fused_train_step_key = cache_key
+        return self._fused_train_step_fn
+
+    def fused_train_step(self, input_ids, target, attention_mask, lr, weight_decay, grad_clip, decoupled):
+        """One fused forward+backward+AdamW step, entirely on-device.
+
+        Call `init_fused_state()` once before the first step and
+        `sync_fused_to_host()` before reading `state_dict()`/`forward()` on
+        the host (e.g. for validation or checkpointing).
+        """
+        if self._fused_params is None:
+            raise RuntimeError("init_fused_state() must be called before fused_train_step().")
+        jax, jnp = _jax()
+        step_fn = self._get_fused_train_step(weight_decay, decoupled)
+        self._fused_step += 1
+        ids_arr = jnp.asarray(input_ids, dtype=jnp.int32)
+        target_arr = jnp.asarray(target, dtype=jnp.int32)
+        mask_arr = jnp.ones_like(ids_arr, dtype=jnp.float32) if attention_mask is None else jnp.asarray(attention_mask, dtype=jnp.float32)
+        new_params, new_m, new_v, loss, grad_norm = step_fn(
+            self._fused_params, self._fused_m, self._fused_v,
+            jnp.asarray(self._fused_step, dtype=jnp.float32),
+            ids_arr, target_arr, mask_arr,
+            jnp.asarray(lr, dtype=jnp.float32), jnp.asarray(grad_clip, dtype=jnp.float32),
+        )
+        self._fused_params, self._fused_m, self._fused_v = new_params, new_m, new_v
+        return float(loss), float(grad_norm)
+
+    def generate(self, input_ids, max_new_tokens, temperature=.8, top_k=40, top_p=.9,
+                 repetition_penalty=1.1, eos_id: Optional[int] = None):
+        # Mirrors `TinyTransformer.generate` (the NumPy path) exactly, so
+        # generation quality doesn't quietly regress just because a model
+        # happens to run on cuda/tpu. Sampling itself stays in NumPy since
+        # it's inherently sequential and tiny relative to the forward pass;
+        # only the forward pass benefits from JIT.
         ids = np.asarray(input_ids, dtype=np.int64).copy()
         for _ in range(max_new_tokens):
             logits = self.forward(ids[:, -self.max_seq_len:])[:, -1, :] / max(temperature, 1e-5)
+            if repetition_penalty and repetition_penalty > 1.0:
+                for row, sequence in enumerate(ids):
+                    seen = np.unique(sequence)
+                    logits[row, seen] = np.where(
+                        logits[row, seen] < 0,
+                        logits[row, seen] * repetition_penalty,
+                        logits[row, seen] / repetition_penalty,
+                    )
             if top_k is not None:
                 k = min(top_k, logits.shape[-1])
                 excluded = np.argpartition(logits, -k, axis=1)[:, :-k]
                 logits[np.arange(len(ids))[:, None], excluded] = -np.inf
             probs = np.exp(logits - logits.max(1, keepdims=True))
             probs /= probs.sum(1, keepdims=True)
+            if top_p is not None and 0 < top_p < 1:
+                order = np.argsort(-probs, axis=1)
+                sorted_probs = np.take_along_axis(probs, order, axis=1)
+                cutoff = np.cumsum(sorted_probs, axis=1) > top_p
+                cutoff[:, 0] = False
+                probs[np.arange(len(ids))[:, None], order] = np.where(cutoff, 0.0, sorted_probs)
+                probs /= probs.sum(1, keepdims=True)
             next_ids = np.array([np.random.choice(logits.shape[1], p=p) for p in probs])[:, None]
             ids = np.concatenate([ids, next_ids], axis=1)
             if eos_id is not None and np.all(next_ids == eos_id):
@@ -205,7 +377,7 @@ class JaxTinyTransformer(TinyTransformer):
         return ids
 
 
-class JaxTabularMLP(TabularMLP):
+class JaxTabularMLP(_JaxFusedAdamMixin, TabularMLP):
     """The native tabular-MLP parameter layout with JAX math and gradients.
 
     Used automatically instead of the plain NumPy `TabularMLP` whenever the
@@ -289,3 +461,58 @@ class JaxTabularMLP(TabularMLP):
         for name, parameter in self.named_parameters():
             parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32)
         return float(loss)
+
+    def _get_fused_train_step(self, weight_decay: float, decoupled: bool):
+        cache_key = (weight_decay, decoupled)
+        if self._fused_train_step_key != cache_key:
+            jax, jnp = _jax()
+            is_regression = self.task == "regression"
+
+            def loss_fn(params, numeric, categorical, dropout_masks, target_arr):
+                logits = self._forward_jax(params, numeric, categorical, dropout_masks)
+                if is_regression:
+                    error = logits - target_arr
+                    return jnp.mean(error * error)
+                log_probs = jax.nn.log_softmax(logits, axis=-1)
+                return -jnp.mean(jnp.take_along_axis(log_probs, target_arr[:, None], axis=1))
+
+            grad_fn = jax.value_and_grad(loss_fn)
+
+            def step(params, m, v, step_count, numeric, categorical, dropout_masks, target_arr, lr, grad_clip):
+                loss, grads = grad_fn(params, numeric, categorical, dropout_masks, target_arr)
+                grad_norm = jnp.sqrt(sum(jnp.sum(g.astype(jnp.float32) ** 2) for g in grads.values()))
+                clip_scale = jnp.where(grad_clip > 0, jnp.minimum(1.0, grad_clip / (grad_norm + 1e-12)), 1.0)
+                grads = {name: g * clip_scale for name, g in grads.items()}
+                new_params, new_m, new_v = _adam_update_tree(
+                    jnp, params, m, v, grads, step_count, lr, weight_decay, decoupled
+                )
+                return new_params, new_m, new_v, loss, grad_norm
+
+            self._fused_train_step_fn = jax.jit(step)
+            self._fused_train_step_key = cache_key
+        return self._fused_train_step_fn
+
+    def fused_train_step(self, numeric, categorical, target, lr, weight_decay, grad_clip, decoupled):
+        """One fused forward+backward+AdamW step, entirely on-device.
+
+        Call `init_fused_state()` once before the first step and
+        `sync_fused_to_host()` before reading `state_dict()`/`forward()` on
+        the host (e.g. for validation or checkpointing).
+        """
+        if self._fused_params is None:
+            raise RuntimeError("init_fused_state() must be called before fused_train_step().")
+        jax, jnp = _jax()
+        step_fn = self._get_fused_train_step(weight_decay, decoupled)
+        self._fused_step += 1
+        batch_size = numeric.shape[0] if self.n_numeric else categorical.shape[0]
+        dropout_masks = self._make_dropout_masks(batch_size)
+        is_regression = self.task == "regression"
+        target_arr = jnp.asarray(target, dtype=jnp.float32 if is_regression else jnp.int32)
+        new_params, new_m, new_v, loss, grad_norm = step_fn(
+            self._fused_params, self._fused_m, self._fused_v,
+            jnp.asarray(self._fused_step, dtype=jnp.float32),
+            numeric, categorical, dropout_masks, target_arr,
+            jnp.asarray(lr, dtype=jnp.float32), jnp.asarray(grad_clip, dtype=jnp.float32),
+        )
+        self._fused_params, self._fused_m, self._fused_v = new_params, new_m, new_v
+        return float(loss), float(grad_norm)

{tensorless-0.9.3 → tensorless-0.9.4}/tensorless/backends/mlx_backend.py

@@ -117,16 +117,35 @@ class MlxTinyTransformer(TinyTransformer):
             parameter.grad[...] = np.asarray(gradients[name], dtype=np.float32) / loss_scale
         return float(np.asarray(loss)) / loss_scale
 
-    def generate(self, input_ids, max_new_tokens, temperature=.8, top_k=40, eos_id: Optional[int] = None):
+    def generate(self, input_ids, max_new_tokens, temperature=.8, top_k=40, top_p=.9,
+                 repetition_penalty=1.1, eos_id: Optional[int] = None):
+        # Mirrors `TinyTransformer.generate` (the NumPy path) exactly, so
+        # generation quality doesn't quietly regress just because a model
+        # happens to run on mps.
         ids = np.asarray(input_ids, dtype=np.int64).copy()
         for _ in range(max_new_tokens):
             logits = self.forward(ids[:, -self.max_seq_len:])[:, -1, :] / max(temperature, 1e-5)
+            if repetition_penalty and repetition_penalty > 1.0:
+                for row, sequence in enumerate(ids):
+                    seen = np.unique(sequence)
+                    logits[row, seen] = np.where(
+                        logits[row, seen] < 0,
+                        logits[row, seen] * repetition_penalty,
+                        logits[row, seen] / repetition_penalty,
+                    )
             if top_k is not None:
                 k = min(top_k, logits.shape[-1])
                 excluded = np.argpartition(logits, -k, axis=1)[:, :-k]
                 logits[np.arange(len(ids))[:, None], excluded] = -np.inf
             probs = np.exp(logits - logits.max(1, keepdims=True))
             probs /= probs.sum(1, keepdims=True)
+            if top_p is not None and 0 < top_p < 1:
+                order = np.argsort(-probs, axis=1)
+                sorted_probs = np.take_along_axis(probs, order, axis=1)
+                cutoff = np.cumsum(sorted_probs, axis=1) > top_p
+                cutoff[:, 0] = False
+                probs[np.arange(len(ids))[:, None], order] = np.where(cutoff, 0.0, sorted_probs)
+                probs /= probs.sum(1, keepdims=True)
             next_ids = np.array([np.random.choice(logits.shape[1], p=p) for p in probs])[:, None]
             ids = np.concatenate([ids, next_ids], axis=1)
             if eos_id is not None and np.all(next_ids == eos_id):

{tensorless-0.9.3 → tensorless-0.9.4}/tensorless/devices/device.py

@@ -37,9 +37,33 @@ def _mps_available() -> bool:
         return False
 
 
+_BF16_CAPABLE_GPU_HINTS = (
+    "a100", "a10", "a30", "a40", "l4", "l40", "l20",
+    "h100", "h200", "h800", "b100", "b200", "gh200",
+    "rtx 30", "rtx 40", "rtx 50", "rtx a", "geforce rtx 30", "geforce rtx 40", "geforce rtx 50",
+)
+
+
 def _cuda_supports_bf16() -> bool:
-    """Use the conservative CUDA default unless hardware is identified."""
-    return False
+    """Best-effort detection of real (fast) bf16 tensor-core support.
+
+    Older GPUs (Turing/Volta/Pascal -- T4, V100, P100, K80, RTX 20xx) either
+    lack bf16 tensor cores or emulate bf16 slowly, so fp16 (with loss
+    scaling) is the safer/faster default there. Ampere-or-newer GPUs (A100,
+    L4, H100, RTX 30xx+, ...) run bf16 at full tensor-core speed and don't
+    need loss scaling, so we opt in automatically when the detected GPU name
+    matches a known Ampere-or-newer architecture. Anything unrecognized
+    conservatively falls back to fp16, matching the previous behavior.
+    """
+    try:
+        import jax
+        gpus = [d for d in jax.devices() if d.platform == "gpu"]
+        if not gpus:
+            return False
+        name = getattr(gpus[0], "device_kind", "").lower()
+        return any(hint in name for hint in _BF16_CAPABLE_GPU_HINTS)
+    except (ImportError, RuntimeError):
+        return False
 
 
 def auto_select_device(user_device: Optional[str], user_precision: Optional[str]) -> Tuple[str, str]:

{tensorless-0.9.3 → tensorless-0.9.4}/tensorless/training/trainer.py

@@ -86,6 +86,26 @@ def run_training(ds: Dataset, cfg: Dict[str, Any], checkpoint_mgr: CheckpointMan
                 "Keep the task, tokenizer, and architecture dimensions compatible."
             ) from exc
     optimizer = _build_optimizer(model, cfg)
+
+    # Fused device-resident training: for jax models trained with Adam/AdamW
+    # on a single accelerator, keep params + optimizer moments on the
+    # GPU/TPU for the whole run instead of re-uploading the full parameter
+    # dict and downloading the full gradient dict through host NumPy on
+    # every step (what the plain `Adam`/`AdamW` classes in `engine.py` force
+    # `loss_and_backward` to do). This is what actually keeps the
+    # accelerator busy once JIT tracing overhead is no longer the
+    # bottleneck. Multi-device (pmap) runs and non-Adam optimizers keep
+    # using the existing host-side path.
+    use_fused = False
+    fused_decoupled = cfg["optimizer"].lower() == "adamw"
+    if backend == "jax" and cfg["optimizer"].lower() in ("adam", "adamw"):
+        from ..backends.jax_backend import supports_fused_training
+        if supports_fused_training(model):
+            try:
+                import jax
+                use_fused = jax.local_device_count() == 1
+            except ImportError:
+                use_fused = False
     total_steps = cfg.get("max_steps") or max(1, len(prepared.train_loader)) * cfg["epochs"]
     scheduler = LambdaScheduler(optimizer, cfg["warmup_steps"], total_steps)
     early_stopper = EarlyStopping(patience=cfg["patience"], min_delta=cfg["min_delta"])
@@ -98,7 +118,30 @@ def run_training(ds: Dataset, cfg: Dict[str, Any], checkpoint_mgr: CheckpointMan
         early_stopper.best = resume_state.get("early_stopping_best", float("inf")); early_stopper.num_bad_checks = resume_state.get("early_stopping_bad_checks", 0)
         if resume_state.get("best_model_state_dict") is not None:
             best_model_state = resume_state["best_model_state_dict"]
+
+    if use_fused:
+        param_names = [name for name, _ in model.named_parameters()]
+        model.init_fused_state(cfg["weight_decay"], fused_decoupled, step=optimizer.step_count)
+        if resume_state:
+            model.load_fused_optimizer_moments(
+                dict(zip(param_names, optimizer.m)), dict(zip(param_names, optimizer.v))
+            )
+
+    def sync_fused():
+        """Bring device-resident params/moments back to host NumPy so
+        `model.state_dict()`/`model.forward()`/`optimizer.state_dict()`
+        reflect the latest step. Only called at checkpoint/validation/
+        end-of-training boundaries -- not every step."""
+        if not use_fused:
+            return
+        param_names = [name for name, _ in model.named_parameters()]
+        m_by_name, v_by_name = model.sync_fused_to_host()
+        optimizer.m = [m_by_name[name] for name in param_names]
+        optimizer.v = [v_by_name[name] for name in param_names]
+        optimizer.step_count = model._fused_step
+
     def checkpoint(epoch, complete):
+        sync_fused()
         checkpoint_mgr.save({"epoch": epoch, "global_step": global_step, "train_loader_epoch": prepared.train_loader.epoch, "model_state_dict": model.state_dict(), "best_model_state_dict": best_model_state, "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "early_stopping_best": early_stopper.best, "early_stopping_bad_checks": early_stopper.num_bad_checks, "config": cfg, "meta": prepared.meta, "tokenizer_state": prepared.tokenizer.state_dict() if prepared.tokenizer else None, "preprocessor_state": prepared.preprocessor.state_dict() if prepared.preprocessor else None, "dataset_fingerprint": dataset_fingerprint, "training_complete": complete})
     last_train_loss = last_val_loss = None
     t0 = time.time(); stop = False
@@ -109,16 +152,37 @@ def run_training(ds: Dataset, cfg: Dict[str, Any], checkpoint_mgr: CheckpointMan
         total_epoch_steps = len(prepared.train_loader)
         for batch in prepared.train_loader:
             epoch_step += 1
-            optimizer.zero_grad(); last_train_loss = _compute_loss(task, model_type, model, batch, prepared.meta.get("pad_id", 0));
-            if not np.isfinite(last_train_loss):
-                raise FloatingPointError(f"Non-finite training loss at step {global_step + 1}: {last_train_loss}")
-            if cfg["grad_clip"]:
-                norm = np.sqrt(sum(float(np.sum(p.grad.astype(np.float64) ** 2)) for p in model.parameters()))
-                if not np.isfinite(norm):
-                    raise FloatingPointError(f"Non-finite gradient norm at step {global_step + 1}: {norm}")
-                if norm > cfg["grad_clip"]:
-                    for p in model.parameters(): p.grad *= cfg["grad_clip"] / (norm + 1e-12)
-            optimizer.step(); scheduler.step(); global_step += 1
+            if use_fused:
+                grad_clip = cfg["grad_clip"] or 0.0
+                if model_type == "transformer" and task == "text-generation":
+                    last_train_loss, grad_norm = model.fused_train_step(
+                        batch[0], batch[1], None, optimizer.lr, cfg["weight_decay"], grad_clip, fused_decoupled
+                    )
+                elif model_type == "transformer":
+                    last_train_loss, grad_norm = model.fused_train_step(
+                        batch[0], batch[2], batch[1], optimizer.lr, cfg["weight_decay"], grad_clip, fused_decoupled
+                    )
+                else:
+                    numeric, categorical, target = batch
+                    last_train_loss, grad_norm = model.fused_train_step(
+                        numeric, categorical, target, optimizer.lr, cfg["weight_decay"], grad_clip, fused_decoupled
+                    )
+                if not np.isfinite(last_train_loss):
+                    raise FloatingPointError(f"Non-finite training loss at step {global_step + 1}: {last_train_loss}")
+                if grad_clip and not np.isfinite(grad_norm):
+                    raise FloatingPointError(f"Non-finite gradient norm at step {global_step + 1}: {grad_norm}")
+            else:
+                optimizer.zero_grad(); last_train_loss = _compute_loss(task, model_type, model, batch, prepared.meta.get("pad_id", 0));
+                if not np.isfinite(last_train_loss):
+                    raise FloatingPointError(f"Non-finite training loss at step {global_step + 1}: {last_train_loss}")
+                if cfg["grad_clip"]:
+                    norm = np.sqrt(sum(float(np.sum(p.grad.astype(np.float64) ** 2)) for p in model.parameters()))
+                    if not np.isfinite(norm):
+                        raise FloatingPointError(f"Non-finite gradient norm at step {global_step + 1}: {norm}")
+                    if norm > cfg["grad_clip"]:
+                        for p in model.parameters(): p.grad *= cfg["grad_clip"] / (norm + 1e-12)
+                optimizer.step()
+            scheduler.step(); global_step += 1
             if progress_enabled:
                 _show_progress(epoch + 1, cfg["epochs"], epoch_step, total_epoch_steps, last_train_loss)
             if global_step % cfg["checkpoint_every"] == 0: checkpoint(epoch, False)
@@ -128,6 +192,7 @@ def run_training(ds: Dataset, cfg: Dict[str, Any], checkpoint_mgr: CheckpointMan
         if stop:
             break
         if prepared.val_loader:
+            sync_fused()
             model.eval(); losses = [_compute_loss(task, model_type, model, batch, prepared.meta.get("pad_id", 0), False) for batch in prepared.val_loader]; last_val_loss = sum(losses) / max(1, len(losses))
             if not np.isfinite(last_val_loss):
                 raise FloatingPointError(f"Non-finite validation loss after epoch {epoch + 1}: {last_val_loss}")
@@ -146,6 +211,7 @@ def run_training(ds: Dataset, cfg: Dict[str, Any], checkpoint_mgr: CheckpointMan
                 f"train_loss={last_train_loss:.4f}"
             )
         checkpoint(epoch + 1, epoch + 1 >= cfg["epochs"])
+    sync_fused()
     elapsed = time.time() - t0
     if best_model_state is not None:
         model.load_state_dict(best_model_state)

{tensorless-0.9.3 → tensorless-0.9.4/tensorless.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorless
-Version: 0.9.3
+Version: 0.9.4
 Summary: ML with maximum automation and minimum setup.
 Author: Tensorless Contributors
 License: MIT