Skip to content

fix rules for fused swiglu scale clamp both forward and backward #698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wanghuancoder merged 2 commits into from
Jul 27, 2026
Merged

fix rules for fused swiglu scale clamp both forward and backward #698

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e2f7f7f
fix rules for fused swiglu scale clamp both forward and backward
greenhandF Jul 27, 2026
84ca5ea
Merge branch 'main' into fixrules
greenhandF Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 68 additions & 59 deletions tester/paddle_to_torch/rules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8253,80 +8253,89 @@ def apply(self, paddle_api: str) -> ConvertResult:

elif op_name == "fused_swiglu_scale_clamp":
# _run_custom_op("fused_swiglu_scale_clamp", x, scale, max_val)
# fwd: split x -> (x1, x2); out = clamp(silu(x1) * x2 * scale, -max_val, max_val)
# Paddle custom op returns a list; we return [out] to match.
# clamp 语义对齐 PaddleFleet fusions/fused_swiglu_scale.py::
# fused_swiglu_scale_forward (gate 只截上限, value 对称截断),
# 精度顺序对齐 CUDA kernel VectorizedFusedSwiGLUFwd 以做到 bit 级一致:
# g = min(gate, cv); v = clamp(value, -cv, cv)
# out = (T)((g * sigmoid(g)) * v * s) # 全程 fp32, 只在最后舍入一次
# 注意: fleet 的 CPU/XPU fallback 写成
# (silu(g)*v).cast(x.dtype) * scale.cast(x.dtype)
# 比 kernel 多两次 bf16 舍入, 与 CUDA 算子本身不 bit 一致, 故这里按 kernel 写。
x = arg1 # shape [..., 2D]
scale = arg2 # scalar or tensor [..., 1] or [...,]
max_val = arg3 # scalar
_cv = float(max_val)
_hidden = x.shape[-1] // 2
if x.shape[-1] == 0:
result = [torch.empty((*x.shape[:-1], 0), dtype=x.dtype, device=x.device)]
elif x.shape[-1] % 2 != 0:
raise ValueError(
f"input shape is invalid for input of size {tuple(x.shape)}: "
"fused_swiglu_scale_clamp requires an even last dimension"
)
_x_fp32 = x.to(torch.float32)
_gate = torch.clamp(_x_fp32[..., :_hidden], max=_cv)
_val = torch.clamp(_x_fp32[..., _hidden:], min=-_cv, max=_cv)
if torch.is_tensor(scale):
_scale_exp = scale.to(torch.float32)
else:
_outer = x.numel() // (_hidden * 2)
_x_2d = x.reshape(_outer, _hidden * 2)
_out = torch.empty((_outer, _hidden), dtype=x.dtype, device=x.device)
if torch.is_tensor(scale):
scale = scale.view(*scale.shape[:-1], 1) if scale.dim() > 1 else scale.unsqueeze(-1)
_workspace_bytes = 32 << 30
_bytes_per_row = max(1, _hidden * (4 * 4 + x.element_size()))
_row_chunk = max(1, min(_outer, _workspace_bytes // _bytes_per_row))
with torch.no_grad():
for _row_start in range(0, _outer, _row_chunk):
_row_end = min(_outer, _row_start + _row_chunk)
_lhs = _x_2d[_row_start:_row_end, :_hidden].to(torch.float32)
_rhs = _x_2d[_row_start:_row_end, _hidden:].to(torch.float32)
_lhs.mul_(torch.sigmoid(_lhs)).mul_(_rhs)
_scale_chunk = scale[_row_start:_row_end] if torch.is_tensor(scale) and scale.numel() > 1 else scale
_lhs.mul_(_scale_chunk).clamp_(min=-float(max_val), max=float(max_val))
_out[_row_start:_row_end] = _lhs.to(x.dtype)
result = [_out.reshape(*x.shape[:-1], _hidden)]
_scale_exp = torch.tensor(float(scale), dtype=torch.float32, device=x.device)
while _scale_exp.dim() < _gate.dim():
_scale_exp = _scale_exp.unsqueeze(-1)
_swiglu = (_gate * torch.sigmoid(_gate)) * _val
result = [(_swiglu * _scale_exp).to(x.dtype)]


elif op_name == "fused_swiglu_scale_clamp_bwd":
# _run_custom_op("fused_swiglu_scale_clamp_bwd", x, scale, dy, max_val)
# arg1=x (original fwd input [..., 2D]), arg2=scale, arg3=dy (grad of output [..., D]),
# arg4=max_val. Returns [dx, d_scale] to match Paddle output.
# clamp 语义对齐 PaddleFleet fusions/fused_swiglu_scale.py::
# fused_swiglu_scale_backward, 精度/舍入点对齐 CUDA kernel
# VectorizedFusedSwiGLUBwd (kHasClamp=true):
# g_eff = min(g, cv); v_eff = clamp(v, -cv, cv)
# g_mask = (g <= cv); v_mask = (-cv <= v <= cv) # fp32 掩码, 边界处梯度通过
# d_u = dout * s # 全程 fp32
# d_v = (T)(d_u * silu_g * v_mask)
# d_g = (T)(d_u * sig * (1 + g_eff*(1-sig)) * v_eff * g_mask)
# d_scale = (ScaleT)sum_fp32( (T)swiglu * (ScaleT)dout )
# kernel 特意把 fp32 swiglu 先压回 x.dtype 再与 dout 相乘, 用 fp32 累加
x = arg1 # shape [..., 2D] (original forward input)
scale = arg2 # scalar or tensor [..., 1]
dy = arg3 # shape [..., D] (gradient of forward output)
max_val = arg4 # scalar
_hidden = dy.shape[-1]
_outer = dy.numel() // _hidden
_x_2d = x.reshape(_outer, _hidden * 2)
_dy_2d = dy.reshape(_outer, _hidden)
_dx = torch.empty_like(_x_2d)
_d_scale = torch.empty((_outer, 1), dtype=torch.float32, device=x.device)
_cv = float(max_val)
_hidden = x.shape[-1] // 2
_x_fp32 = x.to(torch.float32)
_gate_raw = _x_fp32[..., :_hidden]
_val_raw = _x_fp32[..., _hidden:]
_gate = torch.clamp(_gate_raw, max=_cv)
_val = torch.clamp(_val_raw, min=-_cv, max=_cv)
# kernel 中掩码是 float, 保持 fp32 避免额外的 dtype 提升
_g_mask = (_gate_raw <= _cv).to(torch.float32)
_v_mask = ((_val_raw <= _cv) & (_val_raw >= -_cv)).to(torch.float32)
_sig = torch.sigmoid(_gate)
_silu = _gate * _sig
_swiglu_val = _silu * _val
if torch.is_tensor(scale):
scale_v = scale.view(*scale.shape[:-1], 1) if scale.dim() > 1 else scale.unsqueeze(-1)
_scale_dtype = scale.dtype
_scale_exp = scale.to(torch.float32)
else:
scale_v = float(scale)
_workspace_bytes = 32 << 30
_bytes_per_row = max(1, _hidden * (4 * 10 + x.element_size() * 2))
_row_chunk = max(1, min(_outer, _workspace_bytes // _bytes_per_row))
with torch.no_grad():
for _row_start in range(0, _outer, _row_chunk):
_row_end = min(_outer, _row_start + _row_chunk)
_lhs = _x_2d[_row_start:_row_end, :_hidden].to(torch.float32)
_rhs = _x_2d[_row_start:_row_end, _hidden:].to(torch.float32)
_dy = _dy_2d[_row_start:_row_end].to(torch.float32)
_scale_chunk = scale_v[_row_start:_row_end] if torch.is_tensor(scale_v) and scale_v.numel() > 1 else scale_v
_sigmoid = torch.sigmoid(_lhs)
_silu = _lhs * _sigmoid
_pre_clamp = _silu * _rhs * _scale_chunk
_dy.mul_((_pre_clamp.abs() < float(max_val)).to(torch.float32))
_d_scale[_row_start:_row_end] = (_dy * _silu * _rhs).sum(dim=-1, keepdim=True)
_dx[_row_start:_row_end, _hidden:] = (_dy * _silu * _scale_chunk).to(x.dtype)
_lhs.sub_(_silu).add_(1.0).mul_(_sigmoid).mul_(_rhs).mul_(_dy).mul_(_scale_chunk)
_dx[_row_start:_row_end, :_hidden] = _lhs.to(x.dtype)
dx = _dx.reshape(x.shape)
d_scale = _d_scale.reshape(*x.shape[:-1], 1)
if torch.is_tensor(scale):
d_scale = d_scale.to(scale.dtype)
_scale_dtype = torch.float32
_scale_exp = torch.tensor(float(scale), dtype=torch.float32, device=x.device)
while _scale_exp.dim() < dy.dim():
_scale_exp = _scale_exp.unsqueeze(-1)
_d_u = dy * _scale_exp
_d_val = _d_u * _silu * _v_mask
# kernel 里 1.0f + g_eff*(1.0f-sig) 会被 nvcc 收缩成单条 FMA(只舍入一次),
# 用 addcmul 走同一条 fused 路径, 而不是 mul+add 两次舍入。
_d_gate = (
_d_u
* _sig
* torch.addcmul(torch.ones((), dtype=torch.float32, device=x.device), _gate, 1.0 - _sig)
* _val
* _g_mask
)
dx = torch.cat([_d_gate, _d_val], dim=-1).to(x.dtype)
# d_scale: fp32 swiglu 先压回 x.dtype, dout 转 scale dtype, 乘积用 fp32 累加,
# 最后一次性 cast 回 scale dtype(与 kernel 的 shared float 归约一致)
d_scale = (
(_swiglu_val.to(x.dtype) * dy.to(_scale_dtype))
.to(torch.float32)
.sum(dim=-1, keepdim=True)
.to(_scale_dtype)
)
result = [dx, d_scale]

elif op_name in ("fused_swiglu_probs_bwd", "paddlefleet_fused_swiglu_probs_bwd"):
Expand Down
Loading