wehub-resource-sync
197 lines
7.2 KiB
Python
197 lines
7.2 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 paddle
|
|
import paddle.distributed as dist
|
|
import paddle.nn.functional.flash_attention as fa_module
|
|
from paddle.nn.functional.flash_attention import (
|
|
_flashmask_unique_id_cache,
|
|
_get_or_create_unique_id,
|
|
flashmask_attention,
|
|
)
|
|
|
|
|
|
def _is_distributed_env():
|
|
return int(os.environ.get("PADDLE_TRAINERS_NUM", "1")) > 1
|
|
|
|
|
|
class _FakeGroup:
|
|
"""Minimal stand-in for dist.Group."""
|
|
|
|
def __init__(self, gid, rank=0, nranks=1):
|
|
self.id = gid
|
|
self.rank = rank
|
|
self.nranks = nranks
|
|
|
|
|
|
class TestGetOrCreateUniqueIdUnit(unittest.TestCase):
|
|
def setUp(self):
|
|
_flashmask_unique_id_cache.clear()
|
|
self._orig_get_uid = fa_module.flashmask_get_unique_id
|
|
self._orig_all_gather = dist.all_gather_object
|
|
# replace flashmask_get_unique_id with a dummy func
|
|
fa_module.flashmask_get_unique_id = lambda: paddle.ones(
|
|
[128], dtype='uint8'
|
|
)
|
|
# replace all_gather_object with a dummy func
|
|
dist.all_gather_object = lambda rl, obj, group=None: rl.append(obj)
|
|
|
|
def tearDown(self):
|
|
# recover the module functions
|
|
fa_module.flashmask_get_unique_id = self._orig_get_uid
|
|
dist.all_gather_object = self._orig_all_gather
|
|
_flashmask_unique_id_cache.clear()
|
|
|
|
def test_cache_hit(self):
|
|
"""cache hit → return (tensor, False)"""
|
|
_flashmask_unique_id_cache[1] = paddle.ones([128], dtype='uint8')
|
|
uid, is_new = _get_or_create_unique_id(_FakeGroup(1))
|
|
self.assertFalse(is_new)
|
|
self.assertEqual(uid.shape, [128])
|
|
|
|
def test_cache_miss_rank0(self):
|
|
"""rank=0 → flashmask_get_unique_id() → cache → return True"""
|
|
uid, is_new = _get_or_create_unique_id(_FakeGroup(2, rank=0))
|
|
self.assertTrue(is_new)
|
|
self.assertEqual(uid.shape, [128])
|
|
self.assertEqual(uid.dtype, paddle.uint8)
|
|
self.assertIn(2, _flashmask_unique_id_cache)
|
|
|
|
def test_cache_miss_non_rank0(self):
|
|
"""rank!=0 → paddle.zeros → cache → return True"""
|
|
uid, is_new = _get_or_create_unique_id(_FakeGroup(3, rank=1))
|
|
self.assertTrue(is_new)
|
|
self.assertTrue((uid.numpy() == 0).all())
|
|
self.assertIn(3, _flashmask_unique_id_cache)
|
|
|
|
|
|
class TestFlashMaskAttentionGroupUnit(unittest.TestCase):
|
|
"""Cover group-related branches inside flashmask_attention on single GPU."""
|
|
|
|
def setUp(self):
|
|
_flashmask_unique_id_cache.clear()
|
|
self._orig_get_uid = fa_module.flashmask_get_unique_id
|
|
self._orig_all_gather = dist.all_gather_object
|
|
fa_module.flashmask_get_unique_id = lambda: paddle.ones(
|
|
[128], dtype='uint8'
|
|
)
|
|
dist.all_gather_object = lambda rl, obj, group=None: rl.append(obj)
|
|
|
|
def tearDown(self):
|
|
fa_module.flashmask_get_unique_id = self._orig_get_uid
|
|
dist.all_gather_object = self._orig_all_gather
|
|
_flashmask_unique_id_cache.clear()
|
|
|
|
def test_group_none(self):
|
|
q = paddle.rand([1, 128, 2, 64], dtype='bfloat16').cuda()
|
|
k = paddle.rand([1, 128, 2, 64], dtype='bfloat16').cuda()
|
|
v = paddle.rand([1, 128, 2, 64], dtype='bfloat16').cuda()
|
|
out = flashmask_attention(q, k, v, causal=True)
|
|
self.assertEqual(out.shape, [1, 128, 2, 64])
|
|
|
|
def test_group_nranks1(self):
|
|
q = paddle.rand([1, 128, 2, 64], dtype='bfloat16').cuda()
|
|
k = paddle.rand([1, 128, 2, 64], dtype='bfloat16').cuda()
|
|
v = paddle.rand([1, 128, 2, 64], dtype='bfloat16').cuda()
|
|
out = flashmask_attention(
|
|
q, k, v, causal=True, group=_FakeGroup(100, rank=0, nranks=1)
|
|
)
|
|
self.assertEqual(out.shape, [1, 128, 2, 64])
|
|
|
|
|
|
# Multi-GPU integration tests — only under paddle.distributed.launch
|
|
@unittest.skipUnless(
|
|
_is_distributed_env(),
|
|
"Requires multi-GPU (run with: python -m paddle.distributed.launch --gpus=0,1)",
|
|
)
|
|
class TestGetOrCreateUniqueIdDistributed(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
dist.init_parallel_env()
|
|
|
|
def tearDown(self):
|
|
_flashmask_unique_id_cache.clear()
|
|
|
|
def test_first_call_returns_is_new_true(self):
|
|
group = dist.new_group(list(range(dist.get_world_size())))
|
|
uid, is_new = _get_or_create_unique_id(group)
|
|
self.assertTrue(is_new)
|
|
self.assertEqual(uid.shape, [128])
|
|
self.assertEqual(uid.dtype, paddle.uint8)
|
|
|
|
def test_second_call_returns_is_new_false(self):
|
|
group = dist.new_group(list(range(dist.get_world_size())))
|
|
uid1, is_new1 = _get_or_create_unique_id(group)
|
|
uid2, is_new2 = _get_or_create_unique_id(group)
|
|
self.assertTrue(is_new1)
|
|
self.assertFalse(is_new2)
|
|
self.assertTrue(paddle.equal_all(uid1, uid2).item())
|
|
|
|
def test_unique_id_consistent_across_ranks(self):
|
|
group = dist.new_group(list(range(dist.get_world_size())))
|
|
uid, _ = _get_or_create_unique_id(group)
|
|
uid_list = []
|
|
dist.all_gather_object(uid_list, uid.numpy().tolist(), group=group)
|
|
for i in range(1, len(uid_list)):
|
|
self.assertEqual(uid_list[0], uid_list[i])
|
|
|
|
|
|
@unittest.skipUnless(
|
|
_is_distributed_env(),
|
|
"Requires multi-GPU (run with: python -m paddle.distributed.launch --gpus=0,1)",
|
|
)
|
|
class TestFlashMaskAttentionGroupParam(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
dist.init_parallel_env()
|
|
|
|
def tearDown(self):
|
|
_flashmask_unique_id_cache.clear()
|
|
|
|
def test_group_none_no_distributed(self):
|
|
q = paddle.rand([1, 8, 2, 32], dtype='bfloat16')
|
|
k = paddle.rand([1, 8, 2, 32], dtype='bfloat16')
|
|
v = paddle.rand([1, 8, 2, 32], dtype='bfloat16')
|
|
out = flashmask_attention(q, k, v, causal=True)
|
|
self.assertEqual(out.shape, [1, 8, 2, 32])
|
|
|
|
def test_group_extracts_rank_nranks_for_mask_validation(self):
|
|
world_size = dist.get_world_size()
|
|
if world_size < 2:
|
|
self.skipTest("Need at least 2 GPUs")
|
|
group = dist.new_group(list(range(world_size)))
|
|
B, S, H, D = 1, 8, 2, 32
|
|
q = paddle.rand([B, S, H, D], dtype='bfloat16')
|
|
k = paddle.rand([B, S, H, D], dtype='bfloat16')
|
|
v = paddle.rand([B, S, H, D], dtype='bfloat16')
|
|
startend_row_indices = paddle.full(
|
|
[B, 1, S * world_size, 1], S * world_size, dtype='int32'
|
|
)
|
|
try:
|
|
flashmask_attention(
|
|
q, k, v, startend_row_indices, causal=True, group=group
|
|
)
|
|
except Exception as e:
|
|
self.assertNotIn(
|
|
"startend_row_indices.shape[2] must be equal to seqlen_k",
|
|
str(e),
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|