Add FP32 dtype support for RoPE - Part2 (#13328)
This commit is contained in:
@@ -146,6 +146,12 @@ class TestROPE(CustomTestCase):
|
||||
(128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 32, 8),
|
||||
(128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 16, 4),
|
||||
(512, 128, 311, 10000, False, torch.bfloat16, "cpu", 3, 39, 4, 2),
|
||||
(64, 64, 32, 8000, True, torch.float32, "cpu", 32, 32, 1, 1),
|
||||
(256, 128, 4096, 10000, True, torch.float32, "cpu", 2, 512, 32, 8),
|
||||
(512, 128, 311, 10000, True, torch.float32, "cpu", 3, 39, 4, 2),
|
||||
(128, 128, 2048, 10000, False, torch.float32, "cpu", 2, 512, 32, 8),
|
||||
(128, 128, 2048, 10000, False, torch.float32, "cpu", 2, 512, 16, 4),
|
||||
(512, 128, 311, 10000, False, torch.float32, "cpu", 3, 39, 4, 2),
|
||||
]
|
||||
|
||||
for (
|
||||
|
||||
@@ -76,7 +76,7 @@ num_tokens_list = [11, 8192]
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [1, 2])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("num_tokens", num_tokens_list)
|
||||
def test_mrope(
|
||||
model_name: str,
|
||||
|
||||
Reference in New Issue
Block a user