Better use of memory; limit window size and number of tokens

app.py

@@ -54,11 +54,25 @@ if not compact_layout:
 
 model_name = st.selectbox("Model", ["distilgpt2", "gpt2", "EleutherAI/gpt-neo-125m"])
 metric_name = st.selectbox("Metric", ["KL divergence", "Cross entropy"], index=1)
+
+tokenizer = st.cache_resource(AutoTokenizer.from_pretrained, show_spinner=False)(model_name, use_fast=False)
+
+# Make sure the logprobs do not use up more than ~6 GB of memory
+MAX_MEM = 6e9 / (torch.finfo(torch.float16).bits / 8)
+# Select window lengths such that we are allowed to fill the whole window without running out of memory
+# (otherwise the window length is irrelevant)
+window_len_options = [
+    w for w in [8, 16, 32, 64, 128, 256, 512, 1024]
+    if w == 8 or w * (2 * w) * tokenizer.vocab_size <= MAX_MEM
+]
 window_len = st.select_slider(
     r"Window size ($c_\text{max}$)",
-    options=[8, 16, 32, 64, 128, 256, 512, 1024],
-    value=512
+    options=window_len_options,
+    value=min(128, window_len_options[-1])
 )
+# Now figure out how many tokens we are allowed to use:
+# window_len * (num_tokens + window_len) * vocab_size <= MAX_MEM
+max_tokens = int(MAX_MEM / (tokenizer.vocab_size * window_len) - window_len)
 
 DEFAULT_TEXT = """
 We present context length probing, a novel explanation technique for causal
@@ -71,31 +85,38 @@ dependencies.
 """.replace("\n", " ").strip()
 
 text = st.text_area(
-    "Input text",
+    f"Input text (≤{max_tokens} tokens)",
     DEFAULT_TEXT,
 )
 
+inputs = tokenizer([text])
+[input_ids] = inputs["input_ids"]
+
+if len(input_ids) < 2:
+    st.error("Please enter at least 2 tokens.", icon="🚨")
+    st.stop()
+if len(input_ids) > max_tokens:
+    st.error(
+        f"Your input has {len(input_ids)} tokens. Please enter at most {max_tokens} tokens "
+        f"or try reducing the window size.",
+        icon="🚨"
+    )
+    st.stop()
+
 if metric_name == "KL divergence":
     st.error("KL divergence is not supported yet. Stay tuned!", icon="😭")
     st.stop()
 
 with st.spinner("Loading model…"):
-    tokenizer = st.cache_resource(AutoTokenizer.from_pretrained, show_spinner=False)(model_name)
     model = st.cache_resource(AutoModelForCausalLM.from_pretrained, show_spinner=False)(model_name)
 
-inputs = tokenizer([text])
-[input_ids] = inputs["input_ids"]
 window_len = min(window_len, len(input_ids))
 
-if len(input_ids) < 2:
-    st.error("Please enter at least 2 tokens.", icon="🚨")
-    st.stop()
-
 @st.cache_data(show_spinner=False)
 @torch.inference_mode()
-def get_logits(_model, _inputs, cache_key):
+def get_logprobs(_model, _inputs, cache_key):
     del cache_key
-    return _model(**_inputs).logits.to(torch.float16)
+    return _model(**_inputs).logits.log_softmax(dim=-1).to(torch.float16)
 
 @st.cache_data(show_spinner=False)
 @torch.inference_mode()
@@ -108,7 +129,7 @@ def run_context_length_probing(_model, _tokenizer, _inputs, window_len, cache_ke
         pad_id=_tokenizer.eos_token_id
     ).convert_to_tensors("pt")
 
-    logits = []
+    logprobs = []
     with st.spinner("Running model…"):
         batch_size = 8
         num_items = len(inputs_sliding["input_ids"])
@@ -116,27 +137,26 @@ def run_context_length_probing(_model, _tokenizer, _inputs, window_len, cache_ke
         for i in range(0, num_items, batch_size):
             pbar.progress(i / num_items, f"{i}/{num_items}")
             batch = {k: v[i:i + batch_size] for k, v in inputs_sliding.items()}
-            logits.append(
-                get_logits(
+            logprobs.append(
+                get_logprobs(
                     _model,
                     batch,
                     cache_key=(model_name, batch["input_ids"].cpu().numpy().tobytes())
                 )
             )
-        logits = torch.cat(logits, dim=0)
+        logprobs = torch.cat(logprobs, dim=0)
         pbar.empty()
 
     with st.spinner("Computing scores…"):
-        logits = logits.permute(1, 0, 2)
-        logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
-        logits = logits.view(-1, logits.shape[-1])[:-window_len]
-        logits = logits.view(window_len, len(input_ids) + window_len - 2, logits.shape[-1])
+        logprobs = logprobs.permute(1, 0, 2)
+        logprobs = F.pad(logprobs, (0, 0, 0, window_len, 0, 0), value=torch.nan)
+        logprobs = logprobs.view(-1, logprobs.shape[-1])[:-window_len]
+        logprobs = logprobs.view(window_len, len(input_ids) + window_len - 2, logprobs.shape[-1])
 
-        scores = logits.to(torch.float32).log_softmax(dim=-1)
-        scores = scores[:, torch.arange(len(input_ids[1:])), input_ids[1:]]
+        scores = logprobs[:, torch.arange(len(input_ids[1:])), input_ids[1:]]
         scores = scores.diff(dim=0).transpose(0, 1)
         scores = scores.nan_to_num()
-        scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-9
+        scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-6
         scores = scores.to(torch.float16)
 
     return scores