mirror of
https://github.com/facebookresearch/fairseq.git
synced 2024-08-17 04:20:36 +03:00
Make RotaryPositionalEmbedding jit-compatible (#5237)
This commit is contained in:
parent
31fba013a0
commit
100cd91db1
@ -14,27 +14,26 @@ class RotaryPositionalEmbedding(torch.nn.Module):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||||
self.register_buffer("inv_freq", inv_freq)
|
self.register_buffer("inv_freq", inv_freq)
|
||||||
self.seq_len_cached = None
|
self.seq_len_cached = 0
|
||||||
self.cos_cached = None
|
self.cos_cached = torch.empty(self.seq_len_cached, 1, 1, dim)
|
||||||
self.sin_cached = None
|
self.sin_cached = torch.empty(self.seq_len_cached, 1, 1, dim)
|
||||||
self.precision = precision
|
self.precision = precision
|
||||||
|
|
||||||
def forward(self, x, seq_len=None):
|
def forward(self, x, seq_len: int = 0):
|
||||||
"""
|
"""
|
||||||
Args:
|
Args:
|
||||||
x: Input x with T X B X C
|
x: Input x with T X B X C
|
||||||
seq_len: Sequence length of input x
|
seq_len: Sequence length of input x
|
||||||
"""
|
"""
|
||||||
if seq_len != self.seq_len_cached:
|
if seq_len > self.seq_len_cached:
|
||||||
self.seq_len_cached = seq_len
|
self.seq_len_cached = seq_len
|
||||||
t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
|
t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
|
||||||
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
|
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
|
||||||
emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
|
emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
|
||||||
self.cos_cached = emb.cos()[:, None, None, :]
|
self.cos_cached = emb.cos().view(emb.size(0), 1, 1, emb.size(1))
|
||||||
self.sin_cached = emb.sin()[:, None, None, :]
|
self.sin_cached = emb.sin().view(emb.size(0), 1, 1, emb.size(1))
|
||||||
return self.cos_cached, self.sin_cached
|
return self.cos_cached, self.sin_cached
|
||||||
|
|
||||||
|
|
||||||
# rotary pos emb helpers:
|
# rotary pos emb helpers:
|
||||||
def rotate_half(x):
|
def rotate_half(x):
|
||||||
x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
|
x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
|
||||||
|
@ -59,6 +59,27 @@ class TestRotaryPositionalEmbedding(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_jit_compile_rope_module(self):
|
||||||
|
module_scripted = torch.jit.script(self.rope_pos_emd)
|
||||||
|
apply_rotary_scripted = torch.jit.script(apply_rotary_pos_emb)
|
||||||
|
# Test several different lengths
|
||||||
|
for T in [3, 5, 10]:
|
||||||
|
sample = torch.randn(T, self.B, self.C)
|
||||||
|
# Run forward pass with the original module
|
||||||
|
cos_original, sin_original = self.rope_pos_emd(sample, T)
|
||||||
|
query = sample.view(T, self.B, 1, self.C)
|
||||||
|
new_query, new_key = apply_rotary_pos_emb(query, query, cos_original, sin_original)
|
||||||
|
|
||||||
|
# Run forward pass with the scripted module
|
||||||
|
cos_scripted, sin_scripted = module_scripted(sample, T)
|
||||||
|
new_query_scripted, new_key_scripted = apply_rotary_scripted(query, query, cos_scripted, sin_scripted)
|
||||||
|
|
||||||
|
# Ensure the outputs are the same
|
||||||
|
self.assertTrue(torch.allclose(cos_original, cos_scripted))
|
||||||
|
self.assertTrue(torch.allclose(sin_original, sin_scripted))
|
||||||
|
self.assertTrue(torch.allclose(new_query, new_query_scripted))
|
||||||
|
self.assertTrue(torch.allclose(new_key, new_key_scripted))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
Loading…
Reference in New Issue
Block a user