SolitaryThinker
1e58d90d02
Implement the GPU backend the substrate was built for: Platform.detect() -> cuda resolves real torch adapters + torch solver ops instead of the numpy rungs, with the existing loops/policies/scheduler/training unchanged. WRITTEN-NOT-RUN: this environment has no GPU/torch, so the torch code is grounded in the verbatim real fastvideo APIs (DiT forward signature confirmed from source) but cannot be executed/verified here. It is gated available=False (CPU mini stays green; importing the backends never imports torch), with every on-box confirm point marked `# BRINGUP` and an ordered checklist in platform/backends/GPU_BRINGUP.md. - torch_adapters.py: TorchWanDiT / TorchWanVAE / TorchT5Encoder wrap the real module named by each card's load_id and bridge it to the mini's duck-typed surface (numpy<->torch at the boundary; loop math stays numpy fp32). DiT weight-surface (copy_from/blend_from/clone) for serving sync; mse_grad_step raises (GPU training is a separate workstream). - torch_kernels.py: flow_match_step / flow_sde_step as plain torch elementwise. Grounded conclusion from the kernel audit: fastvideo-kernel ships NO fused solver kernel (only attention/norm/quant primitives), so the cuda solver is torch, registered at arch generic with an honest source string. - torch_cuda.py: rewritten as lazy trampolines (torch imported only inside builder/kernel bodies). Adds the missing vae + text_encoder cuda components (they'd otherwise silently fall back to the toy) and corrects the dishonest "fastvideo-kernel:flow_*" labels. - ComponentSpec.checkpoint: the weights source for the torch adapter (risk A; the one field the cards didn't carry). Empty on the CPU toys. - 7 CPU-verifiable wiring tests: honest registration/sources, torch-free import, cuda-resolves-real-cells-not-toy, build-fails-loudly-without-torch. 204 tests pass (CPU). The torch path needs a GPU box to verify (GPU_BRINGUP.md).
58 lines
2.6 KiB
Python
58 lines
2.6 KiB
Python
"""Torch solver-step kernels for the GPU backend (design_v3 §17).
|
|
|
|
GROUND TRUTH (from the fastvideo-kernel audit): there is NO fused flow-match / SDE *solver* kernel.
|
|
fastvideo-kernel ships only primitives — sparse/tiled attention, RMS/LayerNorm, int8 GEMM/quant.
|
|
So the GPU solver step is a plain torch elementwise op with the SAME math as ``loop/sampler.py``'s
|
|
numpy reference (the consistency oracle that the accel stand-in already proves matches bit-for-bit).
|
|
|
|
WRITTEN-NOT-RUN: imports torch; loaded only lazily from ``torch_cuda.py``. Under the current numpy
|
|
loop surface these marshal numpy↔torch at the boundary (a torch-native surface would keep the latent
|
|
on-device through forward→combine→solver — the perf follow-up, Risk G in GPU_BRINGUP.md).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def _as_torch(a):
|
|
"""numpy → cuda torch (passthrough if already torch). Returns (tensor, was_numpy)."""
|
|
if torch.is_tensor(a):
|
|
return a, False
|
|
return torch.as_tensor(np.asarray(a), device="cuda"), True
|
|
|
|
|
|
def _back(t, like_numpy: bool):
|
|
return t.detach().to("cpu", torch.float32).numpy() if like_numpy else t
|
|
|
|
|
|
def flow_match_step(x_t, velocity, sigma_t, sigma_next):
|
|
"""One deterministic flow-match Euler step — identical to sampler.flow_match_euler_step."""
|
|
xt, was_np = _as_torch(x_t)
|
|
v, _ = _as_torch(velocity)
|
|
return _back(xt + (float(sigma_next) - float(sigma_t)) * v, was_np)
|
|
|
|
|
|
def flow_sde_step(x_t, velocity, sigma_t, sigma_next, *, noise=None, prev_sample=None, noise_scale=0.7):
|
|
"""FlowGRPO stochastic step + Gaussian log-prob — the torch port of
|
|
sampler.flow_sde_step_with_logprob. Returns (prev_sample, log_prob, mean, eff_std)."""
|
|
xt, was_np = _as_torch(x_t)
|
|
v, _ = _as_torch(velocity)
|
|
xt = xt.double()
|
|
v = v.double()
|
|
s = min(float(sigma_t), 0.9999)
|
|
dt = float(sigma_next) - float(sigma_t)
|
|
std = float((s / (1.0 - s)) ** 0.5 * noise_scale)
|
|
denom = 2.0 * max(s, 1e-6)
|
|
mean = xt * (1.0 + std ** 2 / denom * dt) + v * (1.0 + std ** 2 * (1.0 - s) / denom) * dt
|
|
eff_std = max(std * max(-dt, 1e-12) ** 0.5, 1e-6)
|
|
if prev_sample is None:
|
|
n = _as_torch(noise)[0].double() if noise is not None else torch.zeros_like(xt)
|
|
prev = mean + eff_std * n
|
|
else:
|
|
prev = _as_torch(prev_sample)[0].double()
|
|
var = eff_std ** 2
|
|
log_prob = float((-((prev - mean) ** 2) / (2.0 * var)
|
|
- float(np.log(eff_std)) - 0.5 * float(np.log(2.0 * np.pi))).mean().item())
|
|
return (_back(prev.float(), was_np), log_prob, _back(mean.float(), was_np), float(eff_std))
|