wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
1590 lines
65 KiB
Python
1590 lines
65 KiB
Python
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. 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 math
|
|
import random
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from nemo.collections.tts.modules.ffn_modules import ConvolutionLayer, PositionwiseConvFF
|
|
from nemo.collections.tts.modules.moe_modules import MoERouter, PositionwiseConvFFMoE
|
|
from nemo.collections.tts.modules.transformer_2501 import CrossAttention, SelfAttention, Transformer, TransformerLayer
|
|
from nemo.collections.tts.parts.utils.tts_dataset_utils import beta_binomial_prior_distribution
|
|
|
|
|
|
def set_seed(seed):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
|
|
|
|
from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestConvolutionLayer:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.in_channels = 3
|
|
cls.out_channels = 6
|
|
cls.kernel_size = 3
|
|
cls.stride = 1
|
|
cls.dilation = 1
|
|
cls.bias = True
|
|
# fmt:off
|
|
cls.input_tensor = torch.Tensor(
|
|
[[[-1.0542, 0.2675, 0.6963, 0.4738, 0.3910, -0.1505, 0.9171, -0.1528, 3.7269, 0.1779],
|
|
[-1.0317, 1.6818, 1.4257, -0.5003, -1.7254, 0.8830, -0.4541, -0.4631, -0.0986, 0.5083],
|
|
[-0.3231, -1.0899, 0.5774, 0.1661, 0.9620, -2.3307, -0.6158, -0.3663, 1.2469, -1.0208]]]
|
|
)
|
|
cls.input_mask = torch.ones(1, cls.input_tensor.shape[2])
|
|
# fmt:on
|
|
|
|
def test_non_causal_forward(self):
|
|
set_seed(0)
|
|
layer = ConvolutionLayer(
|
|
self.in_channels,
|
|
self.out_channels,
|
|
self.kernel_size,
|
|
self.stride,
|
|
dilation=self.dilation,
|
|
bias=self.bias,
|
|
is_causal=False,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output_tensor = layer(self.input_tensor, self.input_mask)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[ 0.1912, -0.0555, -0.2681, -0.2289, 1.0788, -0.3908, 0.0936, -0.7962, 1.3754, -0.0731],
|
|
[-0.3715, -0.6326, -0.9596, -0.0933, -0.1024, -0.2082, -0.5924, 0.1097, -0.5418, -0.0854],
|
|
[ 0.3974, 0.4537, 0.3299, 0.1471, -0.5983, -0.8645, 0.0975, 0.6063, -0.6619, -0.9711],
|
|
[-0.3048, 0.3862, -0.2462, -0.9903, -0.6189, 0.7389, 0.0785, -1.0870, -1.0018, -1.2426],
|
|
[-0.4357, -0.0446, 0.0879, 0.0930, -0.2242, 0.5285, 0.4006, -0.1846, 0.5668, -0.5242],
|
|
[-0.0625, 0.4123, -0.6289, -0.4317, 0.1595, 0.0386, -1.0774, 0.2218, 0.8483, -0.4886]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
|
|
def test_causal_forward(self):
|
|
set_seed(0)
|
|
layer = ConvolutionLayer(
|
|
self.in_channels,
|
|
self.out_channels,
|
|
self.kernel_size,
|
|
self.stride,
|
|
dilation=self.dilation,
|
|
bias=self.bias,
|
|
is_causal=True,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output_tensor = layer(self.input_tensor, self.input_mask)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[ 0.4301, 0.1912, -0.0555, -0.2681, -0.2289, 1.0788, -0.3908, 0.0936, -0.7962, 1.3754],
|
|
[-0.0501, -0.3715, -0.6326, -0.9596, -0.0933, -0.1024, -0.2082, -0.5924, 0.1097, -0.5418],
|
|
[-0.4204, 0.3974, 0.4537, 0.3299, 0.1471, -0.5983, -0.8645, 0.0975, 0.6063, -0.6619],
|
|
[ 0.1543, -0.3048, 0.3862, -0.2462, -0.9903, -0.6189, 0.7389, 0.0785, -1.0870, -1.0018],
|
|
[-0.1337, -0.4357, -0.0446, 0.0879, 0.0930, -0.2242, 0.5285, 0.4006, -0.1846, 0.5668],
|
|
[-0.6127, -0.0625, 0.4123, -0.6289, -0.4317, 0.1595, 0.0386, -1.0774, 0.2218, 0.8483]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
|
|
|
|
class TestPositionwiseConvFF:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.d_model = 3
|
|
cls.d_ffn = 12
|
|
cls.p_dropout = 0.0
|
|
cls.kernel_size = 3
|
|
cls.bias = True
|
|
# fmt:off
|
|
cls.input_tensor = torch.Tensor(
|
|
[[[-1.6682, -0.6069, 0.1321],
|
|
[-1.5489, 0.3279, -0.9159],
|
|
[-0.7490, 1.8984, 0.5030],
|
|
[-0.8130, 0.0058, -1.9979],
|
|
[-1.4994, -0.3270, 1.4961],
|
|
[-1.6613, -1.7827, 0.8932],
|
|
[-0.6276, -1.0770, -0.9971],
|
|
[ 1.5424, 1.3590, 1.2287],
|
|
[-0.1543, 0.3365, 1.7475],
|
|
[-0.1753, 0.4115, 0.0772]]]
|
|
)
|
|
cls.input_mask = torch.ones(1, cls.input_tensor.shape[1])
|
|
# fmt:on
|
|
|
|
def test_causal_forward(self):
|
|
set_seed(0)
|
|
layer = PositionwiseConvFF(
|
|
self.d_model, self.d_ffn, self.p_dropout, self.kernel_size, bias=self.bias, is_causal=True
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output_tensor = layer(self.input_tensor, self.input_mask)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[-0.1242, -0.0114, 0.0212],
|
|
[-0.0441, -0.0555, -0.0795],
|
|
[-0.0282, 0.0366, -0.2033],
|
|
[-0.0421, 0.0305, -0.2573],
|
|
[-0.1877, -0.2492, -0.1638],
|
|
[-0.4300, -0.1160, 0.2177],
|
|
[-0.1652, -0.3130, -0.3329],
|
|
[-0.1737, 0.1133, -0.1802],
|
|
[-0.2599, -0.0381, 0.1362],
|
|
[-0.0584, -0.2936, 0.2719]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
|
|
def test_non_causal_forward(self):
|
|
set_seed(0)
|
|
layer = PositionwiseConvFF(
|
|
self.d_model, self.d_ffn, self.p_dropout, self.kernel_size, bias=self.bias, is_causal=False
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output_tensor = layer(self.input_tensor, self.input_mask)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[-0.0617, -0.0321, -0.1646],
|
|
[-0.0421, 0.0305, -0.2573],
|
|
[-0.1877, -0.2492, -0.1638],
|
|
[-0.4300, -0.1160, 0.2177],
|
|
[-0.1652, -0.3130, -0.3329],
|
|
[-0.1737, 0.1133, -0.1802],
|
|
[-0.2599, -0.0381, 0.1362],
|
|
[-0.0584, -0.2936, 0.2719],
|
|
[ 0.0361, 0.1110, 0.0441],
|
|
[-0.0244, 0.0682, 0.0340]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
|
|
|
|
class TestSelfAttention:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.n_heads = 2
|
|
cls.d_model = 4
|
|
cls.p_dropout = 0.0
|
|
cls.max_length_causal_mask = 6
|
|
# fmt:off
|
|
cls.query_tensor = torch.Tensor(
|
|
[[[ 0.7239, -0.2362, -0.6610, -1.3759],
|
|
[ 1.7381, 0.0793, -1.1241, 0.9529],
|
|
[-1.9809, 0.2217, 0.0795, 0.0307],
|
|
[ 0.3208, 0.4485, 0.3046, -0.0704],
|
|
[-1.4412, 0.8981, 0.1219, 0.0481],
|
|
[ 1.7811, -0.1358, 0.6073, 0.8275]]]
|
|
)
|
|
# fmt:on
|
|
|
|
def test_causal_forward(self):
|
|
set_seed(0)
|
|
layer = SelfAttention(
|
|
self.n_heads,
|
|
self.d_model,
|
|
self.p_dropout,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
)
|
|
query_mask = torch.ones(1, self.max_length_causal_mask).bool()
|
|
with torch.no_grad():
|
|
output_tensor, attn_output = layer(self.query_tensor, query_mask)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[-0.2569, 0.2782, -0.0348, 0.2480],
|
|
[-0.3949, 0.4054, -0.0876, 0.2574],
|
|
[-0.1033, 0.0659, -0.0259, 0.1738],
|
|
[-0.1485, 0.0995, -0.0415, 0.0684],
|
|
[-0.0123, 0.0185, -0.0027, 0.0708],
|
|
[-0.0672, 0.0566, -0.0214, -0.0021]]]
|
|
)
|
|
expected_attn_prob = torch.Tensor(
|
|
[[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5828, 0.4172, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3216, 0.4260, 0.2525, 0.0000, 0.0000, 0.0000],
|
|
[0.2385, 0.2238, 0.2872, 0.2504, 0.0000, 0.0000],
|
|
[0.1807, 0.1973, 0.2057, 0.2045, 0.2118, 0.0000],
|
|
[0.1159, 0.1388, 0.2010, 0.1721, 0.2161, 0.1562]],
|
|
[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.2866, 0.7134, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.2799, 0.2472, 0.4729, 0.0000, 0.0000, 0.0000],
|
|
[0.2964, 0.2535, 0.2075, 0.2427, 0.0000, 0.0000],
|
|
[0.1864, 0.1616, 0.2394, 0.1974, 0.2152, 0.0000],
|
|
[0.1666, 0.2030, 0.1391, 0.1649, 0.1546, 0.1719]]]]
|
|
)
|
|
expected_attn_score = torch.Tensor(
|
|
[[[[ 0.5248, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.1948, -0.5291, float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.0279, 0.2533, -0.2698, float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.0508, -0.1145, 0.1350, -0.0020, float('-inf'), float('-inf')],
|
|
[-0.0985, -0.0105, 0.0315, 0.0257, 0.0604, float('-inf')],
|
|
[-0.3253, -0.1457, 0.2250, 0.0694, 0.2971, -0.0275]],
|
|
[[ 0.5075, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.5541, 0.3578, float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.2499, -0.3738, 0.2746, float('-inf'), float('-inf'), float('-inf')],
|
|
[ 0.2215, 0.0654, -0.1351, 0.0216, float('-inf'), float('-inf')],
|
|
[-0.1011, -0.2439, 0.1488, -0.0438, 0.0425, float('-inf')],
|
|
[ 0.0526, 0.2502, -0.1277, 0.0424, -0.0221, 0.0840]]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
assert torch.allclose(attn_output[0], expected_attn_prob, atol=1e-4)
|
|
assert torch.allclose(attn_output[1], expected_attn_score, atol=1e-4)
|
|
|
|
def test_non_causal_forward(self):
|
|
set_seed(0)
|
|
layer = SelfAttention(
|
|
self.n_heads,
|
|
self.d_model,
|
|
self.p_dropout,
|
|
is_causal=False,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
)
|
|
query_mask = torch.ones(1, self.max_length_causal_mask).bool()
|
|
with torch.no_grad():
|
|
output_tensor, attn_output = layer(self.query_tensor, query_mask)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[-0.0954, 0.1131, -0.0195, 0.0704],
|
|
[-0.0401, 0.0364, -0.0156, 0.0088],
|
|
[ 0.0324, 0.0368, 0.0174, 0.0501],
|
|
[-0.0633, 0.0610, -0.0176, 0.0180],
|
|
[-0.0017, 0.0361, 0.0030, 0.0319],
|
|
[-0.0672, 0.0566, -0.0214, -0.0021]]]
|
|
)
|
|
expected_attn_prob = torch.Tensor(
|
|
[[[[0.2835, 0.1482, 0.1723, 0.1426, 0.1422, 0.1111],
|
|
[0.1254, 0.0898, 0.2819, 0.1493, 0.2682, 0.0854],
|
|
[0.1549, 0.2051, 0.1216, 0.1663, 0.1294, 0.2228],
|
|
[0.1583, 0.1485, 0.1906, 0.1662, 0.1890, 0.1475],
|
|
[0.1498, 0.1635, 0.1705, 0.1696, 0.1755, 0.1711],
|
|
[0.1159, 0.1388, 0.2010, 0.1721, 0.2161, 0.1562]],
|
|
[[0.2536, 0.1945, 0.1079, 0.1628, 0.1233, 0.1579],
|
|
[0.0866, 0.2156, 0.1709, 0.1551, 0.1903, 0.1815],
|
|
[0.1361, 0.1202, 0.2300, 0.1626, 0.1941, 0.1569],
|
|
[0.2038, 0.1744, 0.1427, 0.1669, 0.1488, 0.1634],
|
|
[0.1565, 0.1357, 0.2010, 0.1657, 0.1807, 0.1604],
|
|
[0.1666, 0.2030, 0.1391, 0.1649, 0.1546, 0.1719]]]]
|
|
)
|
|
expected_attn_score = torch.Tensor(
|
|
[[[[ 5.2482e-01, -1.2346e-01, 2.7022e-02, -1.6210e-01, -1.6488e-01, -4.1190e-01],
|
|
[-1.9484e-01, -5.2910e-01, 6.1538e-01, -2.0263e-02, 5.6540e-01, -5.7875e-01],
|
|
[-2.7873e-02, 2.5326e-01, -2.6980e-01, 4.3127e-02, -2.0759e-01, 3.3584e-01],
|
|
[-5.0756e-02, -1.1455e-01, 1.3498e-01, -2.0208e-03, 1.2687e-01, -1.2113e-01],
|
|
[-9.8482e-02, -1.0463e-02, 3.1513e-02, 2.5712e-02, 6.0431e-02, 3.4494e-02],
|
|
[-3.2530e-01, -1.4574e-01, 2.2503e-01, 6.9375e-02, 2.9711e-01, -2.7542e-02]],
|
|
[[ 5.0748e-01, 2.4198e-01, -3.4704e-01, 6.4137e-02, -2.1347e-01, 3.3762e-02],
|
|
[-5.5410e-01, 3.5778e-01, 1.2559e-01, 2.8689e-02, 2.3316e-01, 1.8558e-01],
|
|
[-2.4989e-01, -3.7383e-01, 2.7462e-01, -7.2002e-02, 1.0508e-01, -1.0771e-01],
|
|
[ 2.2154e-01, 6.5375e-02, -1.3510e-01, 2.1609e-02, -9.3194e-02, 3.4042e-04],
|
|
[-1.0105e-01, -2.4395e-01, 1.4884e-01, -4.3842e-02, 4.2481e-02, -7.6735e-02],
|
|
[ 5.2595e-02, 2.5018e-01, -1.2765e-01, 4.2375e-02, -2.2093e-02, 8.4005e-02]]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
assert torch.allclose(attn_output[0], expected_attn_prob, atol=1e-4)
|
|
assert torch.allclose(attn_output[1], expected_attn_score, atol=1e-5)
|
|
|
|
|
|
class TestCrossAttention:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.n_heads = 2
|
|
cls.d_model = 4
|
|
cls.d_memory = 3
|
|
cls.p_dropout = 0.0
|
|
cls.max_length = 6
|
|
# fmt:off
|
|
# shape = (1, cls.max_length, cls.d_model)
|
|
cls.query_tensor = torch.Tensor(
|
|
[[[0.7352, -0.5871, -0.1204, -2.0200],
|
|
[-0.4618, -0.3604, 1.2287, -0.3434],
|
|
[0.7838, -0.7646, -1.3349, -0.1538],
|
|
[-0.9749, -1.0789, -0.0126, -0.7225],
|
|
[2.6929, -0.2091, 2.1242, -1.0123],
|
|
[0.5094, -2.0566, 1.3922, -0.2156]]]
|
|
)
|
|
# fmt:on
|
|
cls.query_mask = torch.ones(1, cls.query_tensor.shape[1]).bool()
|
|
# fmt:off
|
|
# shape = (1, 5, cls.d_memory)
|
|
cls.memory_tensor = torch.Tensor(
|
|
[[[ 2.0132e-01, -5.6582e-01, 1.1191e+00],
|
|
[-6.2371e-01, -9.3398e-02, -1.3744e+00],
|
|
[-9.8265e-01, -8.1742e-01, 4.5611e-01],
|
|
[-5.4802e-01, -1.1218e+00, 7.6138e-01],
|
|
[-1.9899e+00, -1.7910e-03, 9.0718e-01]]]
|
|
)
|
|
# fmt:on
|
|
cls.memory_mask = torch.ones(1, cls.memory_tensor.shape[1]).bool()
|
|
# shape = (1, cls.query_tensor.shape[1], cls.memory_tensor.shape[1])
|
|
cls.attn_prior = torch.from_numpy(
|
|
beta_binomial_prior_distribution(
|
|
phoneme_count=cls.memory_tensor.shape[1], mel_count=cls.query_tensor.shape[1]
|
|
)
|
|
).unsqueeze(0)
|
|
|
|
def test_forward(self):
|
|
set_seed(0)
|
|
layer = CrossAttention(self.n_heads, self.d_model, self.d_memory, self.p_dropout)
|
|
|
|
with torch.no_grad():
|
|
output_tensor, attn_output = layer(
|
|
self.query_tensor, self.query_mask, self.memory_tensor, self.memory_mask, self.attn_prior
|
|
)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = torch.Tensor(
|
|
[[[ 0.2267, -0.2271, 0.0573, -0.0681],
|
|
[ 0.2672, -0.1823, 0.0722, -0.0859],
|
|
[ 0.3212, -0.2218, 0.0835, -0.0715],
|
|
[ 0.3568, -0.2573, 0.0918, -0.0789],
|
|
[ 0.3962, -0.4112, 0.0816, -0.1972],
|
|
[ 0.3457, -0.4253, 0.0568, -0.2216]]]
|
|
)
|
|
expected_attn_prob = torch.Tensor(
|
|
[[[[0.4220, 0.4859, 0.0709, 0.0188, 0.0025],
|
|
[0.3944, 0.3475, 0.1642, 0.0784, 0.0155],
|
|
[0.1335, 0.3448, 0.2794, 0.1752, 0.0671],
|
|
[0.0914, 0.3300, 0.2343, 0.2437, 0.1006],
|
|
[0.0256, 0.1138, 0.2145, 0.3343, 0.3119],
|
|
[0.0117, 0.0617, 0.1112, 0.3354, 0.4800]],
|
|
[[0.8045, 0.1024, 0.0661, 0.0242, 0.0028],
|
|
[0.4020, 0.2953, 0.1914, 0.0907, 0.0207],
|
|
[0.1446, 0.2798, 0.3026, 0.1965, 0.0766],
|
|
[0.0718, 0.2151, 0.2778, 0.2719, 0.1634],
|
|
[0.0673, 0.0341, 0.1929, 0.4534, 0.2522],
|
|
[0.0064, 0.0264, 0.0999, 0.2872, 0.5802]]]]
|
|
)
|
|
expected_attn_score = torch.Tensor(
|
|
[[[[-0.5044, 0.4476, -0.4961, -0.5728, -0.8010],
|
|
[-0.0761, -0.2027, -0.5103, -0.4385, -0.6724],
|
|
[-0.1525, 0.2576, 0.0471, -0.0138, 0.0075],
|
|
[-0.3084, -0.0058, -0.7538, -0.7144, -1.0604],
|
|
[-0.0799, 0.0260, -0.1511, -0.1493, -0.2187],
|
|
[-0.1935, -0.3225, -0.9867, -0.8637, -1.3162]],
|
|
[[ 0.4704, -0.7801, -0.2374, 0.0126, -0.3430],
|
|
[ 0.0623, -0.2461, -0.2380, -0.1743, -0.2658],
|
|
[ 0.0104, 0.1313, 0.2096, 0.1834, 0.2217],
|
|
[-0.0859, 0.0300, -0.1194, -0.1410, -0.1111],
|
|
[ 0.7876, -1.2782, -0.3570, 0.0556, -0.5312],
|
|
[ 0.1046, -0.2652, -0.1856, -0.1104, -0.2180]]]]
|
|
)
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4)
|
|
assert torch.allclose(attn_output[0], expected_attn_prob, atol=1e-4)
|
|
assert torch.allclose(attn_output[1], expected_attn_score, atol=1e-4)
|
|
|
|
|
|
class TestTransformerLayer:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.d_model = 2
|
|
cls.d_ffn = 8
|
|
cls.sa_n_heads = 2
|
|
cls.kernel_size = 3
|
|
cls.p_dropout = 0.0
|
|
cls.max_length_causal_mask = 5
|
|
# fmt:off
|
|
# shape = (1, cls.max_length_causal_mask, cls.d_model)
|
|
cls.x = torch.Tensor(
|
|
[[[ 0.5115, 0.0889],
|
|
[-0.8568, -2.9632],
|
|
[-1.3728, 0.7325],
|
|
[-2.4593, -0.9018],
|
|
[ 0.9621, 0.4212]]]
|
|
)
|
|
# fmt:on
|
|
cls.x_mask = torch.ones(1, cls.max_length_causal_mask).bool()
|
|
# fmt:off
|
|
# shape = (1, 3, cls.d_model)
|
|
cls.cond = torch.Tensor(
|
|
[[[ 1.4441, 0.1393],
|
|
[ 0.2828, -0.2456],
|
|
[-0.3075, 0.6581]]]
|
|
)
|
|
# fmt:on
|
|
cls.cond_mask = torch.ones(1, cls.cond.shape[1]).bool()
|
|
# shape = (1, cls.x.shape[1], cls.cond.shape[1])
|
|
cls.attn_prior = torch.from_numpy(
|
|
beta_binomial_prior_distribution(phoneme_count=cls.cond.shape[1], mel_count=cls.x.shape[1])
|
|
).unsqueeze(0)
|
|
|
|
def test_forward_causal_self_attn_and_has_xattn(self):
|
|
set_seed(0)
|
|
layer = TransformerLayer(
|
|
self.d_model,
|
|
self.d_ffn,
|
|
self.sa_n_heads,
|
|
self.kernel_size,
|
|
self.p_dropout,
|
|
has_xattn=True,
|
|
xa_n_heads=2,
|
|
xa_d_memory=2,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output_dict = layer(self.x, self.x_mask, self.cond, self.cond_mask, self.attn_prior)
|
|
|
|
# fmt:off
|
|
expected_output = {
|
|
'output': torch.Tensor(
|
|
[[[ 0.1936, 0.5387],
|
|
[-1.0270, -2.5452],
|
|
[-1.6884, 0.8765],
|
|
[-2.7496, -0.7887],
|
|
[ 1.2837, 0.1172]]]
|
|
),
|
|
'attn_probabilities': {
|
|
'self_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5000, 0.5000, 0.0000, 0.0000, 0.0000],
|
|
[0.3068, 0.3068, 0.3864, 0.0000, 0.0000],
|
|
[0.2213, 0.2213, 0.2787, 0.2787, 0.0000],
|
|
[0.2180, 0.2180, 0.1730, 0.1730, 0.2180]],
|
|
[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5000, 0.5000, 0.0000, 0.0000, 0.0000],
|
|
[0.3237, 0.3237, 0.3527, 0.0000, 0.0000],
|
|
[0.2393, 0.2393, 0.2607, 0.2607, 0.0000],
|
|
[0.2068, 0.2068, 0.1898, 0.1898, 0.2068]]]]
|
|
),
|
|
torch.Tensor(
|
|
[[[[0.1154, float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[0.1154, 0.1154, float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.1154, -0.1154, 0.1154, float('-inf'), float('-inf')],
|
|
[-0.1154, -0.1154, 0.1154, 0.1154, float('-inf')],
|
|
[0.1154, 0.1154, -0.1154, -0.1154, 0.1154]],
|
|
[[0.0429, float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[0.0429, 0.0429, float('-inf'), float('-inf'), float('-inf')],
|
|
[-0.0429, -0.0429, 0.0429, float('-inf'), float('-inf')],
|
|
[-0.0429, -0.0429, 0.0429, 0.0429, float('-inf')],
|
|
[0.0429, 0.0429, -0.0429, -0.0429, 0.0429]]]]
|
|
)
|
|
],
|
|
'cross_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[0.7181, 0.2394, 0.0426],
|
|
[0.4843, 0.3874, 0.1283],
|
|
[0.2753, 0.4129, 0.3118],
|
|
[0.1344, 0.3583, 0.5074],
|
|
[0.0520, 0.2599, 0.6882]],
|
|
[[0.5959, 0.1987, 0.2054],
|
|
[0.2837, 0.2270, 0.4893],
|
|
[0.3740, 0.5610, 0.0651],
|
|
[0.2355, 0.6280, 0.1365],
|
|
[0.0108, 0.0542, 0.9349]]]]
|
|
),
|
|
torch.Tensor(
|
|
[[[[0.0586, 0.0586, -0.0586],
|
|
[0.0624, 0.0624, -0.0624],
|
|
[-0.0624, -0.0624, 0.0624],
|
|
[-0.0624, -0.0624, 0.0624],
|
|
[0.0624, 0.0624, -0.0624]],
|
|
[[-0.8214, -0.8214, 0.8214],
|
|
[-0.8745, -0.8744, 0.8745],
|
|
[0.8745, 0.8744, -0.8745],
|
|
[0.8745, 0.8744, -0.8745],
|
|
[-0.8744, -0.8744, 0.8744]]]]
|
|
)
|
|
]
|
|
}
|
|
}
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_dict["output"], expected_output["output"], atol=1e-4)
|
|
for i in range(2):
|
|
assert torch.allclose(
|
|
output_dict["attn_probabilities"]["self_attn_probabilities"][i],
|
|
expected_output["attn_probabilities"]["self_attn_probabilities"][i],
|
|
atol=1e-4,
|
|
)
|
|
assert torch.allclose(
|
|
output_dict["attn_probabilities"]["cross_attn_probabilities"][i],
|
|
expected_output["attn_probabilities"]["cross_attn_probabilities"][i],
|
|
atol=1e-4,
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTransformer:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.n_layers = 1
|
|
cls.d_model = 4
|
|
cls.d_ffn = 16
|
|
cls.sa_n_heads = 2
|
|
cls.kernel_size = 3
|
|
cls.p_dropout = 0.0
|
|
cls.p_dropout_out = 0.0
|
|
cls.is_causal = True
|
|
cls.max_length_causal_mask = 6
|
|
|
|
# fmt:off
|
|
cls.input_tensor = torch.Tensor(
|
|
[[[ 0.7049, 0.0305, -0.8542, 0.5388],
|
|
[-0.5265, -1.3320, 1.5451, 0.4086],
|
|
[-2.0546, 0.5259, 0.5995, -0.4078],
|
|
[ 0.4530, -0.3918, 2.1403, -0.2062],
|
|
[-0.0984, 0.4855, 0.7076, 0.0431],
|
|
[-0.4394, -0.6761, 1.7389, -0.9423]]]
|
|
)
|
|
# fmt:on
|
|
|
|
def test_forward_causal_self_attn_and_no_xattn(self):
|
|
set_seed(0)
|
|
model = Transformer(
|
|
n_layers=self.n_layers,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
p_dropout_out=self.p_dropout_out,
|
|
has_xattn=False,
|
|
is_causal=self.is_causal,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
)
|
|
|
|
# Check model init
|
|
assert torch.isclose(torch.mean(model.layers[0].pos_ff.proj.conv.weight), torch.tensor(0.0), atol=1e-2)
|
|
assert torch.isclose(torch.std(model.layers[0].pos_ff.proj.conv.weight), torch.tensor(0.02), atol=1e-2)
|
|
assert torch.isclose(torch.mean(model.layers[0].pos_ff.o_net.conv.weight), torch.tensor(0.0), atol=1e-2)
|
|
assert torch.isclose(
|
|
torch.std(model.layers[0].pos_ff.o_net.conv.weight), torch.tensor(0.02 / math.sqrt(2.0)), atol=1e-3
|
|
)
|
|
|
|
mask_tensor = torch.ones(1, self.max_length_causal_mask).bool()
|
|
with torch.no_grad():
|
|
output_dict = model(x=self.input_tensor, x_mask=mask_tensor)
|
|
|
|
# fmt:off
|
|
expected_output_tensor = {
|
|
'output': torch.Tensor(
|
|
[[[0.7047, 0.0305, -0.8555, 0.5402],
|
|
[-0.5192, -1.3324, 1.5455, 0.4148],
|
|
[-2.0593, 0.5290, 0.5969, -0.4101],
|
|
[0.4517, -0.3968, 2.1392, -0.2041],
|
|
[-0.1019, 0.4854, 0.7077, 0.0431],
|
|
[-0.4458, -0.6789, 1.7447, -0.9447]]]
|
|
),
|
|
'attn_probabilities': [
|
|
{
|
|
'self_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.4998, 0.5002, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3337, 0.3333, 0.3330, 0.0000, 0.0000, 0.0000],
|
|
[0.2498, 0.2500, 0.2501, 0.2501, 0.0000, 0.0000],
|
|
[0.2002, 0.2001, 0.2000, 0.1999, 0.1999, 0.0000],
|
|
[0.1666, 0.1666, 0.1667, 0.1667, 0.1667, 0.1667]],
|
|
[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5005, 0.4995, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3332, 0.3331, 0.3336, 0.0000, 0.0000, 0.0000],
|
|
[0.2507, 0.2494, 0.2510, 0.2489, 0.0000, 0.0000],
|
|
[0.2002, 0.1995, 0.2006, 0.1994, 0.2003, 0.0000],
|
|
[0.1671, 0.1663, 0.1674, 0.1660, 0.1670, 0.1662]]]]
|
|
),
|
|
torch.Tensor(
|
|
[[[[-3.4823e-04, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[-7.0210e-04, -5.6984e-05, float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[1.2551e-03, 1.1431e-04, -5.9334e-04, float('-inf'), float('-inf'), float('-inf')],
|
|
[-8.1514e-04, 2.6650e-05, 5.2952e-04, 5.6903e-04, float('-inf'), float('-inf')],
|
|
[8.0150e-04, 1.4366e-04, -2.7793e-04, -7.0636e-04, -8.2140e-04, float('-inf')],
|
|
[-4.6137e-04, 6.2648e-05, 3.6768e-04, 2.8095e-04, 4.9188e-04, 3.6888e-04]],
|
|
[[-7.5861e-04, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[8.0373e-04, -1.2745e-03, float('-inf'), float('-inf'), float('-inf'), float('-inf')],
|
|
[-4.5038e-04, -6.7648e-04, 7.5978e-04, float('-inf'), float('-inf'), float('-inf')],
|
|
[1.5376e-03, -3.6984e-03, 3.0423e-03, -5.4870e-03, float('-inf'), float('-inf')],
|
|
[3.7014e-04, -2.7010e-03, 2.4310e-03, -3.2604e-03, 9.3840e-04, float('-inf')],
|
|
[1.3868e-03, -3.8372e-03, 3.2144e-03, -5.4860e-03, 5.0273e-04, -4.4343e-03]]]]
|
|
),
|
|
],
|
|
'cross_attn_probabilities': None,
|
|
}
|
|
],
|
|
}
|
|
# fmt:on
|
|
|
|
assert output_dict["output"].shape == expected_output_tensor["output"].shape
|
|
assert torch.allclose(output_dict["output"], expected_output_tensor["output"], atol=1e-4)
|
|
for i in range(2):
|
|
assert torch.allclose(
|
|
output_dict["attn_probabilities"][0]["self_attn_probabilities"][i],
|
|
expected_output_tensor["attn_probabilities"][0]["self_attn_probabilities"][i],
|
|
atol=1e-4,
|
|
)
|
|
assert output_dict["attn_probabilities"][0]["cross_attn_probabilities"] is None
|
|
|
|
def test_forward_causal_self_attn_and_has_xattn(self):
|
|
set_seed(0)
|
|
model = Transformer(
|
|
n_layers=2,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
p_dropout_out=self.p_dropout_out,
|
|
has_xattn=True,
|
|
xa_d_memory=4,
|
|
xa_n_heads=2,
|
|
is_causal=self.is_causal,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
)
|
|
|
|
# fmt:off
|
|
cond = [
|
|
# shape (1, 3, 4)
|
|
torch.Tensor(
|
|
[[[-0.7475, 1.1461, 0.7300, 1.4471],
|
|
[ 1.8744, -0.1654, 1.2418, -1.6983],
|
|
[-0.3123, 0.2320, 0.7457, 1.9868]]]
|
|
),
|
|
# shape (1, 5, 4)
|
|
torch.Tensor(
|
|
[[[-0.6683, -1.2178, 1.3696, 0.9941],
|
|
[ 0.0297, -0.1616, 0.1891, 0.0580],
|
|
[-1.0771, 0.2547, -1.4023, 0.0971],
|
|
[ 1.1132, 0.6311, -0.1449, 0.2351],
|
|
[ 0.8920, 2.3663, 0.2248, -0.7298]]]
|
|
)
|
|
]
|
|
# fmt:on
|
|
|
|
cond_mask = [torch.ones(1, cond[0].shape[1]).bool(), torch.ones(1, cond[1].shape[1]).bool()]
|
|
mask_tensor = torch.ones(1, self.max_length_causal_mask).bool()
|
|
multi_encoder_mapping = [0, 1]
|
|
with torch.no_grad():
|
|
output_dict = model(
|
|
x=self.input_tensor,
|
|
x_mask=mask_tensor,
|
|
cond=cond,
|
|
cond_mask=cond_mask,
|
|
multi_encoder_mapping=multi_encoder_mapping,
|
|
)
|
|
|
|
# fmt:off
|
|
expected_output = {
|
|
'output': torch.Tensor(
|
|
[[[0.7043, 0.0288, -0.8547, 0.5384],
|
|
[-0.5283, -1.3311, 1.5429, 0.4083],
|
|
[-2.0560, 0.5259, 0.6020, -0.4099],
|
|
[0.4554, -0.3829, 2.1433, -0.2036],
|
|
[-0.0986, 0.4794, 0.7067, 0.0432],
|
|
[-0.4392, -0.6772, 1.7428, -0.9393]]]
|
|
),
|
|
'attn_probabilities': [
|
|
{
|
|
'self_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.4989, 0.5011, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3331, 0.3332, 0.3336, 0.0000, 0.0000, 0.0000],
|
|
[0.2495, 0.2496, 0.2504, 0.2505, 0.0000, 0.0000],
|
|
[0.1998, 0.1994, 0.2002, 0.2001, 0.2005, 0.0000],
|
|
[0.1662, 0.1662, 0.1668, 0.1668, 0.1671, 0.1669]],
|
|
[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5008, 0.4992, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3334, 0.3336, 0.3331, 0.0000, 0.0000, 0.0000],
|
|
[0.2504, 0.2500, 0.2496, 0.2500, 0.0000, 0.0000],
|
|
[0.2001, 0.2002, 0.2000, 0.1999, 0.1998, 0.0000],
|
|
[0.1670, 0.1667, 0.1665, 0.1667, 0.1665, 0.1666]]]]
|
|
),
|
|
],
|
|
'cross_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[0.3331, 0.3336, 0.3334],
|
|
[0.3335, 0.3331, 0.3334],
|
|
[0.3336, 0.3331, 0.3332],
|
|
[0.3335, 0.3332, 0.3334],
|
|
[0.3336, 0.3331, 0.3333],
|
|
[0.3335, 0.3332, 0.3333]],
|
|
[[0.3333, 0.3335, 0.3332],
|
|
[0.3334, 0.3335, 0.3331],
|
|
[0.3333, 0.3330, 0.3337],
|
|
[0.3334, 0.3335, 0.3332],
|
|
[0.3333, 0.3331, 0.3336],
|
|
[0.3334, 0.3334, 0.3333]]]]
|
|
)
|
|
]
|
|
},
|
|
{
|
|
'self_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5005, 0.4995, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3336, 0.3330, 0.3334, 0.0000, 0.0000, 0.0000],
|
|
[0.2503, 0.2499, 0.2498, 0.2500, 0.0000, 0.0000],
|
|
[0.2002, 0.1999, 0.2000, 0.2000, 0.2000, 0.0000],
|
|
[0.1669, 0.1666, 0.1666, 0.1667, 0.1666, 0.1666]],
|
|
[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.5001, 0.4999, 0.0000, 0.0000, 0.0000, 0.0000],
|
|
[0.3329, 0.3330, 0.3340, 0.0000, 0.0000, 0.0000],
|
|
[0.2499, 0.2498, 0.2505, 0.2499, 0.0000, 0.0000],
|
|
[0.1997, 0.1997, 0.2003, 0.2000, 0.2004, 0.0000],
|
|
[0.1665, 0.1664, 0.1669, 0.1666, 0.1669, 0.1667]]]]
|
|
),
|
|
],
|
|
'cross_attn_probabilities': [
|
|
torch.Tensor(
|
|
[[[[0.1999, 0.1998, 0.2002, 0.1999, 0.2001],
|
|
[0.2000, 0.1997, 0.2004, 0.1998, 0.2002],
|
|
[0.2001, 0.2000, 0.2001, 0.1999, 0.2000],
|
|
[0.2000, 0.2001, 0.1998, 0.2001, 0.1999],
|
|
[0.2001, 0.2002, 0.1998, 0.2001, 0.1998],
|
|
[0.2000, 0.2002, 0.1998, 0.2001, 0.1999]],
|
|
[[0.1998, 0.1998, 0.2001, 0.2004, 0.2000],
|
|
[0.2003, 0.2003, 0.1998, 0.1995, 0.2001],
|
|
[0.2003, 0.2003, 0.1998, 0.1995, 0.2001],
|
|
[0.2001, 0.2001, 0.2000, 0.1998, 0.2000],
|
|
[0.2002, 0.2001, 0.1999, 0.1997, 0.2000],
|
|
[0.2002, 0.2001, 0.2000, 0.1997, 0.2000]]]]
|
|
),
|
|
],
|
|
}
|
|
],
|
|
}
|
|
# fmt:on
|
|
|
|
assert torch.allclose(output_dict["output"], expected_output["output"], atol=1e-4)
|
|
for i in range(2):
|
|
assert torch.allclose(
|
|
output_dict["attn_probabilities"][i]["self_attn_probabilities"][0],
|
|
expected_output["attn_probabilities"][i]["self_attn_probabilities"][0],
|
|
atol=1e-4,
|
|
)
|
|
assert torch.allclose(
|
|
output_dict["attn_probabilities"][i]["cross_attn_probabilities"][0],
|
|
expected_output["attn_probabilities"][i]["cross_attn_probabilities"][0],
|
|
atol=1e-4,
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTransformerBatchedInference:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.n_layers = 3
|
|
cls.d_model = 4
|
|
cls.d_ffn = 16
|
|
cls.sa_n_heads = 2
|
|
cls.p_dropout = 0.0
|
|
cls.p_dropout_out = 0.0
|
|
cls.max_length_causal_mask = 10
|
|
cls.short_length = 4
|
|
cls.long_length = 10
|
|
|
|
def test_forward(self):
|
|
set_seed(0)
|
|
query_tensor1 = torch.randn(1, self.long_length, self.d_model)
|
|
query_tensor2 = torch.randn(1, self.short_length, self.d_model)
|
|
|
|
padding_tensor = torch.randn(1, self.long_length - self.short_length, self.d_model)
|
|
query_tensor2_padded = torch.cat([query_tensor2, padding_tensor], dim=1)
|
|
lengths = torch.tensor([self.long_length, self.short_length])
|
|
mask_batched = get_mask_from_lengths(lengths)
|
|
|
|
query_batched = torch.cat([query_tensor1, query_tensor2_padded], dim=0)
|
|
|
|
mask_bs1_1 = torch.ones(1, self.long_length)
|
|
mask_bs1_2 = torch.ones(1, self.short_length)
|
|
|
|
for is_causal in [True, False]:
|
|
for kernel_size in [1, 3]:
|
|
model = Transformer(
|
|
n_layers=self.n_layers,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
p_dropout_out=self.p_dropout_out,
|
|
is_causal=is_causal,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
)
|
|
|
|
output_batched = model(query_batched, mask_batched)
|
|
output_bs1_1 = model(query_tensor1, mask_bs1_1)
|
|
output_bs1_2 = model(query_tensor2, mask_bs1_2)
|
|
|
|
assert torch.allclose(
|
|
output_batched['output'][0][: self.long_length, :], output_bs1_1['output'], atol=1e-4
|
|
)
|
|
assert torch.allclose(
|
|
output_batched['output'][1][: self.short_length, :], output_bs1_2['output'], atol=1e-4
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMoERouter:
|
|
"""Test the MoERouter class for expert selection and auxiliary losses."""
|
|
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.d_model = 8
|
|
cls.num_experts = 4
|
|
cls.top_k = 2
|
|
cls.batch_size = 2
|
|
cls.seq_len = 5
|
|
|
|
def test_router_initialization(self):
|
|
"""Test that router initializes correctly with valid parameters."""
|
|
router = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
)
|
|
|
|
assert router.d_model == self.d_model
|
|
assert router.num_experts == self.num_experts
|
|
assert router.top_k == self.top_k
|
|
assert router.router.in_features == self.d_model
|
|
assert router.router.out_features == self.num_experts
|
|
|
|
def test_router_forward_shape(self):
|
|
"""Test that router output shapes are correct."""
|
|
set_seed(42)
|
|
router = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
with torch.no_grad():
|
|
expert_weights, expert_indices, router_logits, router_probs = router(x, x_mask)
|
|
|
|
# Check shapes
|
|
assert expert_weights.shape == (self.batch_size, self.seq_len, self.top_k)
|
|
assert expert_indices.shape == (self.batch_size, self.seq_len, self.top_k)
|
|
assert router_logits.shape == (self.batch_size, self.seq_len, self.num_experts)
|
|
assert router_probs.shape == (self.batch_size, self.seq_len, self.num_experts)
|
|
|
|
# Check that weights sum to 1
|
|
weight_sums = expert_weights.sum(dim=-1)
|
|
assert torch.allclose(weight_sums, torch.ones_like(weight_sums), atol=1e-5)
|
|
|
|
# Check that expert indices are valid (all valid tokens)
|
|
assert expert_indices.min() >= 0
|
|
assert expert_indices.max() < self.num_experts
|
|
|
|
def test_router_returns_logits_and_probs(self):
|
|
"""Test that router returns logits and probs for loss computation."""
|
|
set_seed(42)
|
|
router = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
)
|
|
router.train()
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
expert_weights, expert_indices, router_logits, router_probs = router(x, x_mask)
|
|
|
|
# Check that logits and probs are returned for loss computation
|
|
assert router_logits.shape == (self.batch_size, self.seq_len, self.num_experts)
|
|
assert router_probs.shape == (self.batch_size, self.seq_len, self.num_experts)
|
|
|
|
# Probs should sum to 1 for valid tokens
|
|
assert torch.allclose(router_probs.sum(dim=-1), x_mask, atol=1e-5)
|
|
|
|
def test_router_top_k_selection(self):
|
|
"""Test that exactly top_k experts are selected."""
|
|
set_seed(42)
|
|
for top_k in [1, 2, 3]:
|
|
router = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=top_k,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
with torch.no_grad():
|
|
expert_weights, expert_indices, _, _ = router(x, x_mask)
|
|
|
|
assert expert_weights.shape[-1] == top_k
|
|
assert expert_indices.shape[-1] == top_k
|
|
|
|
def test_router_sinkhorn_strategy(self):
|
|
"""Test that Sinkhorn routing works."""
|
|
set_seed(42)
|
|
router = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
routing_strategy="sinkhorn",
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
with torch.no_grad():
|
|
expert_weights, expert_indices, _, _ = router(x, x_mask)
|
|
|
|
# Just check it runs and produces valid outputs
|
|
assert expert_weights.shape == (self.batch_size, self.seq_len, self.top_k)
|
|
assert torch.allclose(expert_weights.sum(dim=-1), torch.ones(self.batch_size, self.seq_len), atol=1e-5)
|
|
|
|
def test_router_jitter_noise(self):
|
|
"""Test that jitter noise affects routing during training."""
|
|
set_seed(42)
|
|
router_with_noise = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
router_jitter_noise=0.1,
|
|
)
|
|
router_with_noise.train()
|
|
|
|
router_no_noise = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
router_jitter_noise=0.0,
|
|
)
|
|
router_no_noise.train()
|
|
|
|
# Copy weights to make them identical
|
|
router_no_noise.router.weight.data = router_with_noise.router.weight.data.clone()
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
# Get outputs with different random seeds to ensure noise causes differences
|
|
set_seed(100)
|
|
_, indices_with_noise, logits_with, probs_with = router_with_noise(x, x_mask)
|
|
set_seed(101) # Different seed
|
|
_, indices_no_noise, logits_no, probs_no = router_no_noise(x, x_mask)
|
|
|
|
# Shapes should match
|
|
assert indices_with_noise.shape == indices_no_noise.shape
|
|
|
|
# With noise, routing decisions or probabilities should differ
|
|
# (Noise is added to logits, which affects probs and potentially indices)
|
|
# At least one should be different due to noise
|
|
indices_differ = not torch.all(indices_with_noise == indices_no_noise)
|
|
probs_differ = not torch.allclose(probs_with, probs_no, atol=1e-5)
|
|
|
|
assert indices_differ or probs_differ, "Jitter noise should cause different routing decisions or probabilities"
|
|
|
|
def test_router_padding_masking(self):
|
|
"""Test that padded positions are properly masked with expert_indices=-1."""
|
|
set_seed(42)
|
|
router = MoERouter(
|
|
d_model=self.d_model,
|
|
num_experts=self.num_experts,
|
|
top_k=self.top_k,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
# Mask out last 2 positions of first batch
|
|
x_mask[0, 3:] = 0
|
|
|
|
with torch.no_grad():
|
|
expert_weights, expert_indices, router_logits, router_probs = router(x, x_mask)
|
|
|
|
# Check that padded positions have expert_indices = -1
|
|
assert torch.all(expert_indices[0, 3:, :] == -1), "Padded positions should have expert_indices=-1"
|
|
|
|
# Check that padded positions have zero weights
|
|
assert torch.allclose(expert_weights[0, 3:, :], torch.zeros_like(expert_weights[0, 3:, :]))
|
|
|
|
# Check that valid positions have valid expert indices (0 to num_experts-1)
|
|
assert torch.all(expert_indices[0, :3, :] >= 0), "Valid positions should have expert_indices >= 0"
|
|
assert torch.all(
|
|
expert_indices[0, :3, :] < self.num_experts
|
|
), "Valid positions should have valid expert indices"
|
|
|
|
# Check that router_logits are zero at padded positions
|
|
assert torch.allclose(router_logits[0, 3:, :], torch.zeros_like(router_logits[0, 3:, :]))
|
|
|
|
# Check that router_probs are zero at padded positions
|
|
assert torch.allclose(router_probs[0, 3:, :], torch.zeros_like(router_probs[0, 3:, :]))
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestPositionwiseConvFFMoE:
|
|
"""Test the PositionwiseConvFFMoE class."""
|
|
|
|
# Golden expected values captured from the original sequential implementation
|
|
# with set_seed(42), d_model=8, d_ffn=32, batch_size=2, seq_len=10.
|
|
# Each entry: (num_experts, top_k, bias, padding, expected_output_sum,
|
|
# expected_logits_sum, expected_first_expert_indices,
|
|
# expected_output_first_4_elements)
|
|
_GOLDEN_VALUES = {
|
|
"E4_top1_nobias_nopad": {
|
|
"output_sum": 1.4807705879211426,
|
|
"logits_sum": -1.4070085287094116,
|
|
"first_idx": [2],
|
|
"output_0_0": [-0.031408119946718216, 0.14679314196109772, -0.021915754303336143, 0.12932124733924866],
|
|
},
|
|
"E4_top2_nobias_nopad": {
|
|
"output_sum": 1.6757168769836426,
|
|
"logits_sum": -1.4070085287094116,
|
|
"first_idx": [2, 0],
|
|
"output_0_0": [-0.12169260531663895, -0.0002511143684387207, -0.09718462079763412, -0.017200887203216553],
|
|
},
|
|
"E8_top1_nobias_pad": {
|
|
"output_sum": 3.700800895690918,
|
|
"logits_sum": 8.973369598388672,
|
|
"first_idx": [4],
|
|
"output_0_0": [-0.28720206022262573, 0.20689748227596283, 0.40565282106399536, -0.021458642557263374],
|
|
},
|
|
"E8_top2_nobias_pad": {
|
|
"output_sum": 3.0652523040771484,
|
|
"logits_sum": 8.973369598388672,
|
|
"first_idx": [4, 1],
|
|
"output_0_0": [-0.2828606963157654, 0.1696583479642868, 0.2753525376319885, -0.041214004158973694],
|
|
},
|
|
"E4_top2_bias_pad": {
|
|
"output_sum": -1.1706292629241943,
|
|
"logits_sum": -1.2999919652938843,
|
|
"first_idx": [1, 3],
|
|
"output_0_0": [-0.10531097650527954, 0.14638465642929077, -0.1260562241077423, -0.11432743072509766],
|
|
},
|
|
"E16_top1_bias_pad": {
|
|
"output_sum": -0.7250787019729614,
|
|
"logits_sum": 6.78455924987793,
|
|
"first_idx": [8],
|
|
"output_0_0": [0.480951189994812, -0.3138628602027893, -0.010073505342006683, -0.05126545578241348],
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.d_model = 8
|
|
cls.d_ffn = 32
|
|
cls.p_dropout = 0.0
|
|
cls.num_experts = 4
|
|
cls.top_k_experts = 2
|
|
cls.kernel_size = 1 # MoE requires kernel_size=1
|
|
cls.batch_size = 2
|
|
cls.seq_len = 10
|
|
|
|
def test_moe_ffn_initialization(self):
|
|
"""Test that MoE FFN initializes correctly."""
|
|
set_seed(42)
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=self.p_dropout,
|
|
num_experts=self.num_experts,
|
|
top_k_experts=self.top_k_experts,
|
|
kernel_size=self.kernel_size,
|
|
)
|
|
|
|
assert moe_ffn.d_model == self.d_model
|
|
assert moe_ffn.d_ffn == self.d_ffn
|
|
assert moe_ffn.num_experts == self.num_experts
|
|
assert moe_ffn.top_k_experts == self.top_k_experts
|
|
assert len(moe_ffn.experts) == self.num_experts
|
|
|
|
def test_moe_ffn_forward_shape(self):
|
|
"""Test that MoE FFN produces correct output shape."""
|
|
set_seed(42)
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=self.p_dropout,
|
|
num_experts=self.num_experts,
|
|
top_k_experts=self.top_k_experts,
|
|
kernel_size=self.kernel_size,
|
|
is_causal=True,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
with torch.no_grad():
|
|
output, router_logits, router_probs, expert_indices = moe_ffn(x, x_mask)
|
|
|
|
# Check output shape
|
|
assert output.shape == x.shape
|
|
|
|
# Check routing info is returned
|
|
assert router_logits.shape == (self.batch_size, self.seq_len, self.num_experts)
|
|
assert router_probs.shape == (self.batch_size, self.seq_len, self.num_experts)
|
|
assert expert_indices.shape == (self.batch_size, self.seq_len, self.top_k_experts)
|
|
|
|
def test_moe_ffn_masking(self):
|
|
"""Test that masking works correctly and padded outputs are zero."""
|
|
set_seed(42)
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=self.p_dropout,
|
|
num_experts=self.num_experts,
|
|
top_k_experts=self.top_k_experts,
|
|
kernel_size=1,
|
|
is_causal=False,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
|
|
# Create a mask that zeros out some positions
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
x_mask[0, 5:] = 0 # Mask out last 5 positions of first sample
|
|
|
|
with torch.no_grad():
|
|
output, router_logits, router_probs, expert_indices = moe_ffn(x, x_mask)
|
|
|
|
# Check that output shape is correct
|
|
assert output.shape == x.shape
|
|
|
|
# Check that padded positions in output are zero (not processed by any expert)
|
|
assert torch.allclose(output[0, 5:, :], torch.zeros_like(output[0, 5:, :]))
|
|
|
|
# Check that router_logits and router_probs are zero at padded positions
|
|
assert torch.allclose(router_logits[0, 5:, :], torch.zeros_like(router_logits[0, 5:, :]))
|
|
assert torch.allclose(router_probs[0, 5:, :], torch.zeros_like(router_probs[0, 5:, :]))
|
|
|
|
# Check that expert_indices are -1 at padded positions
|
|
assert torch.all(expert_indices[0, 5:, :] == -1)
|
|
|
|
def test_moe_ffn_different_expert_counts(self):
|
|
"""Test that different numbers of experts work."""
|
|
set_seed(42)
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
|
|
for num_experts in [2, 4, 8]:
|
|
for top_k in [1, 2]:
|
|
if top_k <= num_experts:
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=self.p_dropout,
|
|
num_experts=num_experts,
|
|
top_k_experts=top_k,
|
|
kernel_size=1,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output, router_logits, router_probs, expert_indices = moe_ffn(x, x_mask)
|
|
|
|
assert output.shape == x.shape
|
|
|
|
def test_gradient_flow(self):
|
|
"""Backward pass must produce non-zero gradients for router and expert weights."""
|
|
set_seed(42)
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=0.0,
|
|
num_experts=self.num_experts,
|
|
top_k_experts=self.top_k_experts,
|
|
kernel_size=1,
|
|
)
|
|
moe_ffn.train()
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model, requires_grad=True)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
output, _, _, _ = moe_ffn(x, x_mask)
|
|
loss = output.sum()
|
|
loss.backward()
|
|
|
|
assert x.grad is not None and x.grad.abs().sum() > 0, "Input grad must be non-zero"
|
|
assert moe_ffn.router.router.weight.grad is not None, "Router weight grad must exist"
|
|
assert moe_ffn.router.router.weight.grad.abs().sum() > 0, "Router weight grad must be non-zero"
|
|
|
|
has_expert_grad = False
|
|
for expert in moe_ffn.experts:
|
|
for name in ('proj', 'o_net'):
|
|
g = expert[name].conv.weight.grad
|
|
if g is not None and g.abs().sum() > 0:
|
|
has_expert_grad = True
|
|
break
|
|
assert has_expert_grad, "At least one expert weight must receive a gradient"
|
|
|
|
def test_all_padding_produces_zeros(self):
|
|
"""When x_mask is all zeros, output and routing info must be all zeros."""
|
|
set_seed(42)
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=0.0,
|
|
num_experts=self.num_experts,
|
|
top_k_experts=self.top_k_experts,
|
|
kernel_size=1,
|
|
)
|
|
moe_ffn.eval()
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.zeros(self.batch_size, self.seq_len)
|
|
|
|
with torch.no_grad():
|
|
output, router_logits, router_probs, expert_indices = moe_ffn(x, x_mask)
|
|
|
|
assert torch.all(output == 0), "Output must be all zeros when fully padded"
|
|
assert torch.all(router_logits == 0), "Router logits must be all zeros when fully padded"
|
|
assert torch.all(expert_indices == -1), "Expert indices must be -1 when fully padded"
|
|
|
|
@pytest.mark.parametrize(
|
|
"num_experts,top_k,use_bias,use_padding",
|
|
[
|
|
(4, 1, False, False),
|
|
(4, 2, False, False),
|
|
(8, 1, False, True),
|
|
(8, 2, False, True),
|
|
(4, 2, True, True),
|
|
(16, 1, True, True),
|
|
],
|
|
ids=[
|
|
"E4_top1_nobias_nopad",
|
|
"E4_top2_nobias_nopad",
|
|
"E8_top1_nobias_pad",
|
|
"E8_top2_nobias_pad",
|
|
"E4_top2_bias_pad",
|
|
"E16_top1_bias_pad",
|
|
],
|
|
)
|
|
def test_forward_golden_values(self, num_experts, top_k, use_bias, use_padding, request):
|
|
"""Verify forward() output matches golden expected values.
|
|
|
|
Golden values were captured from the original sequential implementation.
|
|
Any refactoring of forward() must reproduce these exact values (within
|
|
floating-point tolerance) to guarantee numerical equivalence.
|
|
"""
|
|
golden = self._GOLDEN_VALUES[request.node.callspec.id]
|
|
|
|
set_seed(42)
|
|
moe_ffn = PositionwiseConvFFMoE(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
p_dropout=0.0,
|
|
num_experts=num_experts,
|
|
top_k_experts=top_k,
|
|
kernel_size=1,
|
|
bias=use_bias,
|
|
)
|
|
moe_ffn.eval()
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len)
|
|
if use_padding:
|
|
x_mask[0, 7:] = 0
|
|
x_mask[1, 5:] = 0
|
|
|
|
with torch.no_grad():
|
|
output, router_logits, router_probs, expert_indices = moe_ffn(x, x_mask)
|
|
|
|
assert (
|
|
expert_indices[0, 0].tolist() == golden["first_idx"]
|
|
), f"Expert indices mismatch: got {expert_indices[0, 0].tolist()}, expected {golden['first_idx']}"
|
|
assert torch.allclose(
|
|
router_logits.sum(), torch.tensor(golden["logits_sum"]), atol=1e-5
|
|
), f"Logits sum mismatch: got {router_logits.sum().item()}, expected {golden['logits_sum']}"
|
|
assert torch.allclose(
|
|
output.sum(), torch.tensor(golden["output_sum"]), atol=1e-5
|
|
), f"Output sum mismatch: got {output.sum().item()}, expected {golden['output_sum']}"
|
|
expected_elems = torch.tensor(golden["output_0_0"])
|
|
assert torch.allclose(
|
|
output[0, 0, :4], expected_elems, atol=1e-5
|
|
), f"Output[0,0,:4] mismatch: got {output[0, 0, :4].tolist()}, expected {golden['output_0_0']}"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTransformerLayerWithMoE:
|
|
"""Test TransformerLayer with MoE enabled."""
|
|
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.d_model = 8
|
|
cls.d_ffn = 32
|
|
cls.sa_n_heads = 2
|
|
cls.kernel_size = 1 # MoE requires kernel_size=1
|
|
cls.p_dropout = 0.0
|
|
cls.max_length_causal_mask = 20
|
|
cls.batch_size = 2
|
|
cls.seq_len = 10
|
|
|
|
def test_transformer_layer_moe_initialization(self):
|
|
"""Test that TransformerLayer initializes with MoE correctly."""
|
|
set_seed(42)
|
|
layer = TransformerLayer(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=False,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
num_experts=4,
|
|
top_k_experts=2,
|
|
)
|
|
|
|
assert layer.use_moe is True
|
|
assert isinstance(layer.pos_ff, PositionwiseConvFFMoE)
|
|
|
|
def test_transformer_layer_without_moe(self):
|
|
"""Test that TransformerLayer without MoE still works."""
|
|
set_seed(42)
|
|
layer = TransformerLayer(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=False,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=False,
|
|
)
|
|
|
|
assert layer.use_moe is False
|
|
assert isinstance(layer.pos_ff, PositionwiseConvFF)
|
|
|
|
def test_transformer_layer_moe_forward(self):
|
|
"""Test forward pass of TransformerLayer with MoE."""
|
|
set_seed(42)
|
|
layer = TransformerLayer(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=False,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
num_experts=4,
|
|
top_k_experts=2,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len).bool()
|
|
|
|
with torch.no_grad():
|
|
output_dict = layer(x, x_mask)
|
|
|
|
# Check output
|
|
assert 'output' in output_dict
|
|
assert output_dict['output'].shape == x.shape
|
|
|
|
# Check MoE routing info is present
|
|
assert 'moe_routing_info' in output_dict
|
|
if output_dict['moe_routing_info'] is not None:
|
|
assert 'router_logits' in output_dict['moe_routing_info']
|
|
assert 'router_probs' in output_dict['moe_routing_info']
|
|
assert 'expert_indices' in output_dict['moe_routing_info']
|
|
|
|
def test_transformer_layer_moe_with_xattn(self):
|
|
"""Test TransformerLayer with both MoE and cross-attention."""
|
|
set_seed(42)
|
|
layer = TransformerLayer(
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=True,
|
|
xa_d_memory=self.d_model,
|
|
xa_n_heads=2,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
num_experts=4,
|
|
top_k_experts=2,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len).bool()
|
|
cond = torch.randn(self.batch_size, 5, self.d_model)
|
|
cond_mask = torch.ones(self.batch_size, 5).bool()
|
|
|
|
with torch.no_grad():
|
|
output_dict = layer(x, x_mask, cond, cond_mask)
|
|
|
|
assert output_dict['output'].shape == x.shape
|
|
assert 'moe_routing_info' in output_dict
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTransformerWithMoE:
|
|
"""Test full Transformer with MoE."""
|
|
|
|
@classmethod
|
|
def setup_class(cls):
|
|
cls.n_layers = 2
|
|
cls.d_model = 8
|
|
cls.d_ffn = 32
|
|
cls.sa_n_heads = 2
|
|
cls.kernel_size = 1 # MoE requires kernel_size=1
|
|
cls.p_dropout = 0.0
|
|
cls.batch_size = 2
|
|
cls.seq_len = 10
|
|
cls.max_length_causal_mask = 20
|
|
|
|
def test_transformer_moe_initialization(self):
|
|
"""Test that Transformer initializes with MoE correctly."""
|
|
set_seed(42)
|
|
model = Transformer(
|
|
n_layers=self.n_layers,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=False,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
num_experts=4,
|
|
top_k_experts=2,
|
|
)
|
|
|
|
assert model.use_moe is True
|
|
assert len(model.layers) == self.n_layers
|
|
for layer in model.layers:
|
|
assert isinstance(layer.pos_ff, PositionwiseConvFFMoE)
|
|
|
|
def test_transformer_moe_forward(self):
|
|
"""Test forward pass of Transformer with MoE."""
|
|
set_seed(42)
|
|
model = Transformer(
|
|
n_layers=self.n_layers,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=False,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
num_experts=4,
|
|
top_k_experts=2,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len).bool()
|
|
|
|
with torch.no_grad():
|
|
output_dict = model(x, x_mask)
|
|
|
|
# Check output
|
|
assert 'output' in output_dict
|
|
assert output_dict['output'].shape == x.shape
|
|
|
|
# Check MoE routing info is collected from all layers
|
|
assert 'moe_routing_info' in output_dict
|
|
assert output_dict['moe_routing_info'] is not None
|
|
assert len(output_dict['moe_routing_info']) == self.n_layers
|
|
|
|
# Each layer should have logits, probs, and indices
|
|
for layer_routing in output_dict['moe_routing_info']:
|
|
assert 'router_logits' in layer_routing
|
|
assert 'router_probs' in layer_routing
|
|
assert 'expert_indices' in layer_routing
|
|
|
|
def test_transformer_moe_with_cross_attention(self):
|
|
"""Test Transformer with both MoE and cross-attention."""
|
|
set_seed(42)
|
|
model = Transformer(
|
|
n_layers=self.n_layers,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=True,
|
|
xa_d_memory=self.d_model,
|
|
xa_n_heads=2,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
num_experts=4,
|
|
top_k_experts=2,
|
|
)
|
|
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len).bool()
|
|
cond = torch.randn(self.batch_size, 5, self.d_model)
|
|
cond_mask = torch.ones(self.batch_size, 5).bool()
|
|
|
|
with torch.no_grad():
|
|
output_dict = model(x, x_mask, cond, cond_mask)
|
|
|
|
assert output_dict['output'].shape == x.shape
|
|
assert 'moe_routing_info' in output_dict
|
|
|
|
def test_transformer_moe_different_configurations(self):
|
|
"""Test various MoE configurations."""
|
|
set_seed(42)
|
|
x = torch.randn(self.batch_size, self.seq_len, self.d_model)
|
|
x_mask = torch.ones(self.batch_size, self.seq_len).bool()
|
|
|
|
configs = [
|
|
{'num_experts': 2, 'top_k_experts': 1},
|
|
{'num_experts': 4, 'top_k_experts': 2},
|
|
{'num_experts': 8, 'top_k_experts': 2},
|
|
{'num_experts': 4, 'top_k_experts': 1, 'routing_strategy': 'sinkhorn'},
|
|
]
|
|
|
|
for config in configs:
|
|
model = Transformer(
|
|
n_layers=self.n_layers,
|
|
d_model=self.d_model,
|
|
d_ffn=self.d_ffn,
|
|
sa_n_heads=self.sa_n_heads,
|
|
kernel_size=self.kernel_size,
|
|
p_dropout=self.p_dropout,
|
|
has_xattn=False,
|
|
is_causal=True,
|
|
max_length_causal_mask=self.max_length_causal_mask,
|
|
use_moe=True,
|
|
**config,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
output_dict = model(x, x_mask)
|
|
|
|
assert output_dict['output'].shape == x.shape
|