titans-pytorch 0.0.53__tar.gz → 0.0.55__tar.gz

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.53
+Version: 0.0.55
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.53"
+version = "0.0.55"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/tests/test_titans.py

@@ -6,17 +6,20 @@ from titans_pytorch import NeuralMemory
 
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('silu', (False, True))
+@pytest.mark.parametrize('learned_mem_model_weights', (False, True))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
 def test_titans(
     seq_len,
     silu,
+    learned_mem_model_weights,
     max_grad_norm,
 ):
     mem = NeuralMemory(
         dim = 384,
         chunk_size = 64,
         activation = nn.SiLU() if silu else None,
-        max_grad_norm = max_grad_norm
+        max_grad_norm = max_grad_norm,
+        learned_mem_model_weights = learned_mem_model_weights
     )
 
     seq = torch.randn(2, seq_len, 384)

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/titans_pytorch/__init__.py

@@ -1,6 +1,7 @@
 from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
+    MemoryAttention
 )
 
 from titans_pytorch.mac_transformer import (

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/titans_pytorch/mac_transformer.py

@@ -227,9 +227,10 @@ class SegmentedAttention(Module):
         self,
         seq,
         value_residual = None,
-        flex_attn_fn: Callable | None = None
+        flex_attn_fn: Callable | None = None,
+        disable_flex_attn = False
     ):
-        if seq.is_cuda and self.use_flex_attn:
+        if seq.is_cuda and self.use_flex_attn and not disable_flex_attn:
             return self.forward_flex(seq, value_residual, flex_attn_fn)
 
         assert not (exists(value_residual) ^ exists(self.to_learned_v_mix))
@@ -303,7 +304,8 @@ class MemoryAsContextTransformer(Module):
         num_residual_streams = 4,
         neural_memory_kwargs: dict = dict(),
         neural_memory_layers: tuple[int, ...] | None = None,
-        aux_kv_recon_loss_weight = 0.
+        aux_kv_recon_loss_weight = 0.,
+        use_flex_attn = False
     ):
         super().__init__()
 
@@ -336,6 +338,8 @@ class MemoryAsContextTransformer(Module):
 
         assert not (num_longterm_mem_tokens > 0 and len(neural_memory_layers) == 0), 'empty `neural_memory_layers` when longterm memory tokens are present'
 
+        # mem, attn, and feedforward layers
+
         for layer in layers:
             is_first = layer == 1
 
@@ -363,6 +367,7 @@ class MemoryAsContextTransformer(Module):
                 dim_head = dim_head,
                 heads = heads,
                 segment_len = segment_len,
+                use_flex_attn = use_flex_attn,
                 accept_value_residual = not is_first,
                 num_longterm_mem_tokens = num_longterm_mem_tokens,
                 num_persist_mem_tokens = num_persist_mem_tokens
@@ -386,11 +391,20 @@ class MemoryAsContextTransformer(Module):
 
         self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
+        # flex attn related
+
+        assert not (use_flex_attn and not exists(flex_attention)), 'you need to be on the latest pytorch with a cuda device available'
+        self.use_flex_attn = use_flex_attn
+
+        self.segment_len = segment_len
+        self.num_persist_mem_tokens = num_persist_mem_tokens
+
     def forward(
         self,
         x,
         return_loss = False,
-        return_loss_breakdown = False
+        return_loss_breakdown = False,
+        disable_flex_attn = False
     ):
 
         if return_loss:
@@ -424,6 +438,17 @@ class MemoryAsContextTransformer(Module):
 
         x = x + pos_emb[:seq_len_with_mem]
 
+        # prep flex attention
+
+        use_flex_attn = x.is_cuda and self.use_flex_attn and not disable_flex_attn
+
+        flex_attn_fn = None
+
+        if use_flex_attn:
+            block_mask = create_mac_block_mask(seq_len, self.segment_len, self.num_persist_mem_tokens)
+
+            flex_attn_fn = partial(flex_attention, block_mask = block_mask)
+
         # value residual
 
         value_residual = None
@@ -442,7 +467,7 @@ class MemoryAsContextTransformer(Module):
                 x, aux_kv_loss = maybe_neural_mem(x, return_aux_kv_loss = True)
                 kv_recon_losses = kv_recon_losses + aux_kv_loss
 
-            x, values = attn(x, value_residual = value_residual)
+            x, values = attn(x, value_residual = value_residual, disable_flex_attn = disable_flex_attn, flex_attn_fn = flex_attn_fn)
 
             value_residual = default(value_residual, values)
 

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/titans_pytorch/titans.py

@@ -123,9 +123,12 @@ class MemoryMLP(Module):
 class MemoryAttention(Module):
     def __init__(
         self,
-        dim
+        dim,
+        scale = 8.
     ):
         super().__init__()
+        self.scale = scale
+
         self.weights = nn.ParameterList([
             nn.Parameter(torch.randn(dim, dim)), # queries
             nn.Parameter(torch.randn(dim, dim)), # keys
@@ -143,6 +146,7 @@ class MemoryAttention(Module):
 
         attn_out = F.scaled_dot_product_attention(
             q, k, v,
+            scale = self.scale,
             is_causal = True
         )
 
@@ -174,6 +178,7 @@ class NeuralMemory(Module):
         default_step_transform_max_lr = 1e-2,
         pre_rmsnorm = True,
         post_rmsnorm = True,
+        learned_mem_model_weights = True,
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
         activation: Module | None = None,
@@ -212,6 +217,9 @@ class NeuralMemory(Module):
         if not exists(model):
             model = MemoryMLP(dim_head, **default_model_kwargs)
 
+        if not learned_mem_model_weights:
+            model.requires_grad_(False)
+
         assert not exists(next(model.buffers(), None)), 'model cannot have buffers for now'
 
         # the memory is the weights of the model

{titans_pytorch-0.0.53 → titans_pytorch-0.0.55}/train_mac.py

@@ -9,7 +9,7 @@ from torch.optim import Adam
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
 
-from titans_pytorch.mac_transformer import MemoryAsContextTransformer
+from titans_pytorch import MemoryAsContextTransformer
 
 # constants
 
@@ -32,6 +32,7 @@ NUM_LONGTERM_MEM = 4
 NEURAL_MEM_LAYERS = (2, 4)
 WINDOW_SIZE = 32
 KV_RECON_LOSS_WEIGHT = 0.
+LEARNED_MEM_MODEL_WEIGHTS = True
 RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS}'
 
 # wandb experiment tracker
@@ -90,7 +91,7 @@ def base_decoding(
     sample_num_times = max(0, seq_len - prompt_seq_len)
 
     for _ in tqdm.tqdm(range(sample_num_times)):
-        logits = net(out)
+        logits = net(out, disable_flex_attn = True)
         logits = logits[:, -1]
 
         logits = min_p_filter(logits, min_p = min_p)
@@ -112,9 +113,11 @@ model = MemoryAsContextTransformer(
     neural_memory_layers = NEURAL_MEM_LAYERS,
     neural_memory_segment_len = WINDOW_SIZE // 2,
     aux_kv_recon_loss_weight = KV_RECON_LOSS_WEIGHT,
+    use_flex_attn = True,
     neural_memory_kwargs = dict(
         dim_head = 64,
         heads = 4,
+        learned_mem_model_weights = LEARNED_MEM_MODEL_WEIGHTS,
         default_model_kwargs = dict(
             depth = NEURAL_MEMORY_DEPTH,
         )