wehub-resource-sync
155 lines
5.1 KiB
Python
155 lines
5.1 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.incubate.nn.functional import batched_gemm as grouped_gemm
|
|
|
|
os.environ["FLAGS_flash_attn_version"] = "v1"
|
|
os.environ["FLAGS_cudnn_deterministic"] = "1"
|
|
os.environ["FLAGS_embedding_deterministic"] = "1"
|
|
|
|
|
|
def allclose(x, y, dtype):
|
|
if dtype == paddle.bfloat16:
|
|
rtol = 1e-5
|
|
else:
|
|
rtol = 1e-5
|
|
np.testing.assert_allclose(x.numpy(), y.numpy(), rtol=rtol)
|
|
|
|
|
|
_TEST_PROBLEMS = (
|
|
(1, 128, 128, 128),
|
|
(8, 128, 128, 128),
|
|
(16, 128, 128, 128),
|
|
(1, 128, 256, 512),
|
|
(8, 128, 256, 512),
|
|
(16, 128, 256, 512),
|
|
)
|
|
|
|
m_group_layout_cases = [(False, True), (False, False)]
|
|
|
|
|
|
def randn(bs, x, y, dtype=paddle.bfloat16):
|
|
out = (paddle.rand([bs, x, y]) - 0.5 * 2) / (y * x)
|
|
return out.astype(dtype)
|
|
|
|
|
|
def pyref_gmm(a, b, batch_sizes, trans_b=False):
|
|
out = []
|
|
start = 0
|
|
for i, size in enumerate(batch_sizes):
|
|
lhs = a[start : start + size, :]
|
|
rhs = b[i, :, :] if not trans_b else b[i, :, :].t()
|
|
out.append(lhs @ rhs)
|
|
start += size
|
|
return paddle.concat(out, axis=0)
|
|
|
|
|
|
def pyref_k_gmm(a, b, batch_sizes):
|
|
out = []
|
|
start = 0
|
|
for i, size in enumerate(batch_sizes):
|
|
lhs = a[start : start + size, :].t()
|
|
rhs = b[start : start + size, :]
|
|
out.append(lhs @ rhs)
|
|
start += size
|
|
return paddle.concat(out, axis=0)
|
|
|
|
|
|
class TestGroupedGemm(unittest.TestCase):
|
|
def setUp(self):
|
|
paddle.seed(42)
|
|
|
|
def test_m_grouped_gemm_fixed_sizes(self):
|
|
"""Test grouped GEMM with fixed sizes"""
|
|
# Test both bfloat16 and float32 dtypes
|
|
dtypes = [paddle.bfloat16, paddle.float32]
|
|
|
|
for dtype in dtypes:
|
|
for z, m, k, n in _TEST_PROBLEMS:
|
|
for trans_lhs, trans_rhs in m_group_layout_cases:
|
|
with self.subTest(
|
|
dtype=dtype,
|
|
z=z,
|
|
m=m,
|
|
k=k,
|
|
n=n,
|
|
trans_a=trans_lhs,
|
|
trans_b=trans_rhs,
|
|
) and paddle.amp.auto_cast(False):
|
|
a = randn(z, m, k, dtype).reshape([-1, k]).astype(dtype)
|
|
b = randn(z, k, n, dtype).astype(dtype)
|
|
if trans_rhs:
|
|
b = b.mT
|
|
batch_sizes = [m] * z
|
|
a.stop_gradient = False
|
|
b.stop_gradient = False
|
|
a_ref = a.clone().detach()
|
|
b_ref = b.clone().detach()
|
|
a_ref.stop_gradient = False
|
|
b_ref.stop_gradient = False
|
|
print(
|
|
f"Testing dtype={dtype}, shape={a.shape}, {b.shape}"
|
|
)
|
|
out = grouped_gemm(a, b, batch_sizes, False, trans_rhs)
|
|
expected_out = pyref_gmm(
|
|
a_ref, b_ref, batch_sizes, trans_rhs
|
|
)
|
|
allclose(out, expected_out.reshape(out.shape), dtype)
|
|
|
|
def test_k_grouped_gemm_variable_sizes(self):
|
|
"""Test grouped GEMM with variable sizes"""
|
|
# Test both bfloat16 and float32 dtypes
|
|
dtypes = [paddle.bfloat16, paddle.float32]
|
|
|
|
for dtype in dtypes:
|
|
for z, m, k, n in _TEST_PROBLEMS:
|
|
with self.subTest(
|
|
dtype=dtype, z=z, m=m, k=k, n=n, trans_a=True, trans_b=False
|
|
) and paddle.amp.auto_cast(False):
|
|
a = randn(z, m, k, dtype).astype(dtype)
|
|
b = randn(z, m, n, dtype).astype(dtype)
|
|
|
|
batch_sizes = [m] * z
|
|
|
|
a.stop_gradient = False
|
|
b.stop_gradient = False
|
|
a_ref = a.clone().detach()
|
|
b_ref = b.clone().detach()
|
|
a_ref.stop_gradient = False
|
|
b_ref.stop_gradient = False
|
|
|
|
out = grouped_gemm(
|
|
a.reshape([-1, k]),
|
|
b.reshape([-1, n]),
|
|
batch_sizes,
|
|
True,
|
|
False,
|
|
)
|
|
expected_out = pyref_k_gmm(
|
|
a_ref.reshape([-1, k]),
|
|
b_ref.reshape([-1, n]),
|
|
batch_sizes,
|
|
)
|
|
allclose(out, expected_out.reshape(out.shape), dtype)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|