Improve load speed of the loss log

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -31,7 +31,11 @@ from torch.utils.checkpoint import checkpoint
 
 
 def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor:
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    # MPS has no float64
+    dt = torch.float32 if pos.device.type == "mps" else torch.float64
+    if dt is torch.float32:
+        pos = pos.to(dt)
+    scale = torch.arange(0, dim, 2, dtype=dt, device=pos.device) / dim
     omega = 1.0 / ((theta * ntk) ** scale)
     out = torch.einsum("...n,d->...nd", pos, omega)
     out = torch.stack(

extensions_built_in/sd_trainer/SDTrainer.py

@@ -170,8 +170,11 @@ class SDTrainer(BaseSDTrainProcess):
                 # needed so we can autoparse the prompt to handle flags
                 gen_img_config = GenerateImageConfig(
                     prompt=prompt, # it will autoparse the prompt
+                    width=sample_item.width,
+                    height=sample_item.height,
                     negative_prompt=sample_item.neg,
                     output_path=output_path,
+                    num_frames=sample_item.num_frames,
                     ctrl_img=sample_item.ctrl_img,
                     ctrl_img_1=sample_item.ctrl_img_1,
                     ctrl_img_2=sample_item.ctrl_img_2,

toolkit/advanced_prompt_embeds.py

@@ -209,3 +209,4 @@ class AdvancedPromptEmbeds:
                 )
             for i in range(num_parts):
                 split_embeds[i]._store[key] = [values[i]]
+        return split_embeds

toolkit/dataloader_mixins.py

@@ -34,10 +34,13 @@ import albumentations as A
 from toolkit.print import print_acc
 from toolkit.accelerator import get_accelerator
 from toolkit.prompt_utils import PromptEmbeds
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 from torchvision.transforms import functional as TF
 
 from toolkit.train_tools import get_torch_dtype
 
+TEXT_EMBED_CACHE_BATCH = 16  # prompts per text-encoder forward when caching
+
 if TYPE_CHECKING:
     from toolkit.data_loader import AiToolkitDataset
     from toolkit.data_transfer_object.data_loader import FileItemDTO
@@ -2489,6 +2492,38 @@ class TextEmbeddingCachingMixin:
             
             did_move = False
 
+            # plain-text prompts are encoded TEXT_EMBED_CACHE_BATCH per forward: a 9B text
+            # encoder at batch 1 is kernel-launch bound (~15% GPU busy)
+            pending = []  # (path, caption), plain text only
+            pending_paths = set()
+            batch_ok = [None]  # None = untested; False = legacy PromptEmbeds, no per-item split
+
+            def flush_pending():
+                if not pending:
+                    return
+                # similar lengths together: a batch is right-padded to its longest prompt
+                pending.sort(key=lambda t: len(t[1]))
+                size = 1 if batch_ok[0] is False else TEXT_EMBED_CACHE_BATCH
+                start = 0
+                while start < len(pending):
+                    chunk = pending[start:start + size]
+                    embeds = self.sd.encode_prompt([c for _, c in chunk])
+                    if isinstance(embeds, AdvancedPromptEmbeds):
+                        batch_ok[0] = True
+                        parts = AdvancedPromptEmbeds.split_prompt_embeds(embeds, len(chunk))
+                    elif len(chunk) == 1:
+                        parts = [embeds]
+                    else:
+                        batch_ok[0] = False
+                        size = 1
+                        continue  # redo this chunk one prompt at a time
+                    for (path, _), part in zip(chunk, parts):
+                        part.save(path)
+                    del embeds, parts
+                    start += len(chunk)
+                pending.clear()
+                pending_paths.clear()
+
             # use tqdm to show progress
             i = 0
             for file_item in tqdm(self.file_list, desc='Caching text embeddings to disk'):
@@ -2616,9 +2651,11 @@ class TextEmbeddingCachingMixin:
                         file_item.tensor = None
                     else:
                         for path, caption in encode_targets:
-                            prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption)
-                            prompt_embeds.save(path)
-                            del prompt_embeds
+                            if path not in pending_paths:
+                                pending_paths.add(path)
+                                pending.append((path, caption))
+                        if len(pending) >= 4 * TEXT_EMBED_CACHE_BATCH:
+                            flush_pending()
                 if getattr(file_item, 'dopsd_self_ref', False):
                     # D-OPSD teacher embeds encode with the item's own media as the reference
                     control_video_paths = getattr(file_item, 'control_video_paths', None) or []
@@ -2665,6 +2702,7 @@ class TextEmbeddingCachingMixin:
                             del prompt_embeds
                 file_item.is_text_embedding_cached = True
                 i += 1
+            flush_pending()
             # restore device state
             # if did_move:
             #     self.sd.restore_device_state()

toolkit/memory_management/manager.py

@@ -457,15 +457,18 @@ class MemoryManager:
         gc.collect()
         torch.cuda.empty_cache()
 
-        # torch's pinned-host cache; private API, name varies by torch version
-        for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"):
-            fn = getattr(torch._C, fn_name, None)
-            if fn is not None:
-                try:
-                    fn()
-                    break
-                except Exception:
-                    pass
+        # torch's pinned-host cache; private API, name varies by torch version.
+        # Skipped on MPS: it segfaults there (pytorch/pytorch#197561), not catchable
+        on_mps = getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available()
+        if not on_mps:
+            for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"):
+                fn = getattr(torch._C, fn_name, None)
+                if fn is not None:
+                    try:
+                        fn()
+                        break
+                    except Exception:
+                        pass
 
         # glibc keeps freed arenas mapped; CDLL(None) resolves malloc_trim in
         # the running process where glibc is present and fails cleanly on

toolkit/models/diffusion_feature_extraction.py

@@ -1221,7 +1221,7 @@ class DiffusionFeatureExtractor9(nn.Module):
         perceptual_loss = torch.nn.functional.mse_loss(
             pred.float(), target.float(), reduction="none"
         )
-        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2)
+        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.01) ** 2)
         loss_perceptual = (perceptual_loss * velocity_equiv_weight).mean()
         
         if self.do_partial_step:
@@ -1379,7 +1379,7 @@ class DiffusionFeatureExtractor10(nn.Module):
 
         pred_feats = self.get_lpips_features(pred_images.float())
 
-        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2)
+        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.01) ** 2)
 
         loss_perceptual = 0
         for idx, pred_feat in enumerate(pred_feats):

toolkit/models/v2/_mixin.py

@@ -903,6 +903,8 @@ class OstrisModelMixin:
                 value = value.to(dtype=dtype)
             state_dict[key] = value.detach().to("cpu").contiguous()
         state_dict.update(q_entries)
+        # block_compile wraps blocks in OptimizedModule: drop its "._orig_mod." key segment
+        state_dict = {k.replace("._orig_mod.", "."): v for k, v in state_dict.items()}
         state_dict = self.convert_state_dict_on_save(state_dict)
 
         parent = os.path.dirname(output_path)

toolkit/optimizer.py

@@ -103,6 +103,9 @@ def get_optimizer(
     elif lower_type == 'automagic3':
         from toolkit.optimizers.automagic3 import Automagic3
         optimizer = Automagic3(params, lr=float(learning_rate), **optimizer_params)
+    elif lower_type == 'automagic4':
+        from toolkit.optimizers.automagic4 import Automagic4
+        optimizer = Automagic4(params, lr=float(learning_rate), **optimizer_params)
     elif lower_type == 'automagicexperiment':
         from toolkit.optimizers.automagicEXPERIMENT import AutomagicEXPERIMENT
         optimizer = AutomagicEXPERIMENT(params, lr=float(learning_rate), **optimizer_params)

toolkit/optimizers/automagic4.py

@@ -0,0 +1,259 @@
+"""
+NOTE: This is experimental and under active development; expect breaking changes and bugs. Feedback welcome.
+"""
+
+import math
+import torch
+
+from toolkit.optimizers.automagic3 import Automagic3
+
+
+class Automagic4(Automagic3):
+    """
+    Automagic v4: Automagic3 with sensitivity-balanced updates.
+
+    v3 gives every element an Adam-normalized step of about lr, so how far one
+    step moves the OUTPUT scales with that weight's sensitivity. A single shared
+    lr cannot fit a model whose tensors differ in sensitivity by orders of
+    magnitude (a head reading a huge raw residual, a modulation shared by every
+    layer): at the lr the bulk needs, the sensitive ones overshoot.
+
+    v4 multiplies each element's (clipped) v3 update by
+
+        balance = min(1, reference / sensitivity)
+
+    sensitivity = sqrt(bias-corrected second moment). For a noise-dominated MSE
+    loss (diffusion) gradient power tracks the Gauss-Newton curvature, i.e. how
+    far the output moves per unit change of the weight. It comes from the same
+    row/col (or 1D) state v3 keeps, so it is free and per element: per column
+    for a head whose input dims differ in scale, per row for a modulation.
+    Above the reference the step becomes Newton-like (grad / v), which is
+    invariant to rescaling the weights -- including coupled Q/K-style
+    rescalings, so it cannot seed the fights the shared group lr prevents.
+
+    One-sided on purpose: low gradient power cannot be told apart from a weight
+    that is gated off (a LoRA down while its zero-initialized up is still ~0,
+    dead units), where Newton would take an unbounded step. Those keep v3's
+    normalized step, which is what gets zero-initialized LoRA pairs moving.
+
+    reference: the ``sensitivity_quantile`` (default 0.9) quantile of the
+    per-tensor geometric-mean sensitivity, weighted by element count, from the
+    previous step. High so a large near-dead mass (LoRA downs at init are about
+    half of a LoRA's elements) cannot drag it down and stall everything above
+    it; the outliers it targets are a few percent of elements.
+
+    The first step only measures (second moments and the reference): without a
+    reference the sensitive tensors would take a full v3 step.
+
+    lr controller: re-centered flip-rate vote. v3's neutral point is a 50% sign
+    flip rate, but a weight bouncing on gradient noise at its minimum has
+    lag-1 gradient correlation -effective_lr * curvature / 2 at any lr > 0, so
+    v3 only balances at lr = 0 and anneals once the signal runs out. v4 votes
+    (neutral - observed) flip rate per element, with
+
+        neutral = 0.5 + asin(step_target / 2) / pi
+
+    the flip rate of a stationary weight whose effective_lr * curvature equals
+    ``step_target`` (default 0.1; 2 is the stability edge). The lr then has a
+    fixed point set by step_target, with or without signal: trends flip less
+    and raise it, overshoot flips more and lowers it. Votes are weighted by the
+    balanced |update| (output impact). No gain: near the fixed point the flip
+    rate moves ~0.16 * step_target per e-fold of lr, so the loop gain over the
+    H-step window lag stays well below 1.
+    """
+
+    def __init__(
+        self,
+        params,
+        *args,
+        sensitivity_quantile: float = 0.9,
+        step_target: float = 0.1,
+        **kwargs,
+    ):
+        if not 0.0 < sensitivity_quantile <= 1.0:
+            raise ValueError(
+                f"sensitivity_quantile ({sensitivity_quantile}) must be in (0, 1]"
+            )
+        if not 0.0 < step_target < 2.0:
+            raise ValueError(f"step_target ({step_target}) must be in (0, 2)")
+        self.sensitivity_quantile = float(sensitivity_quantile)
+        self.step_target = float(step_target)
+        self._neutral_flip = 0.5 + math.asin(self.step_target / 2.0) / math.pi
+        super().__init__(params, *args, **kwargs)
+
+    # ------------------------------------------------------------------ utils
+
+    @staticmethod
+    def _ema_(
+        state_t: torch.Tensor, sample: torch.Tensor, beta2: float
+    ) -> torch.Tensor:
+        # fp32 beta2-EMA; low-precision state is written back
+        v = state_t if state_t.dtype == torch.float32 else state_t.to(torch.float32)
+        v.mul_(beta2).add_(sample, alpha=1.0 - beta2)
+        if v is not state_t:
+            state_t.copy_(v)
+        return v
+
+    def _rebuild_group_index(self) -> None:
+        super()._rebuild_group_index()
+        n = len(self.param_groups)
+        self._group_ref: list = [None] * n
+        self._group_sens_log: list = [[] for _ in range(n)]
+        self._group_sens_numel: list = [[] for _ in range(n)]
+
+    def _init_state(self, p: torch.Tensor, group: dict) -> None:
+        super()._init_state(p, group)
+        gi = self._param_group_index.get(p)
+        ref = self._group_ref[gi] if gi is not None else None
+        if ref is not None:
+            self.state[p]["sens_ref"] = ref.to(p.device, copy=True)
+
+    # -------------------------------------------------------------- per-param
+
+    @torch.no_grad()
+    def _update_param(self, p: torch.Tensor, group: dict) -> None:
+        if p.grad is None:
+            return
+        state = self.state[p]
+        if len(state) == 0:
+            self._init_state(p, group)
+
+        grad = p.grad
+        if grad.is_sparse:
+            raise RuntimeError("Automagic4 does not support sparse gradients.")
+        if grad.dtype != torch.float32:
+            grad = grad.to(torch.float32)
+        grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
+
+        beta2 = group["beta2"]
+        eps = group["eps"]
+        sq = grad * grad
+        if p.dim() >= 2:
+            row = self._ema_(state["exp_avg_sq_row"], sq.mean(dim=-1).add_(eps), beta2)
+            col = self._ema_(state["exp_avg_sq_col"], sq.mean(dim=-2).add_(eps), beta2)
+            del sq
+            # mean log of the factored v = row (x) col / mean(row)
+            log_v = (
+                row.log().mean(-1) - row.mean(-1).log() + col.log().mean(-1)
+            ).mean()
+            inv_sqrt_v = self._approx_sq_grad(row, col)
+        else:
+            v = self._ema_(state["exp_avg_sq"], sq, beta2).add(eps)
+            del sq
+            log_v = v.log().mean()
+            inv_sqrt_v = v.rsqrt_()
+        state["step"] += 1
+        bias_corr = 1.0 - beta2 ** state["step"]
+
+        gi = self._param_group_index.get(p)
+        if gi is not None:
+            self._group_sens_log[gi].append(0.5 * (log_v - math.log(bias_corr)))
+            self._group_sens_numel[gi].append(p.numel())
+
+        ref = state.get("sens_ref")
+        if ref is None:
+            # measurement-only first step (no reference yet)
+            p.grad = None
+            return
+
+        update = grad.mul_(inv_sqrt_v)
+        update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+        update.clamp_(-group["clip_threshold"], group["clip_threshold"])
+        # balance = min(1, ref / sensitivity), sensitivity = sqrt(v / bias_corr)
+        balance = inv_sqrt_v.mul_(ref * math.sqrt(bias_corr)).clamp_(max=1.0)
+        update.mul_(balance)
+        del balance, inv_sqrt_v
+
+        # flip-rate vote (see class docstring), pooled per group, applied in .step()
+        cur_bits = update.gt(0.0)
+        hist = state["sign_history"]
+        idx = state["hist_idx"]
+        H = hist.shape[0]
+        lr_t = state["lr"]
+        hist[idx].copy_(self._pack_bits(cur_bits))
+        state["hist_idx"] = (idx + 1) % H
+        fill = min(H, state["hist_fill"] + 1)
+        state["hist_fill"] = fill
+
+        if fill == H and gi is not None:
+            _, shifts = self._pack_consts(hist.device)
+            chron = torch.roll(hist, -state["hist_idx"], dims=0)
+            bits = (
+                (chron.unsqueeze(-1) >> shifts)
+                .bitwise_and_(1)
+                .view(H, -1)[:, : update.numel()]
+            )
+            flips = (bits[1:] ^ bits[:-1]).sum(0, dtype=torch.int16)
+            w = update.abs().view(-1)
+            den = w.sum()
+            num = den * self._neutral_flip - (w * flips).sum() / (H - 1)
+            if self._group_num[gi] is None:
+                self._group_num[gi] = num
+                self._group_den[gi] = den
+            else:
+                acc = self._group_num[gi]
+                if num.device != acc.device:
+                    num = num.to(acc.device)
+                    den = den.to(acc.device)
+                acc.add_(num)
+                self._group_den[gi].add_(den)
+
+        wd = group["weight_decay"]
+        if p.dtype == torch.float32:
+            if wd != 0.0:
+                update.add_(p, alpha=wd)
+            p.addcmul_(update, lr_t, value=-1.0)
+        else:
+            new_p_fp32 = p.to(torch.float32)
+            if wd != 0.0:
+                update.add_(new_p_fp32, alpha=wd)
+            new_p_fp32.addcmul_(update, lr_t, value=-1.0)
+            self._stochastic_copy_(p, new_p_fp32)
+
+        p.grad = None
+
+    # ----------------------------------------------------------- optimizer API
+
+    def _apply_group_votes(self) -> None:
+        self._update_sensitivity_refs()
+        super()._apply_group_votes()
+
+    def _update_sensitivity_refs(self) -> None:
+        # numel-weighted quantile of this step's per-tensor log sensitivities; no GPU sync
+        for gi, group in enumerate(self.param_groups):
+            logs = self._group_sens_log[gi]
+            if not logs:
+                continue
+            dev = logs[0].device
+            vals = torch.stack([t.to(dev) for t in logs])
+            weights = torch.tensor(
+                self._group_sens_numel[gi], dtype=torch.float64, device=dev
+            )
+            order = vals.argsort()
+            cum = weights[order].cumsum(0)
+            pos = torch.searchsorted(cum, cum[-1:] * self.sensitivity_quantile)
+            ref = vals[order][pos.clamp_(max=len(logs) - 1)].squeeze(0).exp_()
+            self._group_ref[gi] = ref
+            for p in group["params"]:
+                st = self.state.get(p)
+                if not st:
+                    continue
+                cur = st.get("sens_ref")
+                if cur is None:
+                    st["sens_ref"] = ref.to(p.device, copy=True)
+                else:
+                    cur.copy_(ref)
+            self._group_sens_log[gi] = []
+            self._group_sens_numel[gi] = []
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        # parent cast the reference to param dtype; restore fp32 and the group mirror
+        for gi, group in enumerate(self.param_groups):
+            for p in group["params"]:
+                st = self.state.get(p)
+                ref = st.get("sens_ref") if st is not None else None
+                if isinstance(ref, torch.Tensor):
+                    st["sens_ref"] = ref.to(torch.float32)
+                    if self._group_ref[gi] is None:
+                        self._group_ref[gi] = st["sens_ref"].clone()

toolkit/util/unified_memory.py

@@ -1,13 +1,19 @@
-"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10).
+"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10) and
+the Windows RTX Spark (N1X).
 
 On these devices the CPU and GPU share the same physical memory, so moving a
 module to the CPU frees nothing. Worse, device -> host copies are extremely
-slow (~100MB/s on GB10), so every "offload to cpu" of a large transformer or
-text encoder stalls for minutes. When a unified memory device is detected we
-make module level moves to the CPU a no-op (dtype casts still apply). Tensor
-level .cpu()/.to('cpu') is left alone since saving, numpy, etc. need it.
+slow (~100MB/s on GB10 under Linux, ~0.5-2GB/s pageable on Windows/WDDM), so
+every "offload to cpu" of a large transformer or text encoder stalls. When a
+unified memory device is detected we make module level moves to the CPU a
+no-op (dtype casts still apply). Tensor level .cpu()/.to('cpu') is left alone
+since saving, numpy, etc. need it.
 
-Set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
+Note for Windows: WDDM gives CUDA a fixed budget (~80% of system RAM, e.g.
+~45GB of 54GB) that host allocations do not count against. nvidia-smi / Task
+Manager only show the small dedicated segment (8GB); the real pool is what
+torch.cuda.mem_get_info() reports. If a job needs to keep modules on the CPU
+to stay under that budget, set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
 """
 import os
 import sys
@@ -20,8 +26,8 @@ _patched = False
 def is_unified_memory() -> bool:
     if os.environ.get("AITK_DISABLE_UNIFIED_MEMORY", "0") == "1":
         return False
-    # linux only (DGX Spark). The Windows RTX Spark stack does not behave the same way
-    if not sys.platform.startswith("linux"):
+    # macOS has no CUDA; MPS unified memory is handled elsewhere
+    if sys.platform == "darwin":
         return False
     # NVIDIA only. ROCm APUs report integrated too, but may have a dedicated VRAM carveout
     if torch.version.hip is not None:

ui/src/app/api/files/info/route.ts

@@ -0,0 +1,111 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot, getTrainingFolder } from '@/server/settings';
+
+// safetensors header: u64 LE byte length, then that many bytes of JSON.
+const MAX_HEADER_BYTES = 256 * 1024 * 1024;
+
+const DTYPE_BYTES: Record<string, number> = {
+  F64: 8, I64: 8, U64: 8,
+  F32: 4, I32: 4, U32: 4,
+  F16: 2, BF16: 2, I16: 2, U16: 2,
+  F8_E4M3: 1, F8_E5M2: 1, F8_E8M0: 1, I8: 1, U8: 1, BOOL: 1,
+  F4: 0.5, U4: 0.5, I4: 0.5,
+};
+
+export async function POST(request: NextRequest) {
+  try {
+    const body = await request.json();
+    const { filePath } = body;
+
+    if (!filePath || typeof filePath !== 'string') {
+      return new NextResponse('filePath is required', { status: 400 });
+    }
+
+    const decodedFilePath = decodeURIComponent(filePath);
+    const datasetRoot = await getDatasetsRoot();
+    const trainingRoot = await getTrainingFolder();
+    const allowedDirs = [datasetRoot, trainingRoot];
+
+    const resolvedFilePath = path.resolve(decodedFilePath);
+    const isAllowed = allowedDirs.some(
+      allowedDir => resolvedFilePath === allowedDir || resolvedFilePath.startsWith(allowedDir + path.sep),
+    );
+    if (!isAllowed) {
+      console.warn(`Access denied: ${resolvedFilePath} not in ${allowedDirs.join(', ')}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    let stat;
+    try {
+      stat = await fs.promises.stat(resolvedFilePath);
+    } catch {
+      return new NextResponse('File not found', { status: 404 });
+    }
+    if (!stat.isFile()) {
+      return new NextResponse('Not a file', { status: 400 });
+    }
+    if (stat.size < 8) {
+      return new NextResponse('Not a safetensors file', { status: 400 });
+    }
+
+    const handle = await fs.promises.open(resolvedFilePath, 'r');
+    let header: Record<string, any>;
+    let headerSize: number;
+    try {
+      const lenBuf = Buffer.alloc(8);
+      await handle.read(lenBuf, 0, 8, 0);
+      const headerLen = lenBuf.readBigUInt64LE(0);
+      if (headerLen > BigInt(MAX_HEADER_BYTES) || headerLen > BigInt(stat.size - 8)) {
+        return new NextResponse('Not a safetensors file (bad header length)', { status: 400 });
+      }
+      headerSize = Number(headerLen);
+      const headerBuf = Buffer.alloc(headerSize);
+      await handle.read(headerBuf, 0, headerSize, 8);
+      header = JSON.parse(headerBuf.toString('utf8'));
+    } catch (e) {
+      return new NextResponse('Not a safetensors file (unparseable header)', { status: 400 });
+    } finally {
+      await handle.close();
+    }
+
+    const metadata: Record<string, string> = header.__metadata__ ?? {};
+    const tensors: { name: string; dtype: string; shape: number[]; numel: number; bytes: number }[] = [];
+    const dtypeCounts: Record<string, { count: number; numel: number; bytes: number }> = {};
+    let totalParams = 0;
+    let totalBytes = 0;
+
+    for (const [name, info] of Object.entries(header)) {
+      if (name === '__metadata__') continue;
+      const t = info as { dtype: string; shape: number[]; data_offsets?: [number, number] };
+      const shape = Array.isArray(t.shape) ? t.shape : [];
+      const numel = shape.reduce((a, b) => a * b, 1);
+      const bytes = t.data_offsets ? t.data_offsets[1] - t.data_offsets[0] : numel * (DTYPE_BYTES[t.dtype] ?? 0);
+      tensors.push({ name, dtype: t.dtype, shape, numel, bytes });
+      totalParams += numel;
+      totalBytes += bytes;
+      const d = (dtypeCounts[t.dtype] ??= { count: 0, numel: 0, bytes: 0 });
+      d.count += 1;
+      d.numel += numel;
+      d.bytes += bytes;
+    }
+    tensors.sort((a, b) => a.name.localeCompare(b.name));
+
+    return NextResponse.json({
+      filePath: resolvedFilePath,
+      fileSize: stat.size,
+      headerSize,
+      metadata,
+      tensors,
+      tensorCount: tensors.length,
+      totalParams,
+      totalBytes,
+      dtypeCounts,
+    });
+  } catch (error) {
+    console.error('Error reading safetensors info:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

ui/src/app/api/jobs/[jobID]/loss/route.ts

@@ -5,6 +5,8 @@ import fs from 'fs';
 import { getTrainingFolder } from '@/server/settings';
 
 import sqlite3 from 'sqlite3';
+import zlib from 'zlib';
+import { promisify } from 'util';
 
 export const runtime = 'nodejs';
 
@@ -29,6 +31,31 @@ function closeDb(db: sqlite3.Database) {
   });
 }
 
+// Serialize a metric series to JSON arrays inside SQLite. Building the arrays
+// with json_group_array avoids marshaling one JS object per row, which is the
+// dominant cost on multi-million-point logs. Values are truncated to 7
+// significant digits (float32-level precision) to shrink the wire payload;
+// value_text rows are never numeric (see _coerce_value in logging_aitk.py)
+// so they serialize as null.
+const SERIES_SQL = `
+  SELECT json_group_array(step) AS steps, json_group_array(v) AS vals
+  FROM (
+    SELECT step,
+           CASE WHEN value_real IS NULL THEN NULL
+                ELSE CAST(printf('%.7g', value_real) AS REAL) END AS v
+    FROM metrics
+    WHERE key = ? AND step >= ?
+    ORDER BY step ASC
+  )
+`;
+
+const gzip = promisify(zlib.gzip);
+
+// GET /api/jobs/:id/loss?since_step=N
+// Returns every metric key in one response: {keys, since_step, series: {key:
+// {steps: number[], values: (number|null)[]}}}. Points with step >= since_step
+// are included (default 0 = everything). The body is gzipped when the client
+// accepts it; a 150k-step run with a dozen keys is tens of MB uncompressed.
 export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
   // this must be awaited to avoid TS error
   const { jobID } = await params;
@@ -43,58 +70,46 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
   try {
     await fs.promises.access(logPath);
   } catch {
-    return NextResponse.json({ keys: [], key: 'loss', points: [] });
+    return NextResponse.json({ keys: [], since_step: 0, series: {} });
   }
 
   const url = new URL(request.url);
-  const key = url.searchParams.get('key') ?? 'loss';
-  const limit = Math.min(Number(url.searchParams.get('limit') ?? 2000), 20000);
-  const sinceStepParam = url.searchParams.get('since_step');
-  const sinceStep = sinceStepParam != null ? Number(sinceStepParam) : null;
-  const stride = Math.max(1, Number(url.searchParams.get('stride') ?? 1));
+  const sinceStepParam = Number(url.searchParams.get('since_step') ?? 0);
+  const sinceStep = Number.isFinite(sinceStepParam) ? sinceStepParam : 0;
 
   const db = openDb(logPath);
 
+  let body: string;
   try {
     const keysRows = await all<{ key: string }>(db, `SELECT key FROM metric_keys ORDER BY key ASC`);
     const keys = keysRows.map((r) => r.key);
 
-    const points = await all<{
-      step: number;
-      wall_time: number;
-      value: number | null;
-      value_text: string | null;
-    }>(
-      db,
-      `
-      SELECT
-        m.step AS step,
-        s.wall_time AS wall_time,
-        m.value_real AS value,
-        m.value_text AS value_text
-      FROM metrics m
-      JOIN steps s ON s.step = m.step
-      WHERE m.key = ?
-        AND (? IS NULL OR m.step > ?)
-        AND (m.step % ?) = 0
-      ORDER BY m.step ASC
-      LIMIT ?
-      `,
-      [key, sinceStep, sinceStep, stride, limit]
-    );
-
-    return NextResponse.json({
-      key,
-      keys,
-      points: points.map((p) => ({
-        step: p.step,
-        wall_time: p.wall_time,
-        value: p.value ?? (p.value_text ? Number(p.value_text) : null),
-      })),
-    });
+    const parts: string[] = [];
+    for (const key of keys) {
+      const rows = await all<{ steps: string; vals: string }>(db, SERIES_SQL, [key, sinceStep]);
+      const steps = rows[0]?.steps ?? '[]';
+      const vals = rows[0]?.vals ?? '[]';
+      parts.push(`${JSON.stringify(key)}:{"steps":${steps},"values":${vals}}`);
+    }
+    body = `{"keys":${JSON.stringify(keys)},"since_step":${sinceStep},"series":{${parts.join(',')}}}`;
   } finally {
     await closeDb(db);
   }
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+    'Cache-Control': 'no-store',
+    Vary: 'Accept-Encoding',
+  };
+  const acceptEncoding = request.headers.get('accept-encoding') ?? '';
+  if (/\bgzip\b/.test(acceptEncoding)) {
+    // Level 1: the payload is mostly digits, so higher levels buy ~5% for 3x
+    // the CPU time.
+    const gz = await gzip(Buffer.from(body), { level: 1 });
+    headers['Content-Encoding'] = 'gzip';
+    return new NextResponse(gz, { headers });
+  }
+  return new NextResponse(body, { headers });
 }
 
 function run(db: sqlite3.Database, sql: string, params: any[] = []) {

ui/src/app/jobs/new/SimpleJob.tsx

@@ -745,6 +745,7 @@ export default function SimpleJob({
                     { value: 'automagic', label: 'Automagic' },
                     { value: 'automagic2', label: 'Automagic v2' },
                     { value: 'automagic3', label: 'Automagic v3' },
+                    { value: 'automagic4', label: 'Automagic v4' },
                     { value: 'prodigyopt', label: 'Prodigy' },
                     { value: 'prodigy8bit', label: 'Prodigy8Bit' },
                   ]}

ui/src/app/layout.tsx

@@ -10,6 +10,7 @@ import DocModal from '@/components/DocModal';
 import os from 'os';
 import { CaptionDatasetModal } from '@/components/CaptionDatasetModal';
 import MergeLoRAsModal from '@/components/MergeLoRAsModal';
+import SafetensorsInfoModal from '@/components/SafetensorsInfoModal';
 import UpsamplePromptsModal from '@/components/UpsamplePromptsModal';
 import PromptBoxEditorModal from '@/components/PromptBoxEditorModal';
 
@@ -65,6 +66,7 @@ export default function RootLayout({ children }: { children: React.ReactNode })
         <DocModal />
         <CaptionDatasetModal />
         <MergeLoRAsModal />
+        <SafetensorsInfoModal />
         <UpsamplePromptsModal />
         <PromptBoxEditorModal />
       </body>

ui/src/components/FilesWidget.tsx

@@ -1,7 +1,8 @@
 import React from 'react';
 import useFilesList from '@/hooks/useFilesList';
-import { Loader2, AlertCircle, Download, Box, Brain, Trash2, SlidersHorizontal } from 'lucide-react';
+import { Loader2, AlertCircle, Download, Box, Brain, Trash2, SlidersHorizontal, Info } from 'lucide-react';
 import { openMergeLoRAsModal } from './MergeLoRAsModal';
+import { openSafetensorsInfoModal } from './SafetensorsInfoModal';
 import { getFilename, getFoldername, encodeFilePathForUrl } from '@/utils/basic';
 import { openConfirm } from './ConfirmModal';
 import { apiClient } from '@/utils/api';
@@ -114,6 +115,16 @@ export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName
                   </a>
                   <div className="flex items-center space-x-3 flex-shrink-0">
                     <span className="text-xs text-gray-400">{cleanSize(file.size)}</span>
+                    {fileName.endsWith('.safetensors') && (
+                      <button
+                        type="button"
+                        onClick={() => openSafetensorsInfoModal(file.path)}
+                        className="bg-sky-500 bg-opacity-0 group-hover:bg-opacity-10 hover:!bg-opacity-30 rounded-full p-1 transition-all"
+                        title="Safetensors info"
+                      >
+                        <Info className="w-3 h-3 text-sky-400" />
+                      </button>
+                    )}
                     <a
                       target="_blank"
                       href={`/api/files/${encodeFilePathForUrl(file.path)}`}

ui/src/components/JobLossGraph.tsx

@@ -1,7 +1,7 @@
 'use client';
 
 import { Job } from '@prisma/client';
-import useJobLossLog, { LossPoint } from '@/hooks/useJobLossLog';
+import useJobLossLog, { LossSeries } from '@/hooks/useJobLossLog';
 import { useMemo, useState, useEffect, useRef, useCallback } from 'react';
 import uPlot from 'uplot';
 import 'uplot/dist/uPlot.min.css';
@@ -233,7 +233,16 @@ export default function JobLossGraph({ job }: Props) {
     });
   }, [lossKeys]);
 
-  const activeKeys = useMemo(() => lossKeys.filter(k => enabled[k] !== false), [lossKeys, enabled]);
+  // Keys not yet in `enabled` (data and keys now arrive in the same render,
+  // before the sync effect above has run) take their default here rather than
+  // rendering as enabled and getting pruned a frame later.
+  const activeKeys = useMemo(
+    () =>
+      lossKeys.filter(
+        k => enabled[k] ?? persistedEnabledRef.current?.[k] ?? (k === 'loss/loss' || k === 'val/loss'),
+      ),
+    [lossKeys, enabled],
+  );
 
   // Assign palette slots by position in the (sorted) lossKeys list rather than
   // by hashing the key name — hashing let different keys collide onto the same
@@ -257,11 +266,11 @@ export default function JobLossGraph({ job }: Props) {
     // Union of all steps across active series.
     const stepSet = new Set<number>();
     for (const key of activeKeys) {
-      const pts: LossPoint[] = series[key] ?? [];
-      for (const p of pts) {
-        if (p.value === null || !Number.isFinite(p.value as number)) continue;
-        if (useLogScale && (p.value as number) <= 0) continue;
-        stepSet.add(p.step);
+      const s: LossSeries | undefined = series[key];
+      if (!s) continue;
+      for (let i = 0; i < s.steps.length; i++) {
+        if (useLogScale && s.values[i] <= 0) continue;
+        stepSet.add(s.steps[i]);
       }
     }
     let xs = Array.from(stepSet).sort((a, b) => a - b);
@@ -291,13 +300,13 @@ export default function JobLossGraph({ job }: Props) {
     for (let ki = 0; ki < activeKeys.length; ki++) {
       const key = activeKeys[ki];
       const scaleKey = `y::${key}`;
-      const pts: LossPoint[] = series[key] ?? [];
+      const s: LossSeries = series[key] ?? { steps: [], values: [] };
       const map = new Map<number, number>();
-      for (const p of pts) {
-        if (p.value === null || !Number.isFinite(p.value as number)) continue;
-        if (useLogScale && (p.value as number) <= 0) continue;
-        if (!xsSet.has(p.step)) continue;
-        map.set(p.step, p.value as number);
+      for (let i = 0; i < s.steps.length; i++) {
+        const v = s.values[i];
+        if (useLogScale && v <= 0) continue;
+        if (!xsSet.has(s.steps[i])) continue;
+        map.set(s.steps[i], v);
       }
       const raw: (number | null)[] = xs.map(s => (map.has(s) ? (map.get(s) as number) : null));
       // Metrics logged less often than every step are null at most x positions.

ui/src/components/SafetensorsInfoModal.tsx

@@ -0,0 +1,281 @@
+'use client';
+import React, { useEffect, useMemo, useState } from 'react';
+import { createGlobalState } from 'react-global-hooks';
+import { Loader2, AlertCircle, Search } from 'lucide-react';
+import { Modal } from './Modal';
+import { getFilename } from '@/utils/basic';
+import { apiClient } from '@/utils/api';
+
+export interface SafetensorsInfoModalState {
+  filePath: string;
+}
+
+interface TensorInfo {
+  name: string;
+  dtype: string;
+  shape: number[];
+  numel: number;
+  bytes: number;
+}
+
+interface SafetensorsInfo {
+  filePath: string;
+  fileSize: number;
+  headerSize: number;
+  metadata: Record<string, string>;
+  tensors: TensorInfo[];
+  tensorCount: number;
+  totalParams: number;
+  totalBytes: number;
+  dtypeCounts: Record<string, { count: number; numel: number; bytes: number }>;
+}
+
+export const safetensorsInfoModalState = createGlobalState<SafetensorsInfoModalState | null>(null);
+
+export const openSafetensorsInfoModal = (filePath: string) => {
+  safetensorsInfoModalState.set({ filePath });
+};
+
+const fmtBytes = (size: number) => {
+  if (size < 1024) return `${size} B`;
+  if (size < 1024 * 1024) return `${(size / 1024).toFixed(1)} KB`;
+  if (size < 1024 * 1024 * 1024) return `${(size / (1024 * 1024)).toFixed(1)} MB`;
+  return `${(size / (1024 * 1024 * 1024)).toFixed(2)} GB`;
+};
+
+const fmtCount = (n: number) => {
+  if (n < 1000) return `${n}`;
+  if (n < 1e6) return `${(n / 1e3).toFixed(1)}K`;
+  if (n < 1e9) return `${(n / 1e6).toFixed(1)}M`;
+  return `${(n / 1e9).toFixed(2)}B`;
+};
+
+// metadata values are strings; many are JSON blobs worth pretty-printing
+const fmtMetaValue = (value: string) => {
+  const trimmed = value.trim();
+  if ((trimmed.startsWith('{') && trimmed.endsWith('}')) || (trimmed.startsWith('[') && trimmed.endsWith(']'))) {
+    try {
+      return JSON.stringify(JSON.parse(trimmed), null, 2);
+    } catch {
+      /* not json */
+    }
+  }
+  return value;
+};
+
+const MAX_ROWS = 2000;
+
+const SafetensorsInfoModal: React.FC = () => {
+  const [modalInfo, setModalInfo] = safetensorsInfoModalState.use();
+  const isOpen = modalInfo !== null;
+  const [info, setInfo] = useState<SafetensorsInfo | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [isLoading, setIsLoading] = useState(false);
+  const [filter, setFilter] = useState('');
+  const [tab, setTab] = useState<'tensors' | 'metadata'>('tensors');
+
+  useEffect(() => {
+    if (!modalInfo) {
+      setInfo(null);
+      setError(null);
+      setIsLoading(false);
+      setFilter('');
+      setTab('tensors');
+      return;
+    }
+    let cancelled = false;
+    setIsLoading(true);
+    setInfo(null);
+    setError(null);
+    apiClient
+      .post('/api/files/info', { filePath: modalInfo.filePath })
+      .then(res => {
+        if (cancelled) return;
+        setInfo(res.data as SafetensorsInfo);
+      })
+      .catch(err => {
+        if (cancelled) return;
+        setError(err?.response?.data || err?.message || 'Failed to read file');
+      })
+      .finally(() => {
+        if (!cancelled) setIsLoading(false);
+      });
+    return () => {
+      cancelled = true;
+    };
+  }, [modalInfo]);
+
+  const onClose = () => setModalInfo(null);
+
+  const filteredTensors = useMemo(() => {
+    if (!info) return [];
+    const q = filter.trim().toLowerCase();
+    if (!q) return info.tensors;
+    return info.tensors.filter(t => t.name.toLowerCase().includes(q) || t.dtype.toLowerCase().includes(q));
+  }, [info, filter]);
+
+  const metadataEntries = info ? Object.entries(info.metadata) : [];
+  const title = modalInfo ? getFilename(modalInfo.filePath) : 'Safetensors Info';
+
+  return (
+    <Modal isOpen={isOpen} onClose={onClose} title={title} size="xl">
+      {isLoading && (
+        <div className="flex flex-col items-center justify-center py-16 space-y-3 text-gray-300">
+          <Loader2 className="w-8 h-8 text-purple-400 animate-spin" />
+          <span className="text-sm">Gathering info...</span>
+        </div>
+      )}
+
+      {!isLoading && error && (
+        <div className="flex items-center justify-center py-10 text-rose-400 space-x-2">
+          <AlertCircle className="w-5 h-5" />
+          <span className="text-sm">{String(error)}</span>
+        </div>
+      )}
+
+      {!isLoading && info && (
+        <div className="space-y-3">
+          <div className="grid grid-cols-2 sm:grid-cols-4 gap-2 text-xs">
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">File size</div>
+              <div className="text-gray-100 font-mono">{fmtBytes(info.fileSize)}</div>
+            </div>
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">Tensors</div>
+              <div className="text-gray-100 font-mono">{info.tensorCount.toLocaleString()}</div>
+            </div>
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">Parameters</div>
+              <div className="text-gray-100 font-mono" title={info.totalParams.toLocaleString()}>
+                {fmtCount(info.totalParams)}
+              </div>
+            </div>
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">Header</div>
+              <div className="text-gray-100 font-mono">{fmtBytes(info.headerSize)}</div>
+            </div>
+          </div>
+
+          <div className="flex flex-wrap gap-1.5 text-xs">
+            {Object.entries(info.dtypeCounts)
+              .sort((a, b) => b[1].bytes - a[1].bytes)
+              .map(([dtype, d]) => (
+                <span
+                  key={dtype}
+                  className="px-2 py-0.5 rounded-full bg-purple-500/10 text-purple-300 font-mono"
+                  title={`${d.count} tensors · ${d.numel.toLocaleString()} params`}
+                >
+                  {dtype} · {d.count} · {fmtBytes(d.bytes)}
+                </span>
+              ))}
+          </div>
+
+          <div className="flex items-center justify-between border-b border-gray-800">
+            <div className="flex space-x-1">
+              <button
+                type="button"
+                onClick={() => setTab('tensors')}
+                className={`px-3 py-1.5 text-sm border-b-2 ${
+                  tab === 'tensors'
+                    ? 'border-purple-500 text-gray-100'
+                    : 'border-transparent text-gray-400 hover:text-gray-200'
+                }`}
+              >
+                Tensors
+              </button>
+              <button
+                type="button"
+                onClick={() => setTab('metadata')}
+                className={`px-3 py-1.5 text-sm border-b-2 ${
+                  tab === 'metadata'
+                    ? 'border-purple-500 text-gray-100'
+                    : 'border-transparent text-gray-400 hover:text-gray-200'
+                }`}
+              >
+                Metadata
+                <span className="ml-1.5 px-1.5 py-0.5 bg-gray-700 rounded-full text-xs text-gray-300">
+                  {metadataEntries.length}
+                </span>
+              </button>
+            </div>
+            {tab === 'tensors' && (
+              <div className="relative">
+                <Search className="w-3.5 h-3.5 text-gray-500 absolute left-2 top-1/2 -translate-y-1/2" />
+                <input
+                  type="text"
+                  value={filter}
+                  onChange={e => setFilter(e.target.value)}
+                  placeholder="Filter tensors"
+                  className="text-xs pl-7 pr-2 py-1 w-56 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-md text-gray-100 focus:ring-1 focus:ring-gray-600 focus:outline-none"
+                />
+              </div>
+            )}
+          </div>
+
+          {tab === 'tensors' && (
+            <div className="overflow-auto max-h-[60vh] rounded-md border border-gray-800">
+              <table className="w-full text-xs font-mono">
+                <thead className="sticky top-0 bg-gray-800 text-gray-400">
+                  <tr>
+                    <th className="text-left px-2 py-1.5 font-normal">Name</th>
+                    <th className="text-left px-2 py-1.5 font-normal">Shape</th>
+                    <th className="text-left px-2 py-1.5 font-normal">Dtype</th>
+                    <th className="text-right px-2 py-1.5 font-normal">Params</th>
+                    <th className="text-right px-2 py-1.5 font-normal">Size</th>
+                  </tr>
+                </thead>
+                <tbody>
+                  {filteredTensors.slice(0, MAX_ROWS).map(t => (
+                    <tr key={t.name} className="border-t border-gray-800/60 hover:bg-gray-800/50">
+                      <td className="px-2 py-1 text-gray-200 break-all">{t.name}</td>
+                      <td className="px-2 py-1 text-gray-300 whitespace-nowrap">[{t.shape.join(', ')}]</td>
+                      <td className="px-2 py-1 text-purple-300 whitespace-nowrap">{t.dtype}</td>
+                      <td className="px-2 py-1 text-gray-400 text-right whitespace-nowrap">
+                        {t.numel.toLocaleString()}
+                      </td>
+                      <td className="px-2 py-1 text-gray-400 text-right whitespace-nowrap">{fmtBytes(t.bytes)}</td>
+                    </tr>
+                  ))}
+                  {filteredTensors.length === 0 && (
+                    <tr>
+                      <td colSpan={5} className="px-2 py-4 text-center text-gray-500">
+                        No tensors match
+                      </td>
+                    </tr>
+                  )}
+                </tbody>
+              </table>
+              {filteredTensors.length > MAX_ROWS && (
+                <div className="px-2 py-1.5 text-xs text-gray-500 bg-gray-800 border-t border-gray-700">
+                  Showing {MAX_ROWS.toLocaleString()} of {filteredTensors.length.toLocaleString()} — narrow with the
+                  filter
+                </div>
+              )}
+            </div>
+          )}
+
+          {tab === 'metadata' && (
+            <div className="overflow-auto max-h-[60vh] rounded-md border border-gray-800">
+              {metadataEntries.length === 0 ? (
+                <div className="px-2 py-4 text-center text-gray-500 text-xs">No metadata</div>
+              ) : (
+                <div className="divide-y divide-gray-800/60">
+                  {metadataEntries.map(([key, value]) => (
+                    <div key={key} className="px-2 py-1.5">
+                      <div className="text-xs text-purple-300 font-mono">{key}</div>
+                      <pre className="text-xs text-gray-200 font-mono whitespace-pre-wrap break-all mt-0.5">
+                        {fmtMetaValue(String(value))}
+                      </pre>
+                    </div>
+                  ))}
+                </div>
+              )}
+            </div>
+          )}
+        </div>
+      )}
+    </Modal>
+  );
+};
+
+export default SafetensorsInfoModal;

ui/src/hooks/useJobLossLog.tsx

@@ -4,13 +4,20 @@ import { useEffect, useState, useRef, useCallback, useMemo } from 'react';
 import { apiClient } from '@/utils/api';
 import usePollLoop from '@/hooks/usePollLoop';
 
-export interface LossPoint {
-  step: number;
-  wall_time?: number;
-  value: number | null;
+// Columnar series: steps[i] pairs with values[i]. Null / non-finite values are
+// dropped at ingest, so every value here is a finite number.
+export interface LossSeries {
+  steps: number[];
+  values: number[];
 }
 
-type SeriesMap = Record<string, LossPoint[]>;
+type SeriesMap = Record<string, LossSeries>;
+
+interface LossResponse {
+  keys?: string[];
+  since_step?: number;
+  series?: Record<string, { steps: number[]; values: (number | null)[] }>;
+}
 
 export default function useJobLossLog(jobID: string, reloadInterval: null | number = null) {
   const [series, setSeries] = useState<SeriesMap>({});
@@ -20,8 +27,10 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
   const didInitialLoadRef = useRef(false);
   const inFlightRef = useRef(false);
 
-  // track last step per key so polling is incremental per series
-  const lastStepByKeyRef = useRef<Record<string, number | null>>({});
+  // Highest step seen across every key. Polls ask for step >= this (not >) so a
+  // key logged at that step after the previous poll is still picked up; the
+  // per-key merge below drops the one re-sent point.
+  const lastStepRef = useRef<number | null>(null);
 
   const lossKeys = useMemo(() => {
     const base = keys ?? [];
@@ -40,69 +49,60 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
     setStatus(loadStatus);
 
     try {
-      // Step 1: get key list (we can do this by calling endpoint once; it returns keys)
-      // Keep it cheap: limit=1.
-      const first = await apiClient
-        .get(`/api/jobs/${jobID}/loss`, { params: { key: 'loss', limit: 1 } })
-        .then(res => res.data as { keys?: string[] });
+      // One request returns every key. Incremental after the initial load.
+      const params: Record<string, number> = {};
+      if (reloadInterval && didInitialLoadRef.current && lastStepRef.current != null) {
+        params.since_step = lastStepRef.current;
+      }
+      const res = await apiClient.get(`/api/jobs/${jobID}/loss`, { params });
+      const data = res.data as LossResponse;
 
-      const newKeys = first.keys ?? [];
+      const newKeys = data.keys ?? [];
       setKeys(newKeys);
-
       const wantedLossKeys = (newKeys.length ? [...newKeys] : ['loss']).sort();
-
-      // Step 2: fetch each loss key incrementally (since_step per key if polling)
-      const requests = wantedLossKeys.map(k => {
-        const params: Record<string, any> = { key: k };
-
-        if (reloadInterval && lastStepByKeyRef.current[k] != null) {
-          params.since_step = lastStepByKeyRef.current[k];
-        }
-
-        params.limit = 1000000;
-
-        return apiClient
-          .get(`/api/jobs/${jobID}/loss`, { params })
-          .then(res => res.data as { key: string; points?: LossPoint[] });
-      });
-
-      const results = await Promise.all(requests);
+      const incoming = data.series ?? {};
+      const initial = !didInitialLoadRef.current;
 
       setSeries(prev => {
-        const next: SeriesMap = { ...prev };
+        const next: SeriesMap = initial ? {} : { ...prev };
+        let maxStep = lastStepRef.current ?? null;
 
-        for (const r of results) {
-          const k = r.key;
-          const newPoints = (r.points ?? []).filter(p => p.value !== null);
+        for (const k of wantedLossKeys) {
+          const inc = incoming[k];
+          const existing = initial ? undefined : next[k];
+          const prevLast = existing && existing.steps.length ? existing.steps[existing.steps.length - 1] : null;
 
-          if (!didInitialLoadRef.current) {
-            // initial: replace
-            next[k] = newPoints;
-          } else if (newPoints.length) {
-            const existing = next[k] ?? [];
-            const prevLast = existing.length ? existing[existing.length - 1].step : null;
-            const filtered = prevLast == null ? newPoints : newPoints.filter(p => p.step > prevLast);
-            next[k] = filtered.length ? [...existing, ...filtered] : existing;
-          } else {
-            // no new points: keep existing
-            next[k] = next[k] ?? [];
+          const steps: number[] = existing ? existing.steps.slice() : [];
+          const values: number[] = existing ? existing.values.slice() : [];
+          let appended = false;
+          if (inc) {
+            const n = Math.min(inc.steps.length, inc.values.length);
+            for (let i = 0; i < n; i++) {
+              const v = inc.values[i];
+              if (v === null || !Number.isFinite(v)) continue;
+              const s = inc.steps[i];
+              if (prevLast != null && s <= prevLast) continue;
+              steps.push(s);
+              values.push(v as number);
+              appended = true;
+            }
           }
+          if (initial || appended) next[k] = { steps, values };
+          else next[k] = existing ?? { steps: [], values: [] };
 
-          // update last step per key
-          const finalArr = next[k] ?? [];
-          lastStepByKeyRef.current[k] = finalArr.length
-            ? finalArr[finalArr.length - 1].step
-            : (lastStepByKeyRef.current[k] ?? null);
+          const arr = next[k].steps;
+          if (arr.length) {
+            const last = arr[arr.length - 1];
+            if (maxStep == null || last > maxStep) maxStep = last;
+          }
         }
 
         // remove stale keys that no longer exist (rare, but keeps UI clean)
         for (const existingKey of Object.keys(next)) {
-          if (!wantedLossKeys.includes(existingKey)) {
-            delete next[existingKey];
-            delete lastStepByKeyRef.current[existingKey];
-          }
+          if (!wantedLossKeys.includes(existingKey)) delete next[existingKey];
         }
 
+        lastStepRef.current = maxStep;
         return next;
       });
 
@@ -124,13 +124,24 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
       await apiClient.delete(`/api/jobs/${jobID}/loss`, { data: { min_step: minStep, max_step: maxStep } });
       setSeries(prev => {
         const next: SeriesMap = {};
+        let last: number | null = null;
         for (const k of Object.keys(prev)) {
-          const kept = prev[k].filter(p => p.step < minStep || p.step > maxStep);
-          next[k] = kept;
-          // Re-anchor incremental polling to the new tail so a deleted tail
-          // isn't treated as already-fetched.
-          lastStepByKeyRef.current[k] = kept.length ? kept[kept.length - 1].step : null;
+          const steps: number[] = [];
+          const values: number[] = [];
+          const src = prev[k];
+          for (let i = 0; i < src.steps.length; i++) {
+            const s = src.steps[i];
+            if (s < minStep || s > maxStep) {
+              steps.push(s);
+              values.push(src.values[i]);
+            }
+          }
+          next[k] = { steps, values };
+          if (steps.length && (last == null || steps[steps.length - 1] > last)) last = steps[steps.length - 1];
         }
+        // Re-anchor incremental polling to the new tail so a deleted tail
+        // isn't treated as already-fetched.
+        lastStepRef.current = last;
         return next;
       });
     },
@@ -141,7 +152,7 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
   // before the first fetch when jobID changes.
   useEffect(() => {
     didInitialLoadRef.current = false;
-    lastStepByKeyRef.current = {};
+    lastStepRef.current = null;
     setSeries({});
     setKeys([]);
     setStatus('idle');

ui/src/server/monitor.ts

@@ -36,11 +36,20 @@ const NV_BATCH_FLUSH_MS = 100;
 const CPU_TEMP_REFRESH_MS = 5000;
 
 // Power limits for GPUs whose driver reports none (nvidia-smi gives [N/A]).
-// GB10 (DGX Spark) is a unified SoC with no settable GPU limit; 140 W is
-// NVIDIA's rated TDP for the GB10 superchip.
-const KNOWN_POWER_LIMITS_W: Record<string, number> = {
-  'NVIDIA GB10': 140,
-};
+// These unified SoCs have no settable GPU limit, so show the rated TDP of the
+// whole chip. GB10 (DGX Spark) is 140 W. RTX Spark N1X: the 5120-core part is
+// laptop-only, rated 45-80 W (OEM configurable); the 6144-core part is 140 W
+// on desktop and 45-80 W in laptops but shares its name, so it falls through
+// to the desktop figure.
+const KNOWN_POWER_LIMITS_W: [RegExp, number][] = [
+  [/\bGB10\b/, 140],
+  [/RTX Spark.*5120-core/, 80],
+  [/RTX Spark/, 140],
+];
+
+function knownPowerLimit(name: string): number {
+  return KNOWN_POWER_LIMITS_W.find(([re]) => re.test(name))?.[1] ?? NaN;
+}
 
 function parseGpuLine(line: string): GpuInfo | null {
   const [
@@ -76,7 +85,7 @@ function parseGpuLine(line: string): GpuInfo | null {
     },
     power: {
       draw: parseFloat(powerDraw),
-      limit: isNaN(parseFloat(powerLimit)) ? (KNOWN_POWER_LIMITS_W[name] ?? NaN) : parseFloat(powerLimit),
+      limit: isNaN(parseFloat(powerLimit)) ? knownPowerLimit(name) : parseFloat(powerLimit),
     },
     clocks: {
       graphics: parseInt(clockGraphics),
@@ -88,19 +97,28 @@ function parseGpuLine(line: string): GpuInfo | null {
   };
 }
 
+// Integrated NVIDIA GPUs that share system RAM with the CPU.
+const UNIFIED_MEMORY_GPU = /\b(GB10|RTX Spark)\b/i;
+
 /**
- * Unified-memory GPUs (GB10 / DGX Spark) have no dedicated VRAM, so
- * nvidia-smi reports memory as [N/A]. Report the system RAM pool they share
- * instead, like the Mac path does for Apple silicon.
+ * Unified-memory GPUs (GB10 / DGX Spark, RTX Spark) have no dedicated VRAM.
+ * On Linux nvidia-smi reports their memory as [N/A]; on Windows (WDDM) it
+ * reports only the small dedicated segment (8 GB) even though CUDA can use
+ * most of system RAM. Report the system RAM pool they share instead, like
+ * the Mac path does for Apple silicon.
  */
+function isSharedMemoryGpu(g: GpuInfo): boolean {
+  return isNaN(g.memory.total) || UNIFIED_MEMORY_GPU.test(g.name);
+}
+
 function withSharedMemory(gpu: GPUApiResponse, cpu: CpuInfo | null): GPUApiResponse {
-  if (!cpu || !gpu.gpus.some(g => isNaN(g.memory.total))) return gpu;
+  if (!cpu || !gpu.gpus.some(isSharedMemoryGpu)) return gpu;
   const total = Math.round(cpu.totalMemory);
   const free = Math.round(cpu.availableMemory);
   return {
     ...gpu,
     gpus: gpu.gpus.map(g =>
-      isNaN(g.memory.total) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g,
+      isSharedMemoryGpu(g) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g,
     ),
   };
 }

version.py

@@ -1 +1 @@
-VERSION = "0.13.37"
+VERSION = "0.13.40"