streamft 0.2.0__py3-none-any.whl

streamft/metrics.py

@@ -0,0 +1,110 @@
+"""Evaluation metrics: PPL, KL, top-k agreement, routing diagnostics, timing."""
+from __future__ import annotations
+
+import json
+import os
+import statistics
+import time
+
+import torch
+import torch.nn.functional as F
+
+
+def cuda_time_ms(fn, warmup=5, iters=20):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    ts = []
+    for _ in range(iters):
+        s = torch.cuda.Event(True)
+        e = torch.cuda.Event(True)
+        s.record()
+        fn()
+        e.record()
+        torch.cuda.synchronize()
+        ts.append(s.elapsed_time(e))
+    return {"median": round(statistics.median(ts), 4), "min": round(min(ts), 4),
+            "mean": round(statistics.mean(ts), 4)}
+
+
+@torch.no_grad()
+def eval_pair(teacher_logits_fn, student_logits_fn, chunks, device, max_chunks=None,
+              pos_chunk=256):
+    """Compare teacher vs student next-token distributions over token chunks.
+
+    Returns dict with CE/PPL for both, KL(teacher||student) and top-k agreements.
+    """
+    stats = {"ce_t": 0.0, "ce_s": 0.0, "kl_ts": 0.0, "kl_st": 0.0,
+             "top1": 0.0, "top5": 0.0, "n": 0}
+    n_chunks = len(chunks) if max_chunks is None else min(len(chunks), max_chunks)
+    for ci in range(n_chunks):
+        ids = chunks[ci].to(device).unsqueeze(0)  # (1,T)
+        t_logits = teacher_logits_fn(ids)
+        s_logits = student_logits_fn(ids)
+        # predict tokens 1..T-1
+        labels = ids[0, 1:]
+        lt = t_logits[0, :-1]
+        ls = s_logits[0, :-1]
+        T = labels.shape[0]
+        for p0 in range(0, T, pos_chunk):
+            p1 = min(T, p0 + pos_chunk)
+            a = lt[p0:p1].float()
+            b = ls[p0:p1].float()
+            lab = labels[p0:p1]
+            la = F.log_softmax(a, dim=-1)
+            lb = F.log_softmax(b, dim=-1)
+            stats["ce_t"] += F.nll_loss(la, lab, reduction="sum").item()
+            stats["ce_s"] += F.nll_loss(lb, lab, reduction="sum").item()
+            stats["kl_ts"] += F.kl_div(lb, la, log_target=True, reduction="sum").item()
+            stats["kl_st"] += F.kl_div(la, lb, log_target=True, reduction="sum").item()
+            t1 = a.argmax(-1)
+            s1 = b.argmax(-1)
+            stats["top1"] += (t1 == s1).sum().item()
+            tk = torch.topk(a, 5, dim=-1).indices
+            stats["top5"] += (tk == s1.unsqueeze(-1)).any(-1).sum().item()
+            stats["n"] += lab.numel()
+        del t_logits, s_logits
+    n = max(stats["n"], 1)
+    out = {
+        "ntok": n,
+        "ce_teacher": stats["ce_t"] / n,
+        "ce_student": stats["ce_s"] / n,
+        "ppl_teacher": float(torch.exp(torch.tensor(stats["ce_t"] / n))),
+        "ppl_student": float(torch.exp(torch.tensor(stats["ce_s"] / n))),
+        "kl_teacher_student": stats["kl_ts"] / n,
+        "kl_student_teacher": stats["kl_st"] / n,
+        "top1_agreement": stats["top1"] / n,
+        "top5_agreement": stats["top5"] / n,
+    }
+    return out
+
+
+def routing_diagnostics(captured, n_layers=None):
+    """Block-mass recall of the selected blocks per layer (from capture entries)."""
+    per_layer = []
+    for entry in captured:
+        if "mass" not in entry or "idx" not in entry:
+            continue
+        mass = entry["mass"]           # (B,H,NB,NB) normalized over J
+        idx = entry["idx"]             # (B,NB,MMAX)
+        B, H, NB, _ = mass.shape
+        MMAX = idx.shape[2]
+        sel = torch.zeros(B, 1, NB, NB, device=mass.device, dtype=torch.bool)
+        valid = idx >= 0
+        for j in range(MMAX):
+            jj = idx[:, :, j].clamp_min(0).long()  # (B,NB)
+            sel.scatter_(3, jj[:, None, :, None], True)
+        sel = sel & valid.any(-1)[:, None, :, None]
+        recall = (mass * sel).sum(-1) / mass.sum(-1).clamp_min(1e-9)  # (B,H,NB)
+        per_layer.append({"recall_mean": float(recall.mean()), "recall_min": float(recall.min())})
+    return per_layer
+
+
+def append_jsonl(path, record):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "a", encoding="utf-8") as f:
+        f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+
+def now():
+    return time.strftime("%Y-%m-%d %H:%M:%S")

streamft/stream_quant_lora.py

@@ -0,0 +1,216 @@
+"""Streaming QLoRA: train LoRA adapters on a 4-bit base while blocks stream CPU<->GPU.
+
+This is the "our own engine" counterpart to the HF in-memory QLoRA attempt. The
+base model is loaded in 4-bit (NF4) on CPU with `device_map="cpu"`; only one
+decoder block is on the GPU at a time. That keeps VRAM bounded by one block +
+activations + embedding/head, independent of model depth, so 7B can be trained
+on a 6 GB card even though the full 4-bit model + compute temporaries do not fit
+(HF's whole-model QLoRA peaked at 8.4-8.7 GB and spilled).
+
+The backward pass re-streams blocks in reverse order and accumulates LoRA
+gradients on CPU; the base 4-bit weights are never updated. This is the same
+recompute strategy as `lcsa/stream_train.py` (P1), but with a PEFT LoRA on a
+bitsandbytes 4-bit base instead of the minimal local LoRA on a bf16 base.
+"""
+from __future__ import annotations
+
+import time
+
+import torch
+
+from .stream_train import _layer_out, chunked_lm_loss_and_backward
+
+
+def chunked_lm_loss_and_backward_masked(head, h, ys, chunk=128, ignore_index=-100):
+    """Chunked CE with an ignore_index (for SFT prompt masking)."""
+    B, T, _ = h.shape
+    total = torch.zeros((), device=h.device, dtype=torch.float32)
+    valid_total = 0
+    for a in range(0, T, chunk):
+        b = min(a + chunk, T)
+        logits = head(h[:, a:b])
+        tgt = ys[:, a:b].reshape(-1).to(h.device)
+        valid = int((tgt != ignore_index).sum())
+        if valid == 0:
+            continue
+        loss_sum = torch.nn.functional.cross_entropy(
+            logits.reshape(-1, logits.shape[-1]).float(), tgt,
+            reduction="sum", ignore_index=ignore_index,
+        )
+        (loss_sum / valid).backward()
+        with torch.no_grad():
+            total = total + loss_sum.detach()
+        valid_total += valid
+    denom = max(valid_total, 1)
+    return total / denom, total / denom
+
+
+def base_model_of(model):
+    """Return the underlying HF causal-LM (the module owning `.model.layers` + `.lm_head`).
+
+    Works for plain HF models and for peft/wrapper models alike: wrappers delegate
+    attributes, so we take the *deepest* module that owns both the decoder stack
+    and the LM head.
+    """
+    found = None
+    stack = [model]
+    seen = set()
+    while stack:
+        m = stack.pop()
+        if id(m) in seen:
+            continue
+        seen.add(id(m))
+        if hasattr(m, "lm_head") and hasattr(getattr(m, "model", None), "layers"):
+            found = m
+        stack.extend(m.children())
+    return found if found is not None else model
+
+
+class StreamingQuantLoRATrainer:
+    """LoRA fine-tuning on a 4-bit base with block streaming."""
+
+    def __init__(self, model, device: str = "cuda", loss_chunk: int = 128,
+                 lr: float = 1e-4, betas=(0.9, 0.999), eps: float = 1e-8,
+                 weight_decay: float = 0.0, clip: float = 0.0,
+                 verbose: bool = False, stage=None, release=None,
+                 cache_device: str = "gpu"):
+        self.device = device
+        self.loss_chunk = loss_chunk
+        self.clip = clip
+        self.verbose = verbose
+        # Where the per-layer input activations for the recompute-backward live.
+        # "cpu" offloads them (frees VRAM, costs a PCIe round-trip per layer), which
+        # is what makes >1k contexts fit on a 6 GB card.
+        self.cache_device = device if cache_device == "gpu" else "cpu"
+        self.model = model
+        # How a decoder block is brought to the GPU and taken away again. The
+        # defaults assume the whole 4-bit base sits in host RAM; a disk-backed
+        # store can pass its own load/release callables instead.
+        self.stage = stage if stage is not None else (lambda layer, i: layer.to(device))
+        self.release = release if release is not None else (lambda layer, i: layer.to("cpu"))
+        base = base_model_of(model)
+        self.base = base
+        self.layers = list(base.model.layers)
+        self.embed = base.model.embed_tokens
+        self.rotary = base.model.rotary_emb
+        self.norm = base.model.norm
+        self.lm_head = base.lm_head
+
+        # Small/shared modules stay resident on the GPU. The input embedding is
+        # kept on CPU (4-bit/bf16) and only the gathered activation is moved to
+        # the GPU; for 14B the embedding alone is ~1.5 GB bf16. Exception: when
+        # the LM head is *tied* to the embedding (e.g. Qwen2.5-3B/0.5B) the same
+        # tensor feeds the head, so it has to live where the head runs.
+        self.rotary.to(device)
+        self.norm.to(device)
+        self.tied = self.lm_head.weight is self.embed.weight
+        if self.tied:
+            self.embed.to(device)
+        else:
+            self.lm_head.to(device)
+
+        self.trainables = [p for p in model.parameters() if p.requires_grad]
+        self.n_lora = sum(p.numel() for p in self.trainables)
+        self.opt = torch.optim.AdamW(self.trainables, lr=lr, betas=betas, eps=eps,
+                                     weight_decay=weight_decay)
+        self._acc = {id(p): torch.zeros_like(p, device="cpu", dtype=torch.float32)
+                     for p in self.trainables}
+        # layer indices that received a non-zero adapter gradient (verification:
+        # a silently skipped block would never appear here)
+        self.grad_layers = set()
+
+    # ------------------------------------------------------------------ forward
+    @torch.no_grad()
+    def _forward_cached(self, ids):
+        dev = self.device
+        # Embedding usually lives on CPU (only the gathered activation moves to
+        # the GPU); for tied models it lives on the GPU together with the head.
+        emb_dev = self.embed.weight.device
+        h = self.embed(ids.to(emb_dev))
+        if h.device != torch.device(dev):
+            h = h.to(dev)
+        pos_ids = torch.arange(ids.shape[1], device=dev).unsqueeze(0)
+        pos = self.rotary(h, pos_ids)
+        cache = []
+        for i, layer in enumerate(self.layers):
+            self.stage(layer, i)
+            # CPU offload must be synchronous: a non-blocking copy leaves the GPU
+            # source alive until the copy drains, and with 64 layers queued that
+            # keeps ~the whole cache resident (measured: 6.6 GB at seq 2048).
+            cache.append(h if self.cache_device == self.device
+                         else h.to(self.cache_device))
+            h = _layer_out(layer(h, position_embeddings=pos))
+            self.release(layer, i)
+            if self.verbose and (i + 1) % 8 == 0:
+                print(f"    [fwd] {i + 1}/{len(self.layers)} blocks "
+                      f"reserved={torch.cuda.memory_reserved()/1e9:.2f}GB", flush=True)
+        return h, cache, pos
+
+    # ----------------------------------------------------------------- backward
+    def _backward_blocks(self, cache, pos, grad):
+        dev = self.device
+        for i in range(len(self.layers) - 1, -1, -1):
+            layer = self.layers[i]
+            self.stage(layer, i)
+            h_in = cache[i]
+            if h_in.device != torch.device(dev):
+                h_in = h_in.to(dev)
+            h_in.requires_grad_(True)
+            out = _layer_out(layer(h_in, position_embeddings=pos))
+            out.backward(grad)
+            grad = h_in.grad
+            h_in.grad = None
+            h_in.requires_grad_(False)
+            for p in layer.parameters():
+                if p.requires_grad and p.grad is not None:
+                    if float(p.grad.abs().sum()) > 0:
+                        self.grad_layers.add(i)
+                    self._acc[id(p)].add_(p.grad.detach().float().cpu())
+                    p.grad = None
+            self.release(layer, i)
+            if self.verbose and (len(self.layers) - i) % 8 == 0:
+                print(f"    [bwd] {len(self.layers) - i}/{len(self.layers)} blocks "
+                      f"reserved={torch.cuda.memory_reserved()/1e9:.2f}GB", flush=True)
+
+    # --------------------------------------------------------------------- step
+    def backward(self, xs, ys, ignore_index=None):
+        """Forward + backward on one batch, without touching the optimizer.
+
+        Returns (loss, grads, peak_gb) where grads maps id(param) -> cloned tensor.
+        This is the exact path step() uses, exposed so the gradients can be checked
+        against finite differences or another implementation.
+        """
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        with torch.no_grad():
+            h_blocks, cache, pos = self._forward_cached(xs)
+        h_b = h_blocks.detach().requires_grad_(True)
+        z_g = self.norm(h_b)
+        z = z_g.detach().requires_grad_(True)
+        if ignore_index is None:
+            _, loss = chunked_lm_loss_and_backward(self.lm_head, z, ys, self.loss_chunk)
+        else:
+            _, loss = chunked_lm_loss_and_backward_masked(
+                self.lm_head, z, ys, self.loss_chunk, ignore_index=ignore_index
+            )
+        z_g.backward(z.grad)
+        self._backward_blocks(cache, pos, h_b.grad)
+        grads = {id(p): self._acc[id(p)].clone() for p in self.trainables}
+        for p in self.trainables:
+            self._acc[id(p)].zero_()
+        torch.cuda.synchronize()
+        return float(loss), grads, torch.cuda.max_memory_reserved() / 1e9
+
+    def step(self, xs, ys, ignore_index=None):
+        torch.cuda.synchronize()
+        t0 = time.time()
+        loss, grads, peak = self.backward(xs, ys, ignore_index)
+        for p in self.trainables:
+            p.grad = grads[id(p)].to(p.dtype)
+        if self.clip > 0:
+            torch.nn.utils.clip_grad_norm_(self.trainables, self.clip)
+        self.opt.step()
+        for p in self.trainables:
+            p.grad = None
+        torch.cuda.synchronize()
+        return loss, time.time() - t0, peak

streamft/stream_train.py

@@ -0,0 +1,230 @@
+"""P1: block-streaming ("parça parça") LoRA trainer for low-VRAM GPUs.
+
+Base weights live in CPU RAM. Each transformer block is moved to the GPU just
+for its forward pass, moved back, and re-streamed in reverse order during the
+backward pass (per-block activation recompute). Peak VRAM is bounded by
+(one block + its activations + embedding/head + loss chunk), independent of
+model depth — so a 6GB GPU can fine-tune models whose full weights alone do
+not fit (e.g. 3B fp16 = 6.2GB).
+
+Pipeline per step:
+  forward   : embed -> [block_i on GPU, run, off] x L -> final norm
+  loss      : chunked cross-entropy through lm_head (bounded logits memory)
+  backward  : dL/dh_final -> [block_i back on GPU, recompute forward with grad,
+              backward, accumulate LoRA grads on CPU, off] x L (reverse)
+  step      : AdamW on the (small, CPU-resident) LoRA params
+
+Correctness: the recompute backward is mathematically identical to a normal
+autograd backward over the same block list; experiments/exp8_stream.py checks
+the loss trajectory against a fully-on-GPU naive trainer on identical batches.
+"""
+from __future__ import annotations
+
+import time
+
+import torch
+import torch.nn.functional as F
+
+from .lora import apply_lora, count_lora_params, lora_parameters
+
+
+def _layer_out(out):
+    return out[0] if isinstance(out, tuple) else out
+
+
+def chunked_lm_loss_and_backward(head, h, ys, chunk=128):
+    """Cross-entropy in position chunks; accumulates dL/dh into h.grad.
+
+    h: (B, T, d) leaf requiring grad (or the output of a small graph).
+    Targets are pre-shifted (TrainSampler convention): ys[p] is the target of
+    position p, all T positions valid.
+
+    FIX (19.10.2026): the previous version used F.cross_entropy(mean) per chunk
+    and called backward(), which scaled the accumulated gradient by the number
+    of chunks (T/chunk). Now each chunk contributes a cross-entropy SUM divided
+    by n_pos, so the accumulated gradient is exactly d(mean CE)/dh.
+    """
+    B, T, _ = h.shape
+    n_pos = T
+    total = torch.zeros((), device=h.device, dtype=torch.float32)
+    last = None
+    for a in range(0, n_pos, chunk):
+        b = min(a + chunk, n_pos)
+        logits = head(h[:, a:b])
+        tgt = ys[:, a:b].reshape(-1).to(h.device)
+        loss_sum = F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(),
+                                   tgt, reduction="sum")
+        (loss_sum / n_pos).backward()
+        with torch.no_grad():
+            total = total + loss_sum.detach()
+        last = loss_sum
+    return last.detach() / n_pos, total / n_pos
+
+
+@torch.no_grad()
+def stream_eval_ce(trainer, ids, loss_chunk=128):
+    """Mean next-token CE over a (1, T) token chunk, via the streaming forward.
+
+    Positions 0..T-2 predict ids[1..T-1] (standard intra-window LM loss).
+    Reuses the trainer's block streaming; no backward, no cache retained.
+    """
+    h, _cache, _pos = trainer._forward_cached(ids)
+    z = trainer.norm(h)
+    total = 0.0
+    n = 0
+    T = ids.shape[1]
+    for a in range(0, T - 1, loss_chunk):
+        b = min(a + loss_chunk, T - 1)
+        logits = trainer.lm_head(z[:, a:b]).float()
+        tgt = ids[:, a + 1: b + 1].reshape(-1).to(z.device)
+        total += float(F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt,
+                                       reduction="sum"))
+        n += tgt.numel()
+    return total / max(n, 1)
+
+
+class StreamingLoRATrainer:
+    """LoRA fine-tuning while base blocks stream CPU <-> GPU."""
+
+    def __init__(self, model, device: str = "cuda", loss_chunk: int = 128,
+                 lr: float = 1e-3, betas=(0.9, 0.999), eps: float = 1e-8,
+                 weight_decay: float = 0.0, r_attn: int = 16, r_mlp: int = 0,
+                 clip: float = 0.0, verbose: bool = False):
+        self.device = device
+        self.loss_chunk = loss_chunk
+        self.clip = clip
+        self.verbose = verbose
+        apply_lora(model, r_attn=r_attn, r_mlp=r_mlp)
+        self.model = model
+        base = model.model
+        self.layers = list(base.layers)
+        self.embed = base.embed_tokens
+        self.rotary = base.rotary_emb
+        self.norm = base.norm
+        self.lm_head = model.lm_head
+
+        model.to("cpu")
+        for p in model.parameters():
+            p.requires_grad_(False)
+        # shared small modules stay resident on GPU
+        self.embed.to(device)
+        self.rotary.to(device)
+        self.norm.to(device)
+        if self.lm_head.weight is not self.embed.weight:
+            self.lm_head.to(device)
+        for p in self.lm_head.parameters():
+            p.requires_grad_(False)
+
+        self.trainables = list(lora_parameters(model))
+        for p in self.trainables:
+            p.requires_grad_(True)
+        self.n_lora = count_lora_params(model)
+        self.opt = torch.optim.AdamW(self.trainables, lr=lr, betas=betas, eps=eps,
+                                     weight_decay=weight_decay)
+        self._acc = {id(p): torch.zeros_like(p, device="cpu") for p in self.trainables}
+
+    # ------------------------------------------------------------------ forward
+    @torch.no_grad()
+    def _forward_cached(self, ids):
+        dev = self.device
+        h = self.embed(ids.to(dev))
+        pos_ids = torch.arange(ids.shape[1], device=dev).unsqueeze(0)
+        pos = self.rotary(h, pos_ids)
+        cache = []
+        for i, layer in enumerate(self.layers):
+            layer.to(dev)
+            cache.append(h)
+            h = _layer_out(layer(h, position_embeddings=pos))
+            layer.to("cpu")
+            if self.verbose and (i + 1) % 8 == 0:
+                print(f"    [fwd] {i + 1}/{len(self.layers)} blocks "
+                      f"reserved={torch.cuda.memory_reserved()/1e9:.2f}GB", flush=True)
+        return h, cache, pos
+
+    # ----------------------------------------------------------------- backward
+    def _backward_blocks(self, cache, pos, grad):
+        dev = self.device
+        for i in range(len(self.layers) - 1, -1, -1):
+            layer = self.layers[i]
+            layer.to(dev)
+            h_in = cache[i]
+            h_in.requires_grad_(True)
+            out = _layer_out(layer(h_in, position_embeddings=pos))
+            out.backward(grad)
+            grad = h_in.grad
+            h_in.grad = None
+            h_in.requires_grad_(False)
+            for p in layer.parameters():
+                if p.requires_grad and p.grad is not None:
+                    self._acc[id(p)].add_(p.grad.detach().float().cpu())
+                    p.grad = None
+            layer.to("cpu")
+            if self.verbose and (len(self.layers) - i) % 8 == 0:
+                print(f"    [bwd] {len(self.layers) - i}/{len(self.layers)} blocks "
+                      f"reserved={torch.cuda.memory_reserved()/1e9:.2f}GB", flush=True)
+
+    # --------------------------------------------------------------------- step
+    def step(self, xs, ys):
+        torch.cuda.synchronize()
+        t0 = time.time()
+        torch.cuda.reset_peak_memory_stats()
+        with torch.no_grad():
+            h_blocks, cache, pos = self._forward_cached(xs)
+        h_b = h_blocks.detach().requires_grad_(True)
+        z_g = self.norm(h_b)                       # graph h_b -> z_g (elementwise, cheap)
+        z = z_g.detach().requires_grad_(True)      # leaf: chunked head backward accumulates z.grad
+        _, loss = chunked_lm_loss_and_backward(self.lm_head, z, ys, self.loss_chunk)
+        z_g.backward(z.grad)                       # norm graph traversed exactly once
+        grad = h_b.grad
+        self._backward_blocks(cache, pos, grad)
+        for p in self.trainables:
+            p.grad = self._acc[id(p)]
+        if self.clip > 0:
+            torch.nn.utils.clip_grad_norm_(self.trainables, self.clip)
+        self.opt.step()
+        for p in self.trainables:
+            p.grad = None
+            self._acc[id(p)].zero_()
+        torch.cuda.synchronize()
+        peak = torch.cuda.max_memory_reserved() / 1e9
+        return float(loss), time.time() - t0, peak
+
+
+class NaiveLoRATrainer:
+    """Reference: everything on GPU, standard full-graph autograd + full logits."""
+
+    def __init__(self, model, device: str = "cuda", lr: float = 1e-3,
+                 betas=(0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.0,
+                 r_attn: int = 16, r_mlp: int = 0, clip: float = 0.0,
+                 verbose: bool = False):
+        self.device = device
+        self.clip = clip
+        self.verbose = verbose
+        apply_lora(model, r_attn=r_attn, r_mlp=r_mlp)
+        self.model = model.to(device)
+        for p in model.parameters():
+            p.requires_grad_(False)
+        self.trainables = list(lora_parameters(model))
+        for p in self.trainables:
+            p.requires_grad_(True)
+        self.n_lora = count_lora_params(model)
+        self.opt = torch.optim.AdamW(self.trainables, lr=lr, betas=betas, eps=eps,
+                                     weight_decay=weight_decay)
+
+    def step(self, xs, ys):
+        torch.cuda.synchronize()
+        t0 = time.time()
+        torch.cuda.reset_peak_memory_stats()
+        dev = self.device
+        ys = ys.to(dev)
+        h = self.model.model(input_ids=xs.to(dev), use_cache=False).last_hidden_state
+        logits = self.model.lm_head(h).float()
+        loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), ys.reshape(-1))
+        loss.backward()
+        if self.clip > 0:
+            torch.nn.utils.clip_grad_norm_(self.trainables, self.clip)
+        self.opt.step()
+        self.opt.zero_grad(set_to_none=True)
+        torch.cuda.synchronize()
+        peak = torch.cuda.max_memory_reserved() / 1e9
+        return float(loss.detach()), time.time() - t0, peak