| # Owner(s): ["module: nn"] |
| |
| import contextlib |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import unittest |
| |
| from torch.testing._internal.common_nn import NNTestCase |
| from torch.testing._internal.common_utils import TEST_FAIRSEQ, run_tests, parametrize, instantiate_parametrized_tests |
| from torch.testing._internal.common_cuda import TEST_CUDA |
| |
| if TEST_FAIRSEQ: |
| import fairseq.models.transformer as fairseq_transformer |
| |
| @contextlib.contextmanager |
| def set_default_dtype(dtype): |
| saved_dtype = torch.get_default_dtype() |
| torch.set_default_dtype(dtype) |
| try: |
| yield |
| finally: |
| torch.set_default_dtype(saved_dtype) |
| |
| class TestTransformers(NNTestCase): |
| _do_cuda_memory_leak_check = True |
| _do_cuda_non_default_stream = True |
| |
| device_list = ['cpu'] # TODO: is there a way to do parametrize for this? |
| if TEST_CUDA: |
| device_list.append('cuda') |
| |
| @unittest.skip("4D mask not supported yet - activate when 4D mask supported") |
| @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") # TODO: make this work for both cuda and cpu |
| def test_self_attn_TxT_attn_mask(self): |
| embed_dim = 16 |
| num_heads = 4 |
| batch_size = 10 |
| tgt_len = 16 |
| |
| query = torch.rand(batch_size, tgt_len, embed_dim, device="cuda") # [N, T, D] |
| attn_mask = torch.randint(0, 2, (tgt_len, tgt_len)).cuda().float() # [T, T] |
| attn_mask = attn_mask.masked_fill(attn_mask == 0, float('-inf')).masked_fill(attn_mask == 1, float(0.0)) |
| |
| attn_mask_4d = attn_mask.expand(batch_size, num_heads, tgt_len, tgt_len) |
| |
| mta_model = torch.nn.MultiheadAttention(embed_dim, num_heads, batch_first=True).cuda() |
| mta_model.eval() |
| |
| # Generate 3D results |
| with torch.inference_mode(): |
| output_mask_4d = mta_model(query, query, query, attn_mask=attn_mask_4d)[0] |
| output_mask_4d = output_mask_4d.transpose(0, 1) # [N, T, D] |
| |
| output_mask_TxT = mta_model(query, query, query, attn_mask=attn_mask)[0] |
| output_mask_TxT = output_mask_TxT.transpose(0, 1) # [N, T, D] |
| |
| self.assertEqual(output_mask_4d, output_mask_TxT) |
| |
| @parametrize("device", device_list) |
| def test_transformerencoderlayer_src_mask(self, device): |
| batch_size = 2 |
| seqlen = 4 |
| d_model = 8 |
| nhead = 8 |
| dim_feedforward = 32 |
| |
| model = torch.nn.TransformerEncoderLayer( |
| d_model=d_model, |
| nhead=nhead, |
| dim_feedforward=dim_feedforward, |
| batch_first=True).to(device) |
| src = torch.rand(batch_size, seqlen, d_model).to(device) # bs, seqlen, d_model |
| src_mask = torch.zeros(seqlen, seqlen).to(torch.bool).to(device) |
| |
| model(src, src_mask=src_mask) |
| model.eval() |
| with torch.no_grad(): |
| model(src, src_mask=src_mask) |
| |
| @parametrize("use_torchscript", [True, False]) |
| @parametrize("with_no_grad", [True, False]) |
| @parametrize("training", [True, False]) |
| def test_transformerencoder_fastpath_torchscript(self, use_torchscript, with_no_grad, training): |
| """ |
| Test TransformerEncoder does not crash |
| """ |
| model = torch.nn.TransformerEncoder( |
| torch.nn.TransformerEncoderLayer(d_model=2, nhead=2, dim_feedforward=8, batch_first=True), |
| num_layers=2, |
| enable_nested_tensor=True |
| ) |
| |
| if training: |
| model = model.train() |
| else: |
| model = model.eval() |
| |
| if use_torchscript: |
| model = torch.jit.script(model) |
| |
| x = torch.Tensor([[[1, 2], [3, 4]]]).to(torch.float) |
| mask = torch.Tensor([[0, 1]]).to(torch.bool) |
| |
| if with_no_grad: |
| cm = torch.no_grad() |
| else: |
| cm = contextlib.nullcontext() |
| with cm: |
| model(x, src_key_padding_mask=mask) |
| |
| @parametrize("with_no_grad", [True, False]) |
| @parametrize("training", [True, False]) |
| @parametrize("enable_nested_tensor", [False]) |
| @parametrize("device", device_list) |
| def test_transformerencoder_square_input(self, with_no_grad, training, enable_nested_tensor, device): |
| """ |
| Test for edge cases when input of shape (batch size, sequence length, embedding dimension) has |
| batch size == sequence length |
| """ |
| model = torch.nn.TransformerEncoder( |
| torch.nn.TransformerEncoderLayer(d_model=4, nhead=2, dim_feedforward=16, dropout=0.0, batch_first=True), |
| num_layers=2, |
| enable_nested_tensor=enable_nested_tensor |
| ).to(device) |
| |
| with torch.no_grad(): |
| # set constant weights of the model |
| for idx, p in enumerate(model.parameters()): |
| x = p.data |
| sz = x.view(-1).size(0) |
| shape = x.shape |
| x = torch.cos(torch.arange(0, sz).float().view(shape)) |
| p.data.copy_(x) |
| |
| if training: |
| model = model.train() |
| else: |
| model = model.eval() |
| x = torch.arange(0, 16).reshape(2, 2, 4).to(torch.float).to(device) |
| src_mask = torch.Tensor([[0, 1], [0, 0]]).to(torch.bool).to(device) |
| |
| if with_no_grad: |
| cm = torch.no_grad() |
| else: |
| cm = contextlib.nullcontext() |
| with cm: |
| result = model(x, mask=src_mask) |
| |
| ref_output = torch.Tensor([[[2.420306205749512, 0.017629241570830, -0.607857942581177, -0.085519507527351], |
| [2.420306205749512, 0.017629241570830, -0.607857942581177, -0.085519507527351]], |
| [[2.419836044311523, 0.017548924311996, -0.608187675476074, -0.085347734391689], |
| [2.419836044311523, 0.017548924311996, -0.608187675476074, -0.085347734391689]]] |
| ).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| @parametrize("batch_first", [True, False]) |
| @parametrize("training", [True, False]) |
| @parametrize("enable_nested_tensor", [True, False]) |
| @parametrize("device", device_list) |
| def test_transformerencoder(self, batch_first, training, enable_nested_tensor, device): |
| def get_a_test_layer(activation, batch_first=False): |
| d_model = 4 |
| nhead = 2 |
| dim_feedforward = 16 |
| dropout = 0.0 |
| |
| layer = nn.TransformerEncoderLayer( |
| d_model, |
| nhead, |
| dim_feedforward=dim_feedforward, |
| dropout=dropout, |
| activation=activation, |
| batch_first=batch_first, |
| ).to(device) |
| |
| with torch.no_grad(): |
| # set constant weights of the model |
| for idx, p in enumerate(layer.parameters()): |
| x = p.data |
| sz = x.view(-1).size(0) |
| shape = x.shape |
| x = torch.cos(torch.arange(0, sz).float().view(shape)) |
| p.data.copy_(x) |
| |
| return layer |
| |
| # this is a deterministic test for TransformerEncoder |
| activation = F.relu |
| |
| def _test(batch_first, training, enable_nested_tensor): |
| def perm_fn(x): |
| return x.transpose(1, 0) if batch_first else x |
| |
| encoder_layer = get_a_test_layer(activation=activation, |
| batch_first=batch_first) |
| |
| model = nn.TransformerEncoder(encoder_layer, 1).to(device) |
| if not training: |
| model = model.eval() |
| |
| # deterministic input |
| encoder_input = perm_fn(torch.tensor([[[0.7462, 0.6653, 0.5679, 0.4891], |
| [0.5387, 0.1655, 0.3565, 0.0471]], |
| [[0.8335, 0.2799, 0.5031, 0.2947], |
| [0.1402, 0.0318, 0.7636, 0.1346]], |
| [[0.6333, 0.9344, 0.1376, 0.9938], |
| [0.8924, 0.2872, 0.6692, 0.2944]], |
| [[0.9897, 0.6915, 0.3154, 0.1733], |
| [0.8645, 0.3513, 0.3064, 0.0767]], |
| [[0.8117, 0.2366, 0.4838, 0.7881], |
| [0.3718, 0.4945, 0.9511, 0.0864]]] |
| )).to(device) |
| result = model(encoder_input) |
| ref_output = perm_fn(torch.tensor([[[2.428589, 0.020835, -0.602055, -0.085249], |
| [2.427987, 0.021213, -0.602496, -0.084103]], |
| [[2.424689, 0.019155, -0.604793, -0.085672], |
| [2.413863, 0.022211, -0.612486, -0.072490]], |
| [[2.433774, 0.021598, -0.598343, -0.087548], |
| [2.425104, 0.019748, -0.604515, -0.084839]], |
| [[2.436185, 0.022682, -0.596625, -0.087261], |
| [2.433556, 0.021891, -0.598509, -0.086832]], |
| [[2.416246, 0.017512, -0.610712, -0.082961], |
| [2.422901, 0.024187, -0.606178, -0.074929]]] |
| )).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| # all 0 src_mask |
| src_mask = torch.zeros([5, 5]).to(device) == 1 |
| result = model(encoder_input, mask=src_mask) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| # all 0 |
| mask = torch.zeros([2, 5]).to(device) == 1 |
| result = model(encoder_input, src_key_padding_mask=mask) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| mask[0, 1] = 1 |
| mask[1, 3] = 1 |
| mask[1, 4] = 1 |
| # If mask is not left aligned |
| # We disable nested tensor |
| model.enable_nested_tensor = enable_nested_tensor |
| result = model(encoder_input, src_key_padding_mask=mask) |
| ref_output = perm_fn(torch.tensor([[[2.429026, 0.020793, -0.601741, -0.085642], |
| [2.428811, 0.021445, -0.601912, -0.084252]], |
| [[2.425009, 0.019155, -0.604566, -0.085899], |
| [2.415408, 0.02249, -0.611415, -0.073]], |
| [[2.434199, 0.021682, -0.598039, -0.087699], |
| [2.42598, 0.019941, -0.603896, -0.085091]], |
| [[2.436457, 0.022736, -0.59643, -0.08736], |
| [2.434021, 0.022093, -0.598179, -0.08679]], |
| [[2.416531, 0.017498, -0.610513, -0.083181], |
| [2.4242, 0.024653, -0.605266, -0.074959]]] |
| )).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| # test case 2, multiple layers no norm |
| model = nn.TransformerEncoder(encoder_layer, 2, enable_nested_tensor=enable_nested_tensor).to(device) |
| if not training: |
| model = model.eval() |
| result = model(encoder_input, src_key_padding_mask=mask) |
| ref_output = perm_fn(torch.tensor([[[2.419051, 0.017446, -0.608738, -0.085003], |
| [2.419102, 0.017452, -0.608703, -0.085026]], |
| [[2.419043, 0.017445, -0.608744, -0.084999], |
| [2.419052, 0.017446, -0.608738, -0.085004]], |
| [[2.419067, 0.017448, -0.608727, -0.085010], |
| [2.419098, 0.017452, -0.608706, -0.085024]], |
| [[2.419072, 0.017449, -0.608724, -0.085012], |
| [2.419119, 0.017455, -0.608691, -0.085034]], |
| [[2.419019, 0.017442, -0.608761, -0.084989], |
| [2.419075, 0.017449, -0.608722, -0.085014]]] |
| )).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| model = nn.TransformerEncoder(encoder_layer, 6, enable_nested_tensor=enable_nested_tensor).to(device) |
| if not training: |
| model = model.eval() |
| result = model(encoder_input, src_key_padding_mask=mask) |
| ref_output = perm_fn(torch.tensor([[[2.419101, 0.017453, -0.608703, -0.085025], |
| [2.419101, 0.017453, -0.608704, -0.085025]], |
| [[2.419101, 0.017453, -0.608703, -0.085025], |
| [2.419101, 0.017453, -0.608704, -0.085025]], |
| [[2.419101, 0.017453, -0.608703, -0.085025], |
| [2.419101, 0.017453, -0.608704, -0.085025]], |
| [[2.419101, 0.017453, -0.608703, -0.085025], |
| [2.419101, 0.017453, -0.608704, -0.085025]], |
| [[2.419101, 0.017453, -0.608703, -0.085025], |
| [2.419101, 0.017453, -0.608704, -0.085025]]] |
| )).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| # test case 3, multiple layers with norm |
| # d_model = 4 |
| norm = nn.LayerNorm(4) |
| model = nn.TransformerEncoder(encoder_layer, 2, norm=norm, enable_nested_tensor=enable_nested_tensor).to(device) |
| if not training: |
| model = model.eval() |
| result = model(encoder_input, src_key_padding_mask=mask) |
| ref_output = perm_fn(torch.tensor([[[1.695949, -0.357635, -0.893077, -0.445238], |
| [1.695955, -0.357639, -0.893050, -0.445266]], |
| [[1.695948, -0.357634, -0.893082, -0.445233], |
| [1.695950, -0.357635, -0.893077, -0.445238]], |
| [[1.695951, -0.357636, -0.893069, -0.445246], |
| [1.695955, -0.357639, -0.893052, -0.445264]], |
| [[1.695952, -0.357636, -0.893066, -0.445249], |
| [1.695957, -0.357641, -0.893041, -0.445276]], |
| [[1.695946, -0.357632, -0.893095, -0.445220], |
| [1.695952, -0.357637, -0.893065, -0.445251]]] |
| )).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| model = nn.TransformerEncoder(encoder_layer, 6, norm=norm, enable_nested_tensor=enable_nested_tensor).to(device) |
| if not training: |
| model = model.eval() |
| result = model(encoder_input, src_key_padding_mask=mask) |
| ref_output = perm_fn(torch.tensor([[[1.695955, -0.357639, -0.893051, -0.445265], |
| [1.695955, -0.357639, -0.893051, -0.445265]], |
| [[1.695955, -0.357639, -0.893051, -0.445265], |
| [1.695955, -0.357639, -0.893051, -0.445265]], |
| [[1.695955, -0.357639, -0.893051, -0.445265], |
| [1.695955, -0.357639, -0.893051, -0.445265]], |
| [[1.695955, -0.357639, -0.893051, -0.445265], |
| [1.695955, -0.357639, -0.893051, -0.445265]], |
| [[1.695955, -0.357639, -0.893051, -0.445265], |
| [1.695955, -0.357639, -0.893051, -0.445265]]] |
| )).to(device) |
| self.assertEqual(tuple(result.shape), tuple(ref_output.shape)) |
| torch.testing.assert_close(result, ref_output, rtol=1e-7, atol=1e-5) |
| |
| # TODO: remove set default dtype to double by making ref_output more precise. |
| # Added because this test was copied from test_nn.py, which has default |
| # dtype double. If default dtype is float, tests will say tensors not close because |
| # ref output precision too low |
| with set_default_dtype(torch.double): |
| if training: |
| cm = contextlib.nullcontext() |
| else: |
| cm = torch.no_grad() # transformer fast path requires no grad |
| with cm: |
| _test(batch_first, training, enable_nested_tensor) |
| |
| @unittest.skipIf(not TEST_FAIRSEQ, "Fairseq not found") |
| @unittest.skipIf(not TEST_CUDA, 'CUDA not available') |
| def test_decoder_only_layer(self): |
| DEFAULT_PADDING_IDX = 0 |
| |
| class FairseqDecoder(torch.nn.Module): |
| def __init__( |
| self, |
| embed_dim, |
| attention_heads, |
| ffn_embed_dim, |
| num_layers, |
| embedding_layer, # torch.nn.Embedding. Must have a padding_idx field |
| dropout=0, |
| normalize_before=False, |
| torch_encoder=None, # torch encoder that you can map weights from |
| activation="relu", |
| ): |
| super().__init__() |
| |
| cfg = fairseq_transformer.TransformerConfig() |
| cfg.decoder.embed_dim = embed_dim |
| cfg.decoder.output_dim = embed_dim |
| cfg.decoder.attention_heads = attention_heads |
| cfg.decoder.ffn_embed_dim = ffn_embed_dim |
| cfg.dropout = dropout |
| cfg.decoder.normalize_before = normalize_before |
| cfg.decoder.layers = num_layers |
| # make embedding behavior same as other encoders |
| cfg.no_token_positional_embeddings = True |
| cfg.no_scale_embedding = True |
| cfg.activation_fn = activation |
| |
| dictionary = {} # TODO: verify what this is |
| |
| self.decoder = fairseq_transformer.TransformerDecoder( |
| cfg, |
| dictionary, |
| embedding_layer, |
| no_encoder_attn=True, |
| output_projection=None, |
| ) |
| |
| if torch_encoder is not None: |
| self.decoder = torch_to_fairseq(torch_encoder, self.decoder) |
| self.decoder = self.decoder.eval().cuda().half() |
| |
| def forward( |
| self, |
| tokens, |
| src_lengths=None, |
| with_triangle_mask=False, |
| incremental_state=None, |
| ): |
| return self.decoder( |
| prev_output_tokens=tokens, |
| encoder_out=None, |
| incremental_state=incremental_state, |
| features_only=True, |
| full_context_alignment=not with_triangle_mask, |
| alignment_layer=None, |
| alignment_heads=None, |
| src_lengths=src_lengths, |
| return_all_hiddens=False, |
| )[0] |
| |
| class BetterDecoder(torch.nn.Module): |
| """ |
| Only incremental decoder for now |
| """ |
| |
| def __init__(self, transformer, embedding, pad_idx): |
| super().__init__() |
| self.transformer = transformer |
| self.embedding = embedding |
| self.padding_idx = pad_idx |
| |
| def forward( |
| self, |
| x, |
| src_mask=None, |
| include_padding_mask=True, |
| incr_key_lst=None, |
| incr_value_lst=None, |
| is_incremental_decoding=False, |
| ): |
| padding_mask = None |
| if not x.is_nested and include_padding_mask: |
| padding_mask = x.eq(self.padding_idx) |
| if(is_incremental_decoding): |
| x = x[:, -1:] # only take the last token |
| x = self.embedding(x) |
| |
| one_encoder_layer = self.transformer.layers[0] |
| self_attn = one_encoder_layer.self_attn |
| embed_dim = self_attn.embed_dim |
| num_heads = self_attn.num_heads |
| |
| use_gelu = ( |
| one_encoder_layer.activation_relu_or_gelu == 2 |
| ) # see torch/nn/modules/activation attention impl. 1 == relu, 2 == gelu |
| assert ( |
| one_encoder_layer.activation_relu_or_gelu != 0 |
| ) # 0 == not relu or gelu |
| |
| norm_first = one_encoder_layer.norm_first |
| |
| |
| # TODO: make this a bit less janky. but for now we initialize with an empty tensor. |
| if(not is_incremental_decoding): |
| assert len(incr_key_lst) == 0 or incr_key_lst[0] is None |
| assert len(incr_value_lst) == 0 or incr_value_lst[0] is None |
| while len(incr_key_lst) <= len(self.transformer.layers): |
| if(is_incremental_decoding): |
| incr_key_lst.append(torch.Tensor([]).cuda().half()) |
| incr_value_lst.append(torch.Tensor([]).cuda().half()) |
| else: |
| incr_key_lst.append(None) |
| incr_value_lst.append(None) |
| |
| for i, layer in enumerate(self.transformer.layers): |
| incr_key = incr_key_lst[i] |
| incr_value = incr_value_lst[i] |
| |
| x, incr_key, incr_value = torch._transformer_decoder_only_layer_fwd( |
| src=x, |
| embed_dim=embed_dim, |
| num_heads=num_heads, |
| qkv_weight=layer.self_attn.in_proj_weight, |
| qkv_bias=layer.self_attn.in_proj_bias, |
| proj_weight=layer.self_attn.out_proj.weight, |
| proj_bias=layer.self_attn.out_proj.bias, |
| use_gelu=use_gelu, |
| norm_first=norm_first, |
| # TODO: layer_norm_eps hardcoded to be same as nn.TransformerEncoder default. |
| # fix by pulling from self_attn.norm1 |
| eps=1e-5, |
| norm_weight_1=layer.norm1.weight, |
| norm_bias_1=layer.norm1.bias, |
| norm_weight_2=layer.norm2.weight, |
| norm_bias_2=layer.norm2.bias, |
| ffn_weight_1=layer.linear1.weight, |
| ffn_bias_1=layer.linear1.bias, |
| ffn_weight_2=layer.linear2.weight, |
| ffn_bias_2=layer.linear2.bias, |
| mask=src_mask, |
| incr_key=incr_key, # altered in place |
| incr_value=incr_value, |
| ) |
| |
| # not in-place |
| if(not is_incremental_decoding): |
| incr_key = None |
| incr_value = None |
| incr_key_lst[i] = incr_key |
| incr_value_lst[i] = incr_value |
| |
| return x, incr_key_lst, incr_value_lst |
| |
| def torch_to_fairseq(torch_encoder, fairseq_encoder): |
| for src_layer, dst_layer in zip(torch_encoder.layers, fairseq_encoder.layers): |
| w_q, w_k, w_v = src_layer.self_attn.in_proj_weight.chunk(3, dim=0) |
| b_q, b_k, b_v = src_layer.self_attn.in_proj_bias.chunk(3, dim=0) |
| |
| dst_layer.self_attn.q_proj.weight = torch.nn.Parameter(w_q) |
| dst_layer.self_attn.q_proj.bias = torch.nn.Parameter(b_q) |
| dst_layer.self_attn.k_proj.weight = torch.nn.Parameter(w_k) |
| dst_layer.self_attn.k_proj.bias = torch.nn.Parameter(b_k) |
| dst_layer.self_attn.v_proj.weight = torch.nn.Parameter(w_v) |
| dst_layer.self_attn.v_proj.bias = torch.nn.Parameter(b_v) |
| |
| dst_layer.self_attn.out_proj.weight = src_layer.self_attn.out_proj.weight |
| dst_layer.self_attn.out_proj.bias = src_layer.self_attn.out_proj.bias |
| |
| dst_layer.fc1.weight = src_layer.linear1.weight |
| dst_layer.fc1.bias = src_layer.linear1.bias |
| |
| # fairseq may use fusedlayernorm from nvidia apex - diff properties |
| dst_layer.self_attn_layer_norm.load_state_dict(src_layer.norm1.state_dict()) |
| |
| dst_layer.fc2.weight = src_layer.linear2.weight |
| dst_layer.fc2.bias = src_layer.linear2.bias |
| |
| dst_layer.final_layer_norm.load_state_dict(src_layer.norm2.state_dict()) |
| |
| return fairseq_encoder |
| |
| def set_weights_deterministic(model): |
| for idx, p in enumerate(model.parameters()): |
| x = p.data |
| sz = x.view(-1).size(0) |
| shape = x.shape |
| x = torch.cos(torch.arange(0, sz).float().view(shape)) |
| p.data.copy_(x) |
| |
| D = 4 # d_model |
| H = 2 # nhead |
| FD = 16 # dim_feedforward |
| V = 100 # vocab size |
| L = 2 # num layers |
| |
| embedding_layer = torch.nn.Embedding(V, D, DEFAULT_PADDING_IDX) |
| layer = torch.nn.TransformerEncoderLayer( |
| d_model=D, |
| nhead=H, |
| dim_feedforward=FD, |
| batch_first=True, |
| activation="gelu", |
| ) |
| transformer = torch.nn.TransformerEncoder( |
| layer, |
| num_layers=L, |
| ).eval().cuda().half() |
| |
| set_weights_deterministic(embedding_layer) |
| set_weights_deterministic(transformer) |
| |
| better_decoder = ( |
| BetterDecoder(transformer, embedding_layer, DEFAULT_PADDING_IDX) |
| .eval() |
| .cuda() |
| .half() |
| ) |
| fairseq_decoder = ( |
| FairseqDecoder( |
| D, |
| H, |
| FD, |
| L, |
| embedding_layer, |
| dropout=0, |
| normalize_before=False, |
| torch_encoder=transformer, |
| activation="gelu", |
| ) |
| .eval() |
| .cuda() |
| .half() |
| ) |
| |
| tokens = torch.Tensor([ |
| [5, 6, 7, 8], |
| [9, 10, 11, 12] |
| ]).to(torch.int).cuda() |
| lengths_tensor = torch.Tensor([2, 2]).to(torch.int).cuda() |
| # bs = 2, seqlen = 4 |
| bs, seqlen = tokens.shape |
| |
| upper_triangle = torch.zeros(seqlen, seqlen) |
| upper_triangle.fill_(-100000000) |
| upper_triangle = torch.triu(upper_triangle, 1) |
| upper_triangle = upper_triangle.cuda().half() |
| upper_triangle_expanded = upper_triangle.unsqueeze(0).unsqueeze(0) |
| upper_triangle_expanded = upper_triangle_expanded.expand( |
| bs, H, -1, -1 |
| ) |
| |
| # test forced decoding |
| with torch.no_grad(): |
| result, _, _ = better_decoder( |
| tokens, |
| src_mask=upper_triangle_expanded, |
| include_padding_mask=False, |
| incr_key_lst=[], |
| incr_value_lst=[], |
| is_incremental_decoding=False, |
| ) |
| ref_output = fairseq_decoder(tokens, lengths_tensor, with_triangle_mask=True) |
| |
| self.assertEqual(result.shape, ref_output.shape) |
| torch.testing.assert_close(result, ref_output, atol=1e-3, rtol=1e-2) |
| |
| # test incremental decoding |
| bs, seqlen = tokens.shape |
| |
| incr_state = {} |
| ref_outputs = [fairseq_decoder( |
| tokens[:, :i], |
| src_lengths=None, |
| with_triangle_mask=False, |
| incremental_state=incr_state, |
| ) for i in range(1, seqlen + 1)] |
| ref_output = torch.stack(ref_outputs) |
| |
| incr_key_lst = [] |
| incr_value_lst = [] |
| results = [] |
| for i in range(1, seqlen + 1): |
| res, incr_key_lst, incr_value_lst = better_decoder( |
| tokens[:, :i], |
| src_mask=None, |
| include_padding_mask=False, |
| incr_key_lst=incr_key_lst, |
| incr_value_lst=incr_value_lst, |
| is_incremental_decoding=True, |
| ) |
| results.append(res) |
| result = torch.stack(results) |
| |
| self.assertEqual(result.shape, ref_output.shape) |
| torch.testing.assert_close(result, ref_output, atol=1e-3, rtol=1e-2) |
| |
| instantiate_parametrized_tests(TestTransformers) |
| |
| if __name__ == '__main__': |
| run_tests() |
