TransformerEngine Integration

megatron/model/positional_embeddings.py

@@ -37,7 +37,7 @@ def forward(self, x, seq_dim=1):
 
 class RotaryEmbedding(torch.nn.Module):
     def __init__(
-        self, dim, max_seq_len, base=10000, precision=torch.half, save_inv_freqs=False
+        self, dim, max_seq_len, base=10000, precision=torch.half, save_inv_freqs=False, return_embeddings=False
     ):
         super().__init__()
         inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
@@ -49,6 +49,7 @@ def __init__(
         self.max_seq_len = max_seq_len
         self.base = base
         self.dim = dim
+        self.return_embeddings = return_embeddings
 
         # precompute cos_cached, sin_cached in fp32
         cos_cached, sin_cached, inv_freq = self._prepare_cache(
@@ -67,6 +68,8 @@ def _prepare_cache(self, seq_len, precision, base):
         freqs = torch.einsum("i,j->ij", t, inv_freq)
         emb = torch.cat((freqs, freqs), dim=-1)
 
+        self.emb = emb.reshape(emb.size(0), 1, 1, emb.size(1))
+
         cos_cached = emb.cos()[:, None, None, :]
         sin_cached = emb.sin()[:, None, None, :]
 
@@ -77,6 +80,8 @@ def _prepare_cache(self, seq_len, precision, base):
         )
 
     def forward(self, x, seq_dim=0, seq_len=None):
+        if self.return_embeddings:
+            return self.emb.to(self.precision).to(x.device)
         if seq_len is None:
             seq_len = x.shape[seq_dim]
 
@@ -249,4 +254,4 @@ def forward(self, x):
                 a.shape[0], 1, a.shape[2]
             )  # seq_len_k - 1 points to the last token index in the current inference batch.
 
-        return x + a
+        return x + a