[Feature] Infini Attention

.gitignore

@@ -163,3 +163,15 @@ cython_debug/
 
 checkpoints/
 wandb/
+*.pkl
+examples/infinite-context-length/*.txt
+
+*.json
+*.arrow
+*.txt
+*.out
+debug/
+llama3_8b_generated_from_brrr/
+llama3_8b_generated_from_brrr_tp_8/
+llama3_8b_generated_from_brrr_with_pp_4/
+examples/infinite-context-length/configs/exp59/

debug_exp33.ipynb

@@ -0,0 +1,67 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_from_disk"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# DATA_PATH = \"/fsx/phuc/projects/nanotron/examples/infinite-context-length/exp33_eval_data/needle_eval_ctx_len_1024_and_depth_20_and_id_0\"\n",
+    "# DATA_PATH = \"/fsx/phuc/projects/nanotron/examples/infinite-context-length/exp33_eval_data/needle_eval_ctx_len_1024_and_depth_0_and_id_0\"\n",
+    "DATA_PATH = \"/fsx/phuc/projects/nanotron/examples/infinite-context-length/exp33_eval_data/needle_eval_ctx_len_1024_and_depth_90_and_id_0\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = load_from_disk(DATA_PATH)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"There is a pass key hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about what is the pass key later on............................................................... The leaves rustle in the wind. The birds chirp their melodies. The squirrels scamper up the trees. Nature's symphony plays on. Life thrives in the forest.\\nThe city bustles with energy. The skyscrapers touch the clouds. The streets are alive with people. Horns honk and engines roar. The urban jungle never sleeps.\\nThe mountains stand tall and proud. The snow caps their peaks. The valleys are lush and green. Hikers traverse the trails. Adventure awaits in the wilderness.\\nThe stars twinkle in the night sky. The moon casts its gentle glow. The crickets sing their lullabies. Fireflies dance in the darkness. The world is at peace under the celestial blanket.\\nThe waves crash against the shore. The seagulls soar overhead. The sand is warm beneath my feet. I breathe in the salty air. The ocean's embrace is calming.\\nThe leaves rustle in the wind. The birds chirp their melodies. The squirrels scamper up the trees. Nature's symphony plays on. Life thrives in the forest.\\nThe city bustles with energy. The skyscrapers touch the clouds. The streets are alive with people. Horns honk and engines roar. The urban jungle never sleeps.\\nThe mountains stand tall and proud. The snow caps their peaks. The valleys are lush and green. Hikers traverse the trails. Adventure awaits in the wilderness.\\nThe stars twinkle in the night sky. The moon casts its gentle glow. The crickets sing their lullabies. Fireflies dance in the darkness. The world is at peace under the celestial blanket.\\nThe waves crash against the shore. The seagulls soar overhead. The sand is warm beneath my feet. I breathe in the salty air. The ocean's embrace is calming.\\nThe leaves rustle in the wind. The birds chirp their melodies. The squirrels scamper up the trees. Nature's symphony plays on. Life thrives in the forest.\\nThe city bustles with energy. The skyscrapers touch the clouds. The streets are alive with people. Horns honk and engines roar. The urban jungle never sleeps.\\nThe mountains stand tall and proud. The snow caps their peaks. The valleys are lush and green. Hikers traverse the trails. Adventure awaits in the wilderness.\\nThe stars twinkle in the night sky. The moon casts its gentle glow. The crickets sing their lullabies. Fireflies dance in the darkness. The world is at peace under the celestial blanket.\\nThe waves crash against the shore. The seagulls soar overhead. The sand is warm beneath my feet. I breathe in the salty air. The ocean's embrace is calming.\\nThe leaves rustle in the wind. The birds chirp their melodies. The squirrels scamper up the trees. Nature's symphony plays on. Life thrives in the forest.\\nThe city bustles with energy. The skyscrapers touch the clouds. The streets are alive with people. Horns honk and engines roar. The urban jungle never sleeps.\\nThe mountains stand tall and proud. The snow caps their peaks. The valleys are lush and green. Hikers traverse the trails. Adventure awaits in the wilderness.\\nThe stars twinkle in the night sky. The moon casts its gentle glow. The crickets sing their lullabies. Fireflies dance in the darkness. The world is at peace under the celestial blanket.\\nThe waves crash against the shore. The seagulls soar overhead. The sand is warm beneath my feet. I breathe in the salty air. The ocean's embrace is calming.\\nThe leaves rustle in the wind. The birds chirp their melodies. The squirrels scamper up the trees. Nature's symphony plays on. Life thrives in the forest.\\nThe city bustles with energy. The skyscrapers touch the clouds. The streets are alive with people. Horns honk and engines roar.. The pass key is 3. Remember it. 3 is the pass key.  The urban jungle never sleeps.\\nThe mountains stand tall and proud. The snow caps their peaks. The valleys are lush and green. Hikers traverse the trails. Adventure awaits in the wilderness.\\nThe stars twinkle in the night sky. The moon casts its gentle glow. The crickets sing their lullabies. Fireflies dance in the darkness. The world iThe waves crash against the shore. The seagulls soar overhead. The sand is warm beneath my feet. I breathe in the salty air. The. \\n\\nWhat is the pass key? The pass key is \""
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data[0][\"prompt\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}