Как реализовать активацию SwiGLU? Почему SwiGLU принимает два тензора? ⇐ Python

[Anonymous]

Вариант SwiGLU, представленный в https://arxiv.org/pdf/2002.05202, представляет собой просто «божественное благоволение», а реализация Flash-Attention работает «из коробки» https://github.com/Dao-AILab/ flash-attention/tree/main
Сравнение реализации:

Код: Выделить всё

import torch

swiglu_fwd_codestring = """
template  T swiglu_fwd(T x, T y) {
return float(x) * float(y) / (1.0f + ::exp(-float(x)));
}
"""

swiglu_bwd_codestring = """
template  T swiglu_bwd(T x, T y, T g, T& dx, T& dy) {
float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x)));
dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y);
dy = float(x) * x_sigmoid * float(g);
}
"""
swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring)
swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2)

class SwiGLUFunction(torch.autograd.Function):

@staticmethod
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
return swiglu_fwd(x, y)

@staticmethod
def backward(ctx, dout):
x, y = ctx.saved_tensors
return swiglu_bwd(x, y, dout)

swiglu = SwiGLUFunction.apply

swiglu(torch.tensor(1.0).to('cuda'), torch.tensor(0.5).to('cuda'))

и

Код: Выделить всё

import torch.nn.functional as F
swiglu_torch = lambda x, y: F.silu(x) * y

swiglu_torch(torch.tensor(1.0).to('cuda'), torch.tensor(0.5).to('cuda'))

Они оба дают одинаковый результат:

Код: Выделить всё

tensor(0.3655, device='cuda:0')

Описание бумаги, входные данные для активации – один тензор x,


но обе реализации принимают x и y, < strong>Чему соответствует y, когда у нас есть только тензор x, как во входных данных в формуле? Это x = W1.x и y=W2.x?

Подробнее здесь: https://stackoverflow.com/questions/790 ... wo-tensors