diff --git a/configs/acoustic.yaml b/configs/acoustic.yaml index fad75600e..6ae4c7e78 100644 --- a/configs/acoustic.yaml +++ b/configs/acoustic.yaml @@ -3,6 +3,9 @@ base_config: task_cls: training.acoustic_task.AcousticTask +# Enable Triton-fused Linear+SoftSignGLU kernels for LYNXNet2 backbones. +use_fused_kernels: false + dictionaries: {} extra_phonemes: [] merged_phoneme_groups: [] diff --git a/configs/variance.yaml b/configs/variance.yaml index d4e203670..10f90d0f9 100644 --- a/configs/variance.yaml +++ b/configs/variance.yaml @@ -3,6 +3,9 @@ base_config: task_cls: training.variance_task.VarianceTask +# Enable Triton-fused Linear+SoftSignGLU kernels for LYNXNet2 backbones. +use_fused_kernels: false + dictionaries: {} extra_phonemes: [] merged_phoneme_groups: [] diff --git a/modules/backbones/lynxnet2.py b/modules/backbones/lynxnet2.py index 6e55d5f87..e2c717462 100644 --- a/modules/backbones/lynxnet2.py +++ b/modules/backbones/lynxnet2.py @@ -2,7 +2,9 @@ import torch.nn as nn import torch.nn.functional as F -from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, ATanGLU, Transpose, AdamWLinear +from modules.commons.common_layers import ( + SinusoidalPosEmb, SwiGLU, ATanGLU, SoftSignGLU, Transpose, AdamWLinear +) from utils.hparams import hparams @@ -14,6 +16,8 @@ def __init__(self, dim, expansion_factor, kernel_size=31, dropout=0., glu_type=' _glu = SwiGLU() elif glu_type == 'atanglu': _glu = ATanGLU() + elif glu_type == 'softsign_glu': + _glu = SoftSignGLU() else: raise ValueError(f'{glu_type} is not a valid activation') if float(dropout) > 0.: diff --git a/modules/commons/common_layers.py b/modules/commons/common_layers.py index 4da65693a..10852e20e 100644 --- a/modules/commons/common_layers.py +++ b/modules/commons/common_layers.py @@ -175,6 +175,47 @@ def forward(self, x): return out * torch.atan(gate) +class SoftSignGLUFunction(torch.autograd.Function): + """ATanGLUFunction-style memory trick for SoftSignGLU. + + softsign'(x) = 1/(1+|x|)^2 = (1-|softsign(x)|)^2, so both partial + derivatives of y = out * softsign(gate) are precomputable in forward: + dy/dout = softsign(gate) + dy/dgate = out * (1-|softsign(gate)|)^2 + Saves 2 tensors (vs 3 for naive autograd) and backward is two pure + multiplies with no softsign recompute. + """ + @staticmethod + def forward(ctx, out, gate): + ss_gate = torch.nn.functional.softsign(gate) + decay_out = out * (1.0 - ss_gate.abs()).square() + ctx.save_for_backward(ss_gate, decay_out) + return out * ss_gate + + @staticmethod + def backward(ctx, grad_output): + ss_gate, decay_out = ctx.saved_tensors + return grad_output * ss_gate, grad_output * decay_out + + +class SoftSignGLU(nn.Module): + """Gated Linear Unit with SoftSign gate: out * softsign(gate). + + More numerically stable than ATanGLU (no approximation needed in + Triton kernels) while providing similar gating behavior. + """ + def __init__(self, dim=-1): + super().__init__() + self.dim = dim + + def forward(self, x): + out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim) + if self.training: + return SoftSignGLUFunction.apply(out, gate) + else: + return out * torch.nn.functional.softsign(gate) + + class AdamWConv1d(torch.nn.Conv1d): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/modules/kernels/__init__.py b/modules/kernels/__init__.py new file mode 100644 index 000000000..d58de6003 --- /dev/null +++ b/modules/kernels/__init__.py @@ -0,0 +1 @@ +# Fused kernels for LYNXNet2 optimization \ No newline at end of file diff --git a/modules/kernels/fused_linear_softsign_glu.py b/modules/kernels/fused_linear_softsign_glu.py new file mode 100644 index 000000000..b66ce2ebb --- /dev/null +++ b/modules/kernels/fused_linear_softsign_glu.py @@ -0,0 +1,482 @@ +"""Fused Linear + SoftSignGLU for LYNXNet2. + +Computes ``left * softsign(gate)`` directly from a Linear layer whose output +weights are split into left and gate halves. The forward kernel uses fp32 +accumulators, and the backward path uses a Triton element-wise kernel followed +by cuBLAS matrix multiplications. +""" +import torch +import torch.nn.functional as F + +try: + import triton + import triton.language as tl + _TRITON_AVAILABLE = True +except ImportError: # no triton installed (e.g. Windows without the build) + _TRITON_AVAILABLE = False + + +# Minimum CUDA compute capability for Triton tl.dot (tensor cores). +# Volta (sm_70) is the floor; Turing sm_75 has fp16 tensor cores, Ampere +# sm_80 adds bf16. Anything older cannot run the fused kernel. +_MIN_CAPABILITY = (7, 0) + +_FUSED_CAPABLE = None + + +def is_triton_available(): + """Whether the Triton Python package imported successfully.""" + return _TRITON_AVAILABLE + + +def _fused_capable(): + """True only if the current CUDA device can run the fused kernel. + + Caches once per process. Returns False when Triton is missing, the + device is CPU, or the device's compute capability predates tensor cores. + """ + global _FUSED_CAPABLE + if _FUSED_CAPABLE is None: + _FUSED_CAPABLE = False + if _TRITON_AVAILABLE and torch.cuda.is_available(): + cap = torch.cuda.get_device_capability() + if cap >= _MIN_CAPABILITY: + _FUSED_CAPABLE = True + return _FUSED_CAPABLE + + +# --------------------------------------------------------------------------- +# Forward kernel +# --------------------------------------------------------------------------- + +if _TRITON_AVAILABLE: + + @triton.autotune( + configs=[ + # Small tiles + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=4, num_stages=3), + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=4, num_stages=3), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=4, num_stages=3), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=4, num_stages=3), + # Larger tiles; unsupported shared-memory sizes are pruned by Triton. + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64, 'GROUP_M': 8}, num_warps=4, num_stages=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=4, num_stages=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=4, num_stages=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_warps=8, num_stages=3), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64, 'GROUP_M': 8}, num_warps=8, num_stages=3), + ], + key=['M_BUCKET', 'N', 'K'], + ) + @triton.jit + def _fused_linear_softsign_glu_fwd_kernel( + x_ptr, w_left_ptr, w_right_ptr, b_left_ptr, b_right_ptr, + y_ptr, left_ptr, gate_ptr, + M, N, K, + M_BUCKET, # next_power_of_2(M) — autotune key only, not used in body + stride_x_b, stride_x_k, + stride_wl_n, stride_wl_k, + stride_wr_n, stride_wr_k, + stride_y_b, stride_y_n, + stride_l_b, stride_l_n, + stride_g_b, stride_g_n, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + ): + """ + y = (x @ W_left^T + b_left) * softsign(x @ W_right^T + b_right) + + N = output dim per GLU half (= inner_dim = dim × expansion_factor) + K = input feature dim (= dim for first Linear, inner_dim for second) + + 2D grid over (M // BLOCK_M, N // BLOCK_N) with grouped ordering: + programs are swizzled so that GROUP_M row-blocks share column tiles + while they are still hot in L2 (standard Triton matmul swizzle). + """ + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + # Grouped pid swizzle for L2 reuse + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + m_mask_2d = offs_m[:, None] < M + n_mask_nk = offs_n[:, None] < N # [BLOCK_N, 1] for N×K weight access + n_mask_mn = offs_n[None, :] < N # [1, BLOCK_N] for M×N output access + n_mask_1d = offs_n < N + + acc_left = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + acc_gate = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + + for k_start in range(0, K, BLOCK_K): + k_offs = k_start + offs_k + k_mask_2d = k_offs[None, :] < K + + x = tl.load( + x_ptr + offs_m[:, None] * stride_x_b + k_offs[None, :] * stride_x_k, + mask=m_mask_2d & k_mask_2d, other=0.0, + ) + wl = tl.load( + w_left_ptr + offs_n[:, None] * stride_wl_n + k_offs[None, :] * stride_wl_k, + mask=n_mask_nk & k_mask_2d, other=0.0, + ) + acc_left += tl.dot(x, wl.T) + + wr = tl.load( + w_right_ptr + offs_n[:, None] * stride_wr_n + k_offs[None, :] * stride_wr_k, + mask=n_mask_nk & k_mask_2d, other=0.0, + ) + acc_gate += tl.dot(x, wr.T) + + # Bias + b_left = tl.load(b_left_ptr + offs_n, mask=n_mask_1d, other=0.0) + b_right = tl.load(b_right_ptr + offs_n, mask=n_mask_1d, other=0.0) + acc_left += b_left + acc_gate += b_right + + # Computed in fp32 for numerical safety + gate_f32 = acc_gate.to(tl.float32) + ss_gate = gate_f32 / (1.0 + tl.abs(gate_f32)) + gated = acc_left * ss_gate + + # Write output y + tl.store( + y_ptr + offs_m[:, None] * stride_y_b + offs_n[None, :] * stride_y_n, + gated, mask=m_mask_2d & n_mask_mn, + ) + + # Save intermediates for backward + tl.store( + left_ptr + offs_m[:, None] * stride_l_b + offs_n[None, :] * stride_l_n, + acc_left, mask=m_mask_2d & n_mask_mn, + ) + tl.store( + gate_ptr + offs_m[:, None] * stride_g_b + offs_n[None, :] * stride_g_n, + acc_gate, mask=m_mask_2d & n_mask_mn, + ) + + + # --------------------------------------------------------------------------- + # Element-wise backward kernel — grad_left_pre, grad_gate + # + # GEMMs run through cuBLAS; Triton handles the element-wise gradient. + # --------------------------------------------------------------------------- + + @triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128}, num_warps=4, num_stages=2), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_warps=8, num_stages=2), + ], + key=['N'], # element-wise: tile choice is insensitive to M — never key on it + ) + @triton.jit + def _softsign_glu_bwd_elem_kernel( + left_ptr, gate_ptr, grad_y_ptr, + glp_ptr, gg_ptr, + M, N, + stride_l_b, stride_l_n, + stride_g_b, stride_g_n, + stride_gy_b, stride_gy_n, + stride_glp_b, stride_glp_n, + stride_gg_b, stride_gg_n, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + ): + """Element-wise SoftSignGLU backward. + + For y = l * softsign(g): + dy/dl = softsign(g) + dy/dg = l / (1+|g|)^2 + """ + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + m_mask = offs_m[:, None] < M + n_mask = offs_n[None, :] < N + + left = tl.load(left_ptr + offs_m[:, None] * stride_l_b + offs_n[None, :] * stride_l_n, + mask=m_mask & n_mask, other=0.0) + gate = tl.load(gate_ptr + offs_m[:, None] * stride_g_b + offs_n[None, :] * stride_g_n, + mask=m_mask & n_mask, other=0.0) + gy = tl.load(grad_y_ptr + offs_m[:, None] * stride_gy_b + offs_n[None, :] * stride_gy_n, + mask=m_mask & n_mask, other=0.0) + + gate_f32 = gate.to(tl.float32) + left_f32 = left.to(tl.float32) + denom_g = 1.0 / (1.0 + tl.abs(gate_f32)) + denom_g2 = denom_g * denom_g + ss_gate = gate_f32 * denom_g + grad_left_pre = gy * ss_gate + grad_gate = gy * (left_f32 * denom_g2) + + tl.store(glp_ptr + offs_m[:, None] * stride_glp_b + offs_n[None, :] * stride_glp_n, + grad_left_pre, mask=m_mask & n_mask) + tl.store(gg_ptr + offs_m[:, None] * stride_gg_b + offs_n[None, :] * stride_gg_n, + grad_gate, mask=m_mask & n_mask) + + + # --------------------------------------------------------------------------- + # Python wrapper — torch.autograd.Function + # --------------------------------------------------------------------------- + + class FusedLinearSoftSignGLUFn(torch.autograd.Function): + """Fused Linear(2K, K) + SoftSignGLU.""" + + @staticmethod + def forward(ctx, x, weight, bias): + orig_shape = x.shape + K = weight.shape[1] # input feature dim (contraction dim) + N = weight.shape[0] // 2 # output dim per GLU half + x_2d = x.reshape(-1, K) + M = x_2d.shape[0] + + w_left, w_right = weight.split(N, dim=0) + if bias is not None: + b_left, b_right = bias.split(N, dim=0) + else: + b_left = b_right = None + + out = torch.empty(M, N, device=x.device, dtype=x.dtype) + left = torch.empty(M, N, device=x.device, dtype=x.dtype) + gate = torch.empty(M, N, device=x.device, dtype=x.dtype) + + def grid(meta): + return (triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(N, meta['BLOCK_N']),) + + _fused_linear_softsign_glu_fwd_kernel[grid]( + x_2d, w_left, w_right, b_left, b_right, + out, left, gate, + M, N, K, + triton.next_power_of_2(M), # M_BUCKET: bounds autotune re-runs under variable batch frame counts + x_2d.stride(0), x_2d.stride(1), + w_left.stride(0), w_left.stride(1), + w_right.stride(0), w_right.stride(1), + out.stride(0), out.stride(1), + left.stride(0), left.stride(1), + gate.stride(0), gate.stride(1), + ) + + if x.dim() != 2: + out = out.view(*orig_shape[:-1], N) + + ctx.save_for_backward(x_2d, weight, left, gate) + ctx.orig_x_shape = orig_shape + ctx.N = N + return out + + @staticmethod + def backward(ctx, grad_y): + x, weight, left, gate = ctx.saved_tensors + M, K = x.shape + N = ctx.N + w_left, w_right = weight.split(N, dim=0) + + if grad_y.dim() != 2: + grad_y = grad_y.reshape(-1, N) + if not grad_y.is_contiguous(): + grad_y = grad_y.contiguous() + + # Step 1: Fused element-wise GLU backward (single Triton kernel, + # grad_left_pre/grad_gate computed in registers, one HBM write each) + grad_left_pre = torch.empty(M, N, device=x.device, dtype=x.dtype) + grad_gate = torch.empty(M, N, device=x.device, dtype=x.dtype) + + def elem_grid(meta): + return (triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(N, meta['BLOCK_N']),) + + _softsign_glu_bwd_elem_kernel[elem_grid]( + left, gate, grad_y, + grad_left_pre, grad_gate, + M, N, + left.stride(0), left.stride(1), + gate.stride(0), gate.stride(1), + grad_y.stride(0), grad_y.stride(1), + grad_left_pre.stride(0), grad_left_pre.stride(1), + grad_gate.stride(0), grad_gate.stride(1), + ) + + # Step 2/3: All backward GEMMs on cuBLAS (faster than a Triton GEMM + # here, and preserves fp16/bf16/fp32 dtype without forced casts). + # grad_weight assembled without torch.cat: write both halves into one + # preallocated [2N, K] buffer via out= GEMMs. + grad_weight = torch.empty(2 * N, K, device=x.device, dtype=x.dtype) + torch.mm(grad_left_pre.T, x, out=grad_weight[:N]) + torch.mm(grad_gate.T, x, out=grad_weight[N:]) + grad_bias = torch.cat([grad_left_pre.sum(0), grad_gate.sum(0)], dim=0) + + # grad_x = grad_left_pre @ W_left + grad_gate @ W_right + grad_x = torch.mm(grad_left_pre, w_left) + grad_x.addmm_(grad_gate, w_right) + + if len(ctx.orig_x_shape) != 2: + grad_x = grad_x.view(*ctx.orig_x_shape) + + return grad_x, grad_weight, grad_bias + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def _eager_linear_softsign_glu(x, weight, bias): + """Unfused reference path for unsupported dtypes and devices.""" + linear = F.linear(x, weight, bias) + left, gate = torch.split(linear, linear.shape[-1] // 2, dim=-1) + return left * F.softsign(gate) + + +_FUSED_SUPPORTED_DTYPES = None +_FUSED_FALLBACK_LOG = {} + + +def _fused_supported_dtypes(): + """Dtypes the fused kernel can run on the current GPU. + + fp16 tl.dot: all tensor-core GPUs (Volta sm_70 and newer). + bf16 tl.dot: Ampere (sm_80) and newer — RTX 4090 (sm_89) and + RTX 5090 (sm_120) are fine, but Turing debug GPUs (RTX 20xx) are not. + fp32 falls back to eager: training runs 16-mixed/bf16-mixed, and eager + fp32 keeps full precision without a tf32 surprise inside the kernel. + """ + global _FUSED_SUPPORTED_DTYPES + if _FUSED_SUPPORTED_DTYPES is None: + supported = {torch.float16} + if torch.cuda.get_device_capability() >= (8, 0): + supported.add(torch.bfloat16) + _FUSED_SUPPORTED_DTYPES = supported + return _FUSED_SUPPORTED_DTYPES + + +def fused_linear_softsign_glu(x, weight, bias): + """Fused Linear(C, 2*N) + SoftSignGLU. + + Supports expansion_factor != 1 by splitting weight at its midpoint. + Unsupported dtypes and devices use the eager implementation. + + Args: + x: Input [..., K]. + weight: Linear weight [2*N, K]. + bias: Linear bias [2*N] or None. + + Returns: + Output [..., N]. + """ + if not _TRITON_AVAILABLE or not _fused_capable(): + return _eager_linear_softsign_glu(x, weight, bias) + if not x.is_cuda: + return _eager_linear_softsign_glu(x, weight, bias) + if bias is None: + return _eager_linear_softsign_glu(x, weight, bias) + + fallback_key = (x.dtype, x.shape[-1]) + if x.dtype not in _fused_supported_dtypes(): + if _FUSED_FALLBACK_LOG.get(fallback_key, 0) < 1: + _FUSED_FALLBACK_LOG[fallback_key] = 1 + import warnings + warnings.warn( + f'Fused SoftSignGLU: dtype {x.dtype} not supported for this GPU; ' + f'falling back to eager. (This message is shown once per (dtype, K) pair.)', + stacklevel=2, + ) + return _eager_linear_softsign_glu(x, weight, bias) + # Match weight/bias dtype to input (handles 16-mixed precision where + # weights are fp32 but activations are autocast to fp16) + if weight.dtype != x.dtype: + weight = weight.to(x.dtype) + if bias.dtype != x.dtype: + bias = bias.to(x.dtype) + if not weight.is_contiguous(): + weight = weight.contiguous() + if not bias.is_contiguous(): + bias = bias.contiguous() + return FusedLinearSoftSignGLUFn.apply(x, weight, bias) + + +# --------------------------------------------------------------------------- +# Test +# --------------------------------------------------------------------------- + +def _test(): + """Numerical check against an fp32 reference.""" + import time + + torch.manual_seed(42) + device = 'cuda' + margin = 2.0 + + def rel_err(actual, reference): + return (actual.float() - reference).abs().max().item() / reference.abs().mean().item() + + for K in [256, 512, 1024]: + M = 4096 if K == 256 else (2048 if K == 512 else 1024) + x16 = torch.randn(M, K, device=device, dtype=torch.float16, requires_grad=True) + w16 = torch.randn(2 * K, K, device=device, dtype=torch.float16, requires_grad=True) + b16 = torch.randn(2 * K, device=device, dtype=torch.float16, requires_grad=True) + grad = torch.randn(M, K, device=device, dtype=torch.float16) + + x32 = x16.detach().float().requires_grad_(True) + w32 = w16.detach().float().requires_grad_(True) + b32 = b16.detach().float().requires_grad_(True) + l32, g32 = torch.split(F.linear(x32, w32, b32), K, dim=-1) + ref = l32 * F.softsign(g32) + ref.backward(grad.float()) + + l16, g16 = torch.split(F.linear(x16, w16, b16), K, dim=-1) + y_eager = l16 * F.softsign(g16) + y_eager.backward(grad) + eager_errors = ( + rel_err(y_eager, ref), + rel_err(x16.grad, x32.grad), + rel_err(w16.grad, w32.grad), + rel_err(b16.grad, b32.grad), + ) + + x16.grad = w16.grad = b16.grad = None + y_fused = fused_linear_softsign_glu(x16, w16, b16) + y_fused.backward(grad) + fused_errors = ( + rel_err(y_fused, ref), + rel_err(x16.grad, x32.grad), + rel_err(w16.grad, w32.grad), + rel_err(b16.grad, b32.grad), + ) + + for name, fused_error, eager_error in zip( + ('fwd', 'grad_x', 'grad_w', 'grad_b'), fused_errors, eager_errors + ): + assert fused_error <= eager_error * margin + 1e-6, ( + f'K={K} {name}: fused={fused_error:.4e} vs eager={eager_error:.4e}' + ) + + torch.cuda.synchronize() + t0 = time.time() + for _ in range(50): + fused_linear_softsign_glu(x16, w16, b16) + torch.cuda.synchronize() + fused_t = (time.time() - t0) / 50 + + print( + f"K={K:4d} fwd={fused_errors[0]:.2e} dx={fused_errors[1]:.2e} " + f"dw={fused_errors[2]:.2e} db={fused_errors[3]:.2e} " + f"t={fused_t * 1000:.2f}ms" + ) + + print("\nAll tests passed: fused error is within fp16 rounding of the eager path.") + + +if __name__ == '__main__': + _test() diff --git a/modules/kernels/integration.py b/modules/kernels/integration.py new file mode 100644 index 000000000..7525e08ea --- /dev/null +++ b/modules/kernels/integration.py @@ -0,0 +1,346 @@ +""" +Drop-in replacement for LYNXNet2Block with fused Linear+SoftSignGLU kernels. + +The fused kernel replaces: + nn.Linear(dim, inner_dim*2) + SoftSignGLU → one fused kernel call +(training mode only; eval mode uses the original nn.Sequential path). + +Only softsign_glu is supported — other GLU types are left unpatched +(warning at patch time, block runs the original forward). + +Numerical accuracy: + SoftSignGLU is exact in Triton (no approximation). Differences vs the + eager path are fp16 rounding only (~1e-3 max on unit-scale activations). + +HBM savings (per fused call, M=50000, N=1024, fp16): + Eager: Linear writes [M, 2N] (200 MB), GLU reads [M, 2N] + writes [M, N] + Fused: writes y/left/gate = 3×[M, N] — saves the [M, 2N] round-trip +Backward saves the softsign/denominator intermediates by fusing the +element-wise gradient into one kernel; all GEMMs stay on cuBLAS. + +ONNX export: + Use `model.eval()` → falls back to original path → ONNX export works +""" +import contextlib +import traceback +import warnings + +import torch +import torch.nn as nn + +from modules.backbones.lynxnet2 import LYNXNet2Block +from modules.kernels.fused_linear_softsign_glu import ( + fused_linear_softsign_glu, + is_triton_available, +) + + +_FUSABLE_GLU_TYPES = ('softsign_glu',) + + +class _FusedLYNXNet2BlockMixin: + """Pickle-safe fused forward mixed into an existing LYNXNet2Block.""" + + def forward(self, x): + if not self.training: + return super().forward(x) + + residual = x + x = self.net[0](x) + x = self.net[1](x) + x = self.net[2](x) + x = self.net[3](x) + x = fused_linear_softsign_glu(x, self.net[4].weight, self.net[4].bias) + x = fused_linear_softsign_glu(x, self.net[6].weight, self.net[6].bias) + x = self.net[8](x) + x = self.net[9](x) + return x + residual + + +class FusedLYNXNet2Block(_FusedLYNXNet2BlockMixin, LYNXNet2Block): + """LYNXNet2Block variant with a pickle-safe fused training forward.""" + + +def wrap_lynxnet2_block(block, glu_type='softsign_glu'): + """Wrap an existing LYNXNet2Block to use fused forward. + + Keeps all weights in-place (state_dict compatible). + Only modifies the forward pass. + + Only 'softsign_glu' is fused. Other GLU types are returned unpatched. + + Args: + block: LYNXNet2Block instance + glu_type: GLU type configured for this block + + Returns: + The same block, with patched forward if glu_type is supported. + """ + if glu_type not in _FUSABLE_GLU_TYPES: + warnings.warn( + f"Fused kernels support only {_FUSABLE_GLU_TYPES}; leaving block " + f"with glu_type={glu_type!r} unpatched.", + stacklevel=2, + ) + return block + + net = block.net + if not ( + len(net) == 10 + and isinstance(net[4], nn.Linear) + and isinstance(net[6], nn.Linear) + and isinstance(net[8], nn.Linear) + and net[4].out_features == 2 * net[6].in_features + and net[6].out_features == 2 * net[8].in_features + ): + warnings.warn( + 'Unexpected LYNXNet2Block.net layout; leaving block unpatched.', + stacklevel=2, + ) + return block + + block.__class__ = FusedLYNXNet2Block + return block + + +def patch_lynxnet2_model(model, glu_type='softsign_glu'): + """Patch all LYNXNet2Blocks in a LYNXNet2 model. + + Args: + model: LYNXNet2 instance + glu_type: GLU type configured for the model (only softsign_glu fuses) + + Returns: + Number of blocks patched (0 if glu_type unsupported). + """ + if glu_type not in _FUSABLE_GLU_TYPES: + warnings.warn( + f"Fused kernels require glu_type in {_FUSABLE_GLU_TYPES}; " + f"got {glu_type!r}. Skipping patch.", + stacklevel=2, + ) + return 0 + if not is_triton_available(): + raise RuntimeError( + 'Fused kernels require a working Triton installation. ' + 'Install Triton for this platform or set use_fused_kernels=false.' + ) + patched = 0 + for i, layer in enumerate(model.residual_layers): + if isinstance(layer, LYNXNet2Block): + layer = wrap_lynxnet2_block(layer, glu_type=glu_type) + model.residual_layers[i] = layer + patched += isinstance(layer, FusedLYNXNet2Block) + return patched + + +# --------------------------------------------------------------------------- +# Safe patching — handles both DDPM (denoise_fn) and ReFlow (velocity_fn), +# and checks that the backbone is actually a LYNXNet2 before patching. +# --------------------------------------------------------------------------- + +def _patch_backbone_fn(backbone_fn, glu_type): + """Patch a single backbone function/module if it's a LYNXNet2. + + Args: + backbone_fn: The backbone module (e.g., diffusion.denoise_fn) + glu_type: GLU type (only softsign_glu fuses) + + Returns: + Number of blocks patched (0 if not a LYNXNet2). + """ + from modules.backbones.lynxnet2 import LYNXNet2 + if not isinstance(backbone_fn, LYNXNet2): + return 0 + return patch_lynxnet2_model(backbone_fn, glu_type=glu_type) + + +def _try_patch(module, attr, glu_type): + """Try to patch backbone at module.attr if it's a LYNXNet2. Safe to call + even if attr doesn't exist — returns 0 silently.""" + backbone = getattr(module, attr, None) + if backbone is None: + return 0 + return _patch_backbone_fn(backbone, glu_type) + + +def patch_diffusion_module(diffusion, glu_type='softsign_glu'): + """Patch a diffusion module's backbone (DDPM or ReFlow). + + Handles both: + GaussianDiffusion / PitchDiffusion / MultiVarianceDiffusion → .denoise_fn + RectifiedFlow / PitchRectifiedFlow / MultiVarianceRectifiedFlow → .velocity_fn + + Returns: + Number of blocks patched. + """ + return ( + _try_patch(diffusion, 'denoise_fn', glu_type) + + _try_patch(diffusion, 'velocity_fn', glu_type) + ) + + +# --------------------------------------------------------------------------- +# Warmup — trigger Triton autotune before training starts +# --------------------------------------------------------------------------- + +def warmup_fused_backbone(backbone, max_frames=None, autocast_dtype=None): + """Run dummy forward passes to trigger Triton autotune compilation + for all fused kernels (fwd + bwd elem). Call after patching, before + the first real training step (model must already be on its CUDA device). + + Only forward is executed (``torch.no_grad``) — the element-wise backward + kernel's autotune key depends on ``N`` (a single fixed value per model), + so its one-off compile cost is paid on the first real step instead. + + Autotune timings are cached in process memory only (Triton persists + compiled binaries to disk, but re-runs the config benchmark per process), + so this runs once per training process. The forward kernel's autotune + key buckets M by next_power_of_2, so we sweep the power-of-two buckets + a real run will hit: from a small bucket up to next_power_of_2(max_frames). + + Args: + backbone: LYNXNet2 model (already patched). + max_frames: max total frames per batch (hparams['max_batch_frames']). + If None, warms a single small bucket only. + autocast_dtype: torch.float16 for '16-mixed', torch.bfloat16 for + 'bf16-mixed'. If None, no autocast — with fp32 parameters the + fused path falls back to eager and the warmup is a no-op. + """ + device = next(backbone.parameters()).device + dtype = next(backbone.parameters()).dtype + + if device.type != 'cuda': + return 0 + if not is_triton_available(): + raise RuntimeError( + 'Fused kernel warmup requires a working Triton installation. ' + 'Install Triton for this platform or set use_fused_kernels=false.' + ) + + import triton + + # cond hidden size from the conditioner projection (Linear or Conv1d) + proj = backbone.conditioner_projection + hidden = getattr(proj, 'in_features', None) or proj.in_channels + + B = 4 + # Sweep M buckets: 2048 up to next_power_of_2(max_frames) + if max_frames is not None: + top = triton.next_power_of_2(int(max_frames)) + bucket = 2048 + t_list = [] + while bucket <= top: + # M = B * T lands in this bucket (M just above the previous bucket) + t_list.append(bucket // B // 2 + 1) + bucket *= 2 + else: + t_list = [500] + + ac_factory = ( + (lambda: torch.autocast(device_type=device.type, dtype=autocast_dtype)) + if autocast_dtype is not None else contextlib.nullcontext + ) + # Fork the RNG so dummy inputs do not advance the training noise stream. + with torch.random.fork_rng(devices=[device]): + for T in t_list: + # spec shape: [B, n_feats, in_dims, T] + spec = torch.randn(B, backbone.n_feats, backbone.in_dims, T, + device=device, dtype=dtype) + t = torch.randint(0, 1000, (B,), device=device).float() + cond = torch.randn(B, hidden, T, device=device, dtype=dtype) + + try: + with torch.no_grad(): + with ac_factory(): + backbone(spec, t, cond=cond) + except Exception as e: # noqa: BLE001 - warmup must remain non-fatal + # Autotune failure should not crash training — Triton cache + # can be built on the first real step instead. + warnings.warn( + f'Fused kernel warmup skipped at T={T} ' + f'({type(e).__name__}: {e})\n{traceback.format_exc()}', + stacklevel=2, + ) + break + finally: + del spec, cond + torch.cuda.empty_cache() + return len(t_list) + + +def warmup_fused_backbones(backbones, max_frames, precision): + """Warm all patched backbones using Lightning's effective precision.""" + precision = str(precision) + autocast_dtype = ( + torch.float16 if '16' in precision and 'bf16' not in precision + else torch.bfloat16 if 'bf16' in precision + else None + ) + if autocast_dtype is None: + from lightning.pytorch.utilities.rank_zero import rank_zero_info + rank_zero_info( + 'Fused kernels: precision=%s has no autocast dtype; ' + 'fused kernel will fall back to eager at runtime.', precision + ) + for backbone in backbones: + warmup_fused_backbone( + backbone, + max_frames=max_frames, + autocast_dtype=autocast_dtype, + ) + + +# --------------------------------------------------------------------------- +# Test +# --------------------------------------------------------------------------- + +def _test(): + import torch + from modules.backbones.lynxnet2 import LYNXNet2Block + + device = 'cuda' + torch.manual_seed(42) + + # Create a single block + block = LYNXNet2Block(dim=256, expansion_factor=1, glu_type='softsign_glu').to(device).half() + + # Copy weights + block_ref = LYNXNet2Block(dim=256, expansion_factor=1, glu_type='softsign_glu').to(device).half() + block_ref.load_state_dict(block.state_dict()) + + # Patch + wrap_lynxnet2_block(block, glu_type='softsign_glu') + + B, T = 2, 500 + x = torch.randn(B, T, 256, device=device, dtype=torch.float16) + + # Forward + out_orig = block_ref(x) + out_fused = block(x) + + fwd_diff = (out_fused - out_orig).abs().max().item() + print(f"Block forward max diff: {fwd_diff:.4e}") + + # Backward + grad = torch.randn_like(out_orig) + out_orig.backward(grad) + grads_ref = {n: p.grad.clone() for n, p in block_ref.named_parameters() if p.grad is not None} + + for p in block.parameters(): + p.grad = None + + out_fused = block(x) + out_fused.backward(grad) + grads_fused = {n: p.grad.clone() for n, p in block.named_parameters() if p.grad is not None} + + max_w_diff = max( + (grads_fused[n] - grads_ref[n]).abs().max().item() + for n in grads_ref + ) + print(f"Block weight grad max diff: {max_w_diff:.4e}") + print("\nIntegration works! Use model.eval() for ONNX export fallback.") + + +if __name__ == '__main__': + _test() diff --git a/training/acoustic_task.py b/training/acoustic_task.py index ca6a71c65..1f2e4f362 100644 --- a/training/acoustic_task.py +++ b/training/acoustic_task.py @@ -94,6 +94,38 @@ def __init__(self): self.required_variances.append('tension') super()._finish_init() + # ── Fuse LYNXNet2 backbone kernels (in-place) ── + # Only SoftSignGLU backbones are patched. + self._fused_kernels_patched = 0 + if hparams.get('use_fused_kernels', False): + try: + from modules.kernels.integration import patch_diffusion_module + from lightning.pytorch.utilities.rank_zero import rank_zero_info + # NOTE: LYNXNet2 defaults to swiglu when glu_type is unset + self._fused_kernels_patched = patch_diffusion_module( + self.model.diffusion, + glu_type=hparams['backbone_args'].get('glu_type', 'swiglu'), + ) + rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks', self._fused_kernels_patched) + except ImportError as e: + from lightning.pytorch.utilities.rank_zero import rank_zero_info + rank_zero_info('Fused kernels unavailable (ImportError: %s); running eager.', e) + + def on_fit_start(self): + # Warm Triton autotune caches after the model is on its CUDA device, + # so the first training steps don't pay the per-bucket benchmark cost. + if self._fused_kernels_patched > 0 and self.device.type == 'cuda': + from modules.kernels.integration import warmup_fused_backbones + backbones = [ + backbone for attr in ('denoise_fn', 'velocity_fn') + if (backbone := getattr(self.model.diffusion, attr, None)) is not None + ] + warmup_fused_backbones( + backbones, + max_frames=hparams['max_batch_frames'], + precision=self.trainer.precision, + ) + def _build_model(self): return DiffSingerAcoustic( vocab_size=len(self.phoneme_dictionary), diff --git a/training/variance_task.py b/training/variance_task.py index 646d9540a..032acfc72 100644 --- a/training/variance_task.py +++ b/training/variance_task.py @@ -115,6 +115,53 @@ def __init__(self): self.lambda_var_loss = hparams['lambda_var_loss'] super()._finish_init() + # ── Fuse LYNXNet2 backbone kernels (in-place) ── + self._fused_kernels_patched = 0 + self._fused_kernel_backbones = [] + if hparams.get('use_fused_kernels', False): + try: + from modules.backbones.lynxnet2 import LYNXNet2 + from modules.kernels.integration import patch_diffusion_module + from lightning.pytorch.utilities.rank_zero import rank_zero_info + # Each predictor has its own backbone config; patch only the ones + # actually configured with softsign_glu (others are skipped with + # a warning instead of silently changing their math). + # NOTE: LYNXNet2 defaults to swiglu when glu_type is unset. + for predictor_attr, args_key in ( + ('pitch_predictor', 'pitch_prediction_args'), + ('variance_predictor', 'variances_prediction_args'), + ): + predictor = getattr(self.model, predictor_attr, None) + if predictor is None: + continue + glu = (hparams.get(args_key) or {}).get('backbone_args', {}).get('glu_type', 'swiglu') + n = patch_diffusion_module(predictor, glu_type=glu) + self._fused_kernels_patched += n + if n > 0: + for attr in ('denoise_fn', 'velocity_fn'): + backbone = getattr(predictor, attr, None) + if isinstance(backbone, LYNXNet2): + self._fused_kernel_backbones.append(backbone) + rank_zero_info( + 'Fused kernels: patched %d LYNXNet2 blocks in %s (glu_type=%s)', + n, predictor_attr, glu + ) + except ImportError as e: + from lightning.pytorch.utilities.rank_zero import rank_zero_info + rank_zero_info('Fused kernels unavailable (ImportError: %s); running eager.', e) + + def on_fit_start(self): + # Warm Triton autotune caches after the model is on its CUDA device, + # so the first training steps don't pay the per-bucket benchmark cost. + # Mirrors AcousticTask.on_fit_start, but sweeps both predictors. + if self._fused_kernels_patched > 0 and self.device.type == 'cuda': + from modules.kernels.integration import warmup_fused_backbones + warmup_fused_backbones( + self._fused_kernel_backbones, + max_frames=hparams['max_batch_frames'], + precision=self.trainer.precision, + ) + def _build_model(self): return DiffSingerVariance( vocab_size=len(self.phoneme_dictionary),