Let's dump extra tensors

Author: pwilkin

examples/model-conversion/scripts/causal/run-org-model-multi-token.py

@@ -40,9 +40,11 @@
 ### == END ROPE DEBUG ===
 
 token_counter = {}
+layer_counter = {}
+num_model_layers = 0
 
 def summarize(tensor: torch.Tensor, name: str, max_seq: int = 3, max_vals: int = 3):
-    global token, token_counter
+    global num_model_layers, layer_counter, token_counter
     """
     Print a tensor in llama.cpp debug style.
 
@@ -120,15 +122,27 @@ def summarize(tensor: torch.Tensor, name: str, max_seq: int = 3, max_vals: int =
     print("                                     ]")
     print(f"                                     sum = {t.sum().item():.6f}\n")
 
-    pattern = r"model\.layers\.[0-9]+_out"
-    pattern2 = r"recurrent_cache_[0-9]+"
-    if re.fullmatch(pattern, name) or re.fullmatch(pattern2, name):
+    indexed_patterns = [ r"model\.layers\.[0-9]+_out", r"recurrent_cache_[0-9]+" ]
+    non_indexed_patterns = [ r"k_pad", r"v_pad", r"q_pad" ]
+    
+    if any(re.fullmatch(p, name) for p in indexed_patterns):
         if name not in token_counter:
             token_counter[name] = 1
         else:
             token_counter[name] = token_counter[name] + 1
         save_tensor(t, f"reference/tensors/org/{name}_{token_counter[name]}.bin")
 
+    if any(re.fullmatch(p, name) for p in non_indexed_patterns):
+        if name not in token_counter:
+            token_counter[name] = 1
+        else:
+            token_counter[name] = token_counter[name] + 1
+        if name not in layer_counter or layer_counter[name] == num_model_layers - 1:
+            layer_counter[name] = 0
+        else:
+            layer_counter[name] = layer_counter[name] + 1
+        save_tensor(t, f"reference/tensors/org/{name}_{layer_counter[name] - 1}_{token_counter[name]}.bin")
+
 from transformers.models.qwen3_next.modeling_qwen3_next import torch_causal_conv1d_update, apply_rotary_pos_emb, l2norm  # noqa: E402
 orig_conv1d_update = torch_causal_conv1d_update
 orig_rope = apply_rotary_pos_emb
@@ -223,10 +237,8 @@ def patched_torch_chunk_gated_delta_rule(
     chunk_size=64,
     initial_state=None,
     output_final_state=False,
-    use_qk_l2norm_in_kernel=False,
-    long=False
+    use_qk_l2norm_in_kernel=False
 ):
-    torch.set_printoptions(threshold=10_000_000, sci_mode=False, precision=10, linewidth=200)
     initial_dtype = query.dtype
     [ summarize(x, y) for (x, y) in ((query, "q_prenorm"), (key, "k_prenorm")) ]
     if use_qk_l2norm_in_kernel:
@@ -359,13 +371,10 @@ def patched_torch_chunk_gated_delta_rule(
     core_attn_out = core_attn_out[:, :, :num_heads]
     core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
     summarize(core_attn_out, "attn_out")
-    if not long:
-        print(f"attn_out:\n{core_attn_out}\n\n")
 
     if isinstance(last_recurrent_state, torch.Tensor):
         summarize(last_recurrent_state, "state_out")
-        if not long:
-            print(f"state_out:\n{last_recurrent_state}\n\n")
+
     return core_attn_out, last_recurrent_state
 
 
@@ -667,6 +676,8 @@ def save_kv_cache(past_key_values, step_num, data_dir, model_name):
 print("BOS token id:     ", config.bos_token_id)
 print("EOS token id:     ", config.eos_token_id)
 
+num_model_layers = config.num_hidden_layers
+
 print("Loading model and tokenizer using AutoTokenizer:", model_path)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 config = AutoConfig.from_pretrained(model_path)

tools/main/main.cpp

@@ -246,6 +246,9 @@ static bool ggml_debug(struct ggml_tensor * t, bool ask, void * user_data) {
         uint8_t * data = is_host ? (uint8_t *) t->data : cb_data->data.data();
         std::string tensor_name(t->name);
         if (std::string(tensor_name).substr(0, std::string("post_moe-").size()) == "post_moe-" || 
+            std::string(tensor_name).substr(0, std::string("k_pad-").size()) == "k_pad-" || 
+            std::string(tensor_name).substr(0, std::string("q_pad-").size()) == "q_pad-" || 
+            std::string(tensor_name).substr(0, std::string("v_pad-").size()) == "v_pad-" || 
             std::string(tensor_name).substr(0, std::string("state_1d-").size()) == "state_1d-") {
 
             if (cb_data->tensors.count(tensor_name) == 0) {