mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
## Summary Resolves #936. Adds Liger SwiGLU support for Falcon H1's MLP block, which pre-scales the gate pre-activation and post-scales the down-projection output: y = down_proj(silu(gate_proj(x) * gate_mult) * up_proj(x)) * down_mult Previously `apply_liger_kernel_to_falcon_h1(swiglu=True)` was a no-op that just logged a warning; now it actually patches the MLP. ## Details - `src/liger_kernel/ops/swiglu.py` — forward/backward Triton kernels extended with a `gate_multiplier` scalar applied to the gate pre-activation in fp32 before silu. Backward chain rule pulls an extra `gate_multiplier` factor into `da`. `LigerSiLUMulFunction` accepts `gate_multiplier` and `down_multiplier` as optional args (defaults 1.0), so all existing callers (Llama, Qwen, Granite, Phi3, Olmo, Hunyuan, etc.) are unaffected. `down_multiplier` is applied outside the kernel in the autograd function — a single elementwise multiply on the output, near-zero overhead. - `src/liger_kernel/transformers/swiglu.py` — new `LigerFalconH1SwiGLUMLP` reads `config.mlp_multipliers` and forwards them to `LigerSiLUMulFunction`. Falls back to `self.config.mlp_multipliers` when patched onto an already-instantiated HF module via `_patch_swiglu_module` (which only rebinds `forward`). - `src/liger_kernel/transformers/monkey_patch.py` — flip `apply_liger_kernel_to_falcon_h1`'s `swiglu` default from `False` to `True`, wire it to `LigerFalconH1SwiGLUMLP`, and fix the decoder-layer attribute to `feed_forward` (Falcon H1's actual HF attribute name). - `src/liger_kernel/transformers/__init__.py` — export `LigerFalconH1SwiGLUMLP`. - `test/transformers/test_swiglu.py` — three new parametrized tests (21 total cases): - `test_correctness_silumul_with_multipliers` — 12 cases (3 multiplier combos × 2 shapes × 2 dtypes) validating kernel math against a pure PyTorch reference. - `test_silumul_default_multipliers_backward_compat` — confirms no-arg `LigerSiLUMulFunction.apply(a, b)` is bit-identical to explicit `(1.0, 1.0)`. - `test_correctness_falcon_h1_mlp` — 8 cases end-to-end parity for `LigerFalconH1SwiGLUMLP` against a pure-PyTorch Falcon H1 MLP reference. Out of scope: Qwen3-VL `swiglu=True` no-op (#956) — separate monkey-patch wiring issue, not kernel math. ## Testing Done - Hardware Type: NVIDIA A10G (g5.xlarge) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Notes on `make test`: - All 21 new tests pass. - All pre-existing SwiGLU tests that exercise `LigerSiLUMulFunction` pass (Llama, Phi3, Mixtral fp32, functional, falcon_h1 monkey_patch). - Unrelated pre-existing failures on `main` that also show up here (confirmed by running the same tests on `main`): - `test_correctness_mixtralexperts[bf16]` — Triton 3.2 `atomic_add does not support bf16` in `src/liger_kernel/ops/fused_moe.py`. - `test_apply_liger_kernel_to_instance_for_paligemma` — `SiglipVisionModel` attribute change in transformers v5. - `test_apply_liger_kernel_to_instance_for_gemma3_conditional_generation` — `Gemma3Config` v5 API change. - `test_grpo_loss::test_correctness[sequence-luspo-...*-8-128-1024-4096]` (3 large-shape cases) — pre-existing GRPO numerical sensitivity. Convergence tests were not run locally (single-GPU box). Happy to run them if a maintainer can advise on the expected workflow, or to add a Falcon H1 mini-model entry in a follow-up.
This commit is contained in:
Shivam Sharma
@@ -12,7 +12,9 @@ def silu(x):
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _swiglu_forward_kernel(a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr):
|
||||
def _swiglu_forward_kernel(
|
||||
a_ptr, b_ptr, c_ptr, stride, gate_multiplier, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
|
||||
):
|
||||
program_id = tl.program_id(0).to(tl.int64)
|
||||
|
||||
# locate start index
|
||||
@@ -24,14 +26,16 @@ def _swiglu_forward_kernel(a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BL
|
||||
mask = col_offsets < n_cols
|
||||
|
||||
# sigmoid requires type float32
|
||||
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
|
||||
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier
|
||||
b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
|
||||
c_row = silu(a_row).cast(b_row.dtype) * b_row
|
||||
tl.store(c_ptr + col_offsets, c_row, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _swiglu_backward_kernel(dc_ptr, a_ptr, b_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr):
|
||||
def _swiglu_backward_kernel(
|
||||
dc_ptr, a_ptr, b_ptr, stride, gate_multiplier, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
|
||||
):
|
||||
program_id = tl.program_id(0).to(tl.int64)
|
||||
|
||||
# locate start index
|
||||
@@ -44,20 +48,21 @@ def _swiglu_backward_kernel(dc_ptr, a_ptr, b_ptr, stride, n_cols: tl.constexpr,
|
||||
|
||||
dc_row = tl.load(dc_ptr + col_offsets, mask=mask, other=0)
|
||||
# sigmoid requires type float32
|
||||
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
|
||||
a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32) * gate_multiplier
|
||||
b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
|
||||
|
||||
# recomputation to save memory
|
||||
# recomputation to save memory. a_row already holds a * gate_multiplier.
|
||||
sig_a = tl.sigmoid(a_row)
|
||||
silu_a = a_row * sig_a
|
||||
db_row = dc_row * silu_a
|
||||
da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row
|
||||
# chain rule pulls an extra factor of gate_multiplier through the pre-activation scaling
|
||||
da_row = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row * gate_multiplier
|
||||
|
||||
tl.store(a_ptr + col_offsets, da_row, mask=mask)
|
||||
tl.store(b_ptr + col_offsets, db_row, mask=mask)
|
||||
|
||||
|
||||
def swiglu_forward(a, b):
|
||||
def swiglu_forward(a, b, gate_multiplier: float = 1.0):
|
||||
ori_shape = a.shape
|
||||
|
||||
n_cols = ori_shape[-1]
|
||||
@@ -73,6 +78,7 @@ def swiglu_forward(a, b):
|
||||
b,
|
||||
c,
|
||||
c.stride(-2),
|
||||
float(gate_multiplier),
|
||||
n_cols=n_cols,
|
||||
BLOCK_SIZE=BLOCK_SIZE,
|
||||
num_warps=num_warps,
|
||||
@@ -80,7 +86,7 @@ def swiglu_forward(a, b):
|
||||
return a, b, c.view(*ori_shape)
|
||||
|
||||
|
||||
def swiglu_backward(a, b, dc):
|
||||
def swiglu_backward(a, b, dc, gate_multiplier: float = 1.0):
|
||||
ori_shape = dc.shape
|
||||
n_cols = ori_shape[-1]
|
||||
dc = dc.view(-1, n_cols)
|
||||
@@ -93,6 +99,7 @@ def swiglu_backward(a, b, dc):
|
||||
a,
|
||||
b,
|
||||
dc.stride(-2),
|
||||
float(gate_multiplier),
|
||||
n_cols=n_cols,
|
||||
BLOCK_SIZE=BLOCK_SIZE,
|
||||
num_warps=num_warps,
|
||||
@@ -103,7 +110,12 @@ def swiglu_backward(a, b, dc):
|
||||
class LigerSiLUMulFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
@ensure_contiguous
|
||||
def forward(ctx, a, b):
|
||||
def forward(ctx, a, b, gate_multiplier: float = 1.0, down_multiplier: float = 1.0):
|
||||
gate_multiplier = float(gate_multiplier)
|
||||
down_multiplier = float(down_multiplier)
|
||||
ctx.gate_multiplier = gate_multiplier
|
||||
ctx.down_multiplier = down_multiplier
|
||||
|
||||
if isinstance(a, torch.distributed.tensor.DTensor) or isinstance(b, torch.distributed.tensor.DTensor):
|
||||
device_mesh, placements = (
|
||||
(a.device_mesh, a.placements)
|
||||
@@ -117,12 +129,16 @@ class LigerSiLUMulFunction(torch.autograd.Function):
|
||||
a = torch.distributed.tensor.distribute_tensor(a, device_mesh=device_mesh, placements=placements)
|
||||
if not isinstance(b, torch.distributed.tensor.DTensor):
|
||||
b = torch.distributed.tensor.distribute_tensor(b, device_mesh=device_mesh, placements=placements)
|
||||
a_local, b_local, c_local = swiglu_forward(a.to_local(), b.to_local())
|
||||
a_local, b_local, c_local = swiglu_forward(a.to_local(), b.to_local(), gate_multiplier)
|
||||
if down_multiplier != 1.0:
|
||||
c_local = c_local * down_multiplier
|
||||
ctx.save_for_backward(a_local, b_local)
|
||||
ctx.dtensor_metadata = (device_mesh, placements)
|
||||
return torch.distributed.tensor.DTensor.from_local(c_local, device_mesh, placements)
|
||||
else:
|
||||
a, b, c = swiglu_forward(a, b)
|
||||
a, b, c = swiglu_forward(a, b, gate_multiplier)
|
||||
if down_multiplier != 1.0:
|
||||
c = c * down_multiplier
|
||||
ctx.save_for_backward(a, b)
|
||||
ctx.dtensor_metadata = None
|
||||
return c
|
||||
@@ -131,6 +147,9 @@ class LigerSiLUMulFunction(torch.autograd.Function):
|
||||
@ensure_contiguous
|
||||
def backward(ctx, dc):
|
||||
a, b = ctx.saved_tensors
|
||||
gate_multiplier = ctx.gate_multiplier
|
||||
down_multiplier = ctx.down_multiplier
|
||||
|
||||
if ctx.dtensor_metadata is not None:
|
||||
device_mesh, placements = ctx.dtensor_metadata
|
||||
|
||||
@@ -141,11 +160,17 @@ class LigerSiLUMulFunction(torch.autograd.Function):
|
||||
if isinstance(dc, torch.distributed.tensor.DTensor)
|
||||
else torch.distributed.tensor.distribute_tensor(dc, device_mesh=device_mesh, placements=placements)
|
||||
)
|
||||
a_local, b_local = swiglu_backward(a, b, dc_local)
|
||||
if down_multiplier != 1.0:
|
||||
dc_local = dc_local * down_multiplier
|
||||
a_local, b_local = swiglu_backward(a, b, dc_local, gate_multiplier)
|
||||
return (
|
||||
torch.distributed.tensor.DTensor.from_local(a_local, device_mesh, placements),
|
||||
torch.distributed.tensor.DTensor.from_local(b_local, device_mesh, placements),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
|
||||
a, b = swiglu_backward(a, b, dc)
|
||||
return a, b
|
||||
if down_multiplier != 1.0:
|
||||
dc = dc * down_multiplier
|
||||
a, b = swiglu_backward(a, b, dc, gate_multiplier)
|
||||
return a, b, None, None
|
||||
|
||||
@@ -24,6 +24,7 @@ from liger_kernel.transformers.softmax import LigerSoftmax # noqa: F401
|
||||
from liger_kernel.transformers.sparsemax import LigerSparsemax # noqa: F401
|
||||
from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP # noqa: F401
|
||||
from liger_kernel.transformers.swiglu import LigerExperts # noqa: F401
|
||||
from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP # noqa: F401
|
||||
from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP # noqa: F401
|
||||
from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP # noqa: F401
|
||||
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP # noqa: F401
|
||||
@@ -180,6 +181,7 @@ __all__ = [
|
||||
"liger_llama4_text_rotary_pos_emb",
|
||||
"liger_llama4_vision_rotary_pos_emb",
|
||||
"LigerBlockSparseTop2MLP",
|
||||
"LigerFalconH1SwiGLUMLP",
|
||||
"LigerPhi3SwiGLUMLP",
|
||||
"LigerQwen3MoeSwiGLUMLP",
|
||||
"LigerSwiGLUMLP",
|
||||
|
||||
@@ -2773,7 +2773,7 @@ def apply_liger_kernel_to_falcon_h1(
|
||||
cross_entropy: bool = False,
|
||||
fused_linear_cross_entropy: bool = True,
|
||||
rms_norm: bool = True,
|
||||
swiglu: bool = False,
|
||||
swiglu: bool = True,
|
||||
model: PreTrainedModel = None,
|
||||
) -> None:
|
||||
"""
|
||||
@@ -2798,6 +2798,8 @@ def apply_liger_kernel_to_falcon_h1(
|
||||
from transformers.models.falcon_h1 import modeling_falcon_h1
|
||||
from transformers.models.falcon_h1.modeling_falcon_h1 import FalconH1Model
|
||||
|
||||
from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP
|
||||
|
||||
if rope:
|
||||
logger.info("Apply liger rotary pos emb.")
|
||||
modeling_falcon_h1.apply_rotary_pos_emb = liger_rotary_pos_emb
|
||||
@@ -2805,7 +2807,8 @@ def apply_liger_kernel_to_falcon_h1(
|
||||
logger.info("Apply liger RMSNorm")
|
||||
modeling_falcon_h1.FalconH1RMSNorm = LigerRMSNorm
|
||||
if swiglu:
|
||||
logger.warning("LigerSwiGLUMLP is not available for Falcon-H1 models. There will be no effect.")
|
||||
logger.info("Apply liger SwiGLU")
|
||||
modeling_falcon_h1.FalconH1MLP = LigerFalconH1SwiGLUMLP
|
||||
|
||||
if cross_entropy:
|
||||
logger.info("Apply liger cross entropy")
|
||||
@@ -2831,7 +2834,7 @@ def apply_liger_kernel_to_falcon_h1(
|
||||
|
||||
for decoder_layer in base_model.layers:
|
||||
if swiglu:
|
||||
_patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)
|
||||
_patch_swiglu_module(decoder_layer.feed_forward, LigerFalconH1SwiGLUMLP)
|
||||
if rms_norm:
|
||||
_patch_rms_norm_module(decoder_layer.input_layernorm)
|
||||
_patch_rms_norm_module(decoder_layer.pre_ff_layernorm)
|
||||
|
||||
@@ -135,3 +135,41 @@ class LigerHunyuanV1SwiGLUMLP(nn.Module):
|
||||
|
||||
def forward(self, x):
|
||||
return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
|
||||
|
||||
|
||||
class LigerFalconH1SwiGLUMLP(nn.Module):
|
||||
"""
|
||||
Patch FalconH1MLP to use LigerSiLUMulFunction with gate / down multipliers.
|
||||
|
||||
Falcon H1's MLP block pre-scales the gate pre-activation and post-scales the
|
||||
down projection output:
|
||||
|
||||
y = down_proj(silu(gate_proj(x) * gate_mult) * up_proj(x)) * down_mult
|
||||
|
||||
https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon_h1/modeling_falcon_h1.py
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
bias = getattr(config, "mlp_bias", False)
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
|
||||
if config.hidden_act not in ["silu", "swish"]:
|
||||
raise ValueError(f"Activation function {config.hidden_act} not supported.")
|
||||
gate_multiplier, down_multiplier = config.mlp_multipliers
|
||||
self.gate_multiplier = float(gate_multiplier)
|
||||
self.down_multiplier = float(down_multiplier)
|
||||
|
||||
def forward(self, x):
|
||||
return self.down_proj(
|
||||
LigerSiLUMulFunction.apply(
|
||||
self.gate_proj(x),
|
||||
self.up_proj(x),
|
||||
float(self.gate_multiplier),
|
||||
float(self.down_multiplier),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -17,6 +17,7 @@ from liger_kernel.ops import LigerSiLUMulFunction
|
||||
from liger_kernel.transformers.functional import liger_swiglu
|
||||
from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP
|
||||
from liger_kernel.transformers.swiglu import LigerExperts
|
||||
from liger_kernel.transformers.swiglu import LigerFalconH1SwiGLUMLP
|
||||
from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP
|
||||
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
|
||||
from liger_kernel.utils import infer_comm_backend
|
||||
@@ -416,6 +417,180 @@ def test_correctness_functional(bsz, seq_len, size, dtype, atol, rtol):
|
||||
assert torch.allclose(b1.grad, b2.grad, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def _torch_silu_mul_ref(a, b, gate_multiplier, down_multiplier):
|
||||
"""Pure-PyTorch reference for silu(a * gate_mult) * b * down_mult."""
|
||||
scaled = a * gate_multiplier
|
||||
return torch.nn.functional.silu(scaled) * b * down_multiplier
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bsz, seq_len, size",
|
||||
[
|
||||
(2, 8, 8),
|
||||
(9, 7, 41),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"gate_multiplier, down_multiplier",
|
||||
[
|
||||
(0.7, 1.3),
|
||||
(1.5, 0.5),
|
||||
(1.0, 1.0), # degenerate case — must match the no-multiplier path
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype, atol, rtol",
|
||||
[
|
||||
(torch.float32, 1e-3, 1e-5),
|
||||
pytest.param(
|
||||
torch.bfloat16,
|
||||
1e-1,
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_correctness_silumul_with_multipliers(bsz, seq_len, size, gate_multiplier, down_multiplier, dtype, atol, rtol):
|
||||
_a = torch.randn(bsz, seq_len, size, device=device, dtype=dtype)
|
||||
_b = torch.randn(bsz, seq_len, size, device=device, dtype=dtype)
|
||||
|
||||
a1 = _a.clone().detach().requires_grad_(True)
|
||||
b1 = _b.clone().detach().requires_grad_(True)
|
||||
a2 = _a.clone().detach().requires_grad_(True)
|
||||
b2 = _b.clone().detach().requires_grad_(True)
|
||||
|
||||
y_ref = _torch_silu_mul_ref(a1, b1, gate_multiplier, down_multiplier)
|
||||
y_liger = LigerSiLUMulFunction.apply(a2, b2, gate_multiplier, down_multiplier)
|
||||
|
||||
torch.testing.assert_close(y_ref, y_liger, atol=atol, rtol=rtol)
|
||||
|
||||
grad = torch.randn_like(y_ref)
|
||||
y_ref.backward(grad.clone())
|
||||
y_liger.backward(grad.clone())
|
||||
|
||||
torch.testing.assert_close(a1.grad, a2.grad, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(b1.grad, b2.grad, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def test_silumul_default_multipliers_backward_compat():
|
||||
"""Calling LigerSiLUMulFunction.apply(a, b) without multipliers must behave exactly as before."""
|
||||
_a = torch.randn(4, 16, 32, device=device, dtype=torch.float32)
|
||||
_b = torch.randn(4, 16, 32, device=device, dtype=torch.float32)
|
||||
|
||||
a1 = _a.clone().detach().requires_grad_(True)
|
||||
b1 = _b.clone().detach().requires_grad_(True)
|
||||
a2 = _a.clone().detach().requires_grad_(True)
|
||||
b2 = _b.clone().detach().requires_grad_(True)
|
||||
|
||||
y_default = LigerSiLUMulFunction.apply(a1, b1)
|
||||
y_explicit = LigerSiLUMulFunction.apply(a2, b2, 1.0, 1.0)
|
||||
|
||||
torch.testing.assert_close(y_default, y_explicit)
|
||||
|
||||
grad = torch.randn_like(y_default)
|
||||
y_default.backward(grad.clone())
|
||||
y_explicit.backward(grad.clone())
|
||||
|
||||
torch.testing.assert_close(a1.grad, a2.grad)
|
||||
torch.testing.assert_close(b1.grad, b2.grad)
|
||||
|
||||
|
||||
class _FalconH1MLPRef(torch.nn.Module):
|
||||
"""Pure-PyTorch reference matching Falcon H1's MLP forward from issue #936."""
|
||||
|
||||
def __init__(self, hidden_size, intermediate_size, gate_multiplier, down_multiplier):
|
||||
super().__init__()
|
||||
self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False)
|
||||
self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False)
|
||||
self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False)
|
||||
self.gate_multiplier = gate_multiplier
|
||||
self.down_multiplier = down_multiplier
|
||||
|
||||
def forward(self, x):
|
||||
gate = self.gate_proj(x)
|
||||
up = self.up_proj(x)
|
||||
activated = torch.nn.functional.silu(gate * self.gate_multiplier) * up
|
||||
return self.down_proj(activated) * self.down_multiplier
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bsz, seq_len, hidden_size, intermediate_size",
|
||||
[
|
||||
(2, 256, 256, 512),
|
||||
(6, 42, 123, 431),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"gate_multiplier, down_multiplier",
|
||||
[
|
||||
(0.7, 1.3),
|
||||
(1.5, 0.5),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype, atol, rtol",
|
||||
[
|
||||
(torch.float32, 1e-0, 1e-5),
|
||||
pytest.param(
|
||||
torch.bfloat16,
|
||||
1e4,
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_correctness_falcon_h1_mlp(
|
||||
bsz, seq_len, hidden_size, intermediate_size, gate_multiplier, down_multiplier, dtype, atol, rtol
|
||||
):
|
||||
"""Parity test for LigerFalconH1SwiGLUMLP against a pure-PyTorch reference.
|
||||
|
||||
A pure-PyTorch reference is used rather than HF's FalconH1MLP so the test
|
||||
doesn't depend on transformers exposing FalconH1MLP at module scope.
|
||||
"""
|
||||
|
||||
class _FakeConfig:
|
||||
def __init__(self):
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.hidden_act = "silu"
|
||||
self.mlp_bias = False
|
||||
self.mlp_multipliers = (gate_multiplier, down_multiplier)
|
||||
|
||||
config = _FakeConfig()
|
||||
|
||||
_input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype)
|
||||
x1 = _input.clone().requires_grad_(True)
|
||||
x2 = _input.clone().requires_grad_(True)
|
||||
|
||||
G = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype)
|
||||
U = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype)
|
||||
D = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype)
|
||||
|
||||
ref_mlp = _FalconH1MLPRef(hidden_size, intermediate_size, gate_multiplier, down_multiplier).to(device).to(dtype)
|
||||
ref_mlp.gate_proj.weight.data = G.T.contiguous()
|
||||
ref_mlp.up_proj.weight.data = U.T.contiguous()
|
||||
ref_mlp.down_proj.weight.data = D.T.contiguous()
|
||||
|
||||
liger_mlp = LigerFalconH1SwiGLUMLP(config=config).to(device).to(dtype)
|
||||
liger_mlp.gate_proj.weight.data = G.T.contiguous()
|
||||
liger_mlp.up_proj.weight.data = U.T.contiguous()
|
||||
liger_mlp.down_proj.weight.data = D.T.contiguous()
|
||||
|
||||
y1 = ref_mlp(x1)
|
||||
y2 = liger_mlp(x2)
|
||||
|
||||
torch.testing.assert_close(y1, y2, atol=atol, rtol=rtol)
|
||||
|
||||
dy = torch.randn_like(y1)
|
||||
y1.backward(dy.clone(), retain_graph=True)
|
||||
y2.backward(dy.clone(), retain_graph=True)
|
||||
|
||||
torch.testing.assert_close(ref_mlp.gate_proj.weight.grad, liger_mlp.gate_proj.weight.grad, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(ref_mlp.up_proj.weight.grad, liger_mlp.up_proj.weight.grad, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(ref_mlp.down_proj.weight.grad, liger_mlp.down_proj.weight.grad, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(x1.grad, x2.grad, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def _test_dtensor_liger_silumul(rank, world_size, bsz, seq_len, hidden_size, dtype, atol, rtol, file_name):
|
||||
torch.distributed.init_process_group(
|
||||
backend=infer_comm_backend(),
|
||||
|
||||
Reference in New Issue
Block a user