Spaces:
Sleeping
Sleeping
| from typing import Any | |
| import torch | |
| import triton | |
| import triton.language as tl | |
| def _moments(X: tl.tensor, P: tl.tensor, N: tl.tensor, PARTS: tl.tensor, BLOCK: tl.constexpr) -> None: | |
| c = tl.program_id(0) | |
| k = tl.program_id(1) | |
| i = k * BLOCK + tl.arange(0, BLOCK) | |
| x = tl.load(X + c * N + i, i < N, 0).to(tl.float32) | |
| tl.store(P + (c * PARTS + k) * 2, tl.sum(x, 0)) | |
| tl.store(P + (c * PARTS + k) * 2 + 1, tl.sum(x * x, 0)) | |
| def _coefficients( | |
| P: tl.tensor, COEF: tl.tensor, N: tl.tensor, PARTS: tl.tensor, LO: tl.tensor, HI: tl.tensor, RLO: tl.tensor, RHI: tl.tensor, SCALE: tl.tensor, BLOCK: tl.constexpr | |
| ) -> None: | |
| c = tl.program_id(0) | |
| i = tl.arange(0, BLOCK) | |
| mean = tl.sum(tl.load(P + (c * PARTS + i) * 2, i < PARTS, 0), 0) / N | |
| second = tl.sum(tl.load(P + (c * PARTS + i) * 2 + 1, i < PARTS, 0), 0) / N | |
| variance = second - mean * mean | |
| std = tl.sqrt(tl.maximum(variance, 1.0e-20)) | |
| rms = tl.sqrt(tl.maximum(second, 1.0e-20)) | |
| ratio = mean / std | |
| direction = tl.where(ratio < LO, -1.0, tl.where(ratio > HI, 1.0, 0.0)) | |
| rdir = tl.where(rms < RLO, -1.0, tl.where(rms > RHI, 1.0, 0.0)) | |
| a = -direction * mean / (N * std * std * std) * (variance >= 1.0e-20) | |
| r = rdir / (N * tl.maximum(second, 1.0e-20)) * (second >= 1.0e-20) | |
| a += r | |
| b = direction / (N * std) + r * mean | |
| norm = tl.maximum(tl.sqrt(a * a * tl.maximum(variance, 0.0) + b * b), 1.0e-20) | |
| tl.store(COEF + c * 3, mean) | |
| tl.store(COEF + c * 3 + 1, a * (SCALE / norm)) | |
| tl.store(COEF + c * 3 + 2, b * (SCALE / norm)) | |
| def _apply(X: tl.tensor, G: tl.tensor, COEF: tl.tensor, Y: tl.tensor, N: tl.tensor, BLOCK: tl.constexpr) -> None: | |
| c = tl.program_id(0) | |
| i = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) | |
| x = tl.load(X + c * N + i, i < N, 0).to(tl.float32) | |
| g = tl.load(G + c * N + i, i < N, 0).to(tl.float32) | |
| mean = tl.load(COEF + c * 3) | |
| a = tl.load(COEF + c * 3 + 1) | |
| b = tl.load(COEF + c * 3 + 2) | |
| tl.store(Y + c * N + i, g + tl.abs(g) * (a * (x - mean) + b), i < N) | |
| def balancer_backward( | |
| x: torch.Tensor, grad: torch.Tensor, min_mean: float, max_mean: float, min_rms: float, max_rms: float, grad_scale: float, channel_dim: int | |
| ) -> torch.Tensor: | |
| xc = x.movedim(channel_dim, 0).contiguous() | |
| gc = grad.movedim(channel_dim, 0).contiguous() | |
| channels = xc.shape[0] | |
| n = xc.numel() // channels | |
| parts = triton.cdiv(n, 1024) | |
| partial = torch.empty((channels, parts, 2), device=x.device) | |
| coefficients = torch.empty((channels, 3), device=x.device) | |
| output = torch.empty_like(gc) | |
| _moments[(channels, parts)](xc, partial, n, parts, 1024, enable_fp_fusion=False) | |
| _coefficients[(channels,)](partial, coefficients, n, parts, min_mean, max_mean, min_rms, max_rms, grad_scale, triton.next_power_of_2(parts), enable_fp_fusion=False) | |
| _apply[(channels, parts)](xc, gc, coefficients, output, n, 1024, enable_fp_fusion=False) | |
| return output.movedim(0, channel_dim) | |
| class _Balancer(torch.autograd.Function): | |
| def forward(ctx: Any, x: torch.Tensor, config: tuple[float, float, float, float, float, int]) -> torch.Tensor: | |
| ctx.save_for_backward(x) | |
| ctx.config = config | |
| return x | |
| def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None]: | |
| (x,) = ctx.saved_tensors | |
| return balancer_backward(x, grad, *ctx.config), None | |
| def balance(x: torch.Tensor, min_mean: float, max_mean: float, min_rms: float, max_rms: float, grad_scale: float, channel_dim: int) -> torch.Tensor: | |
| return _Balancer.apply(x, (min_mean, max_mean, min_rms, max_rms, grad_scale, channel_dim)) | |