add debug flag

modeling_llama.py

@@ -195,9 +195,10 @@ class LoRAMoeLayer(torch.nn.Module):
         weights, selected_experts = torch.topk(
             gate_logits, self.num_experts_per_tok
         )  # b,s,n
-        if self._layer_idx == 0 or self._layer_idx == 16 or self._layer_idx == 31:
-            print(f"{self._name}_{self._layer_idx}: {selected_experts}")
-            print("-"*10)
+        if hasattr(config, "show_debug") and config["show_debug"] == True:
+            if self._layer_idx == 0 or self._layer_idx == 16 or self._layer_idx == 31:
+                print(f"{self._name}_{self._layer_idx}: {selected_experts}")
+                print("-"*10)
         weights = F.softmax(weights, dim=2, dtype=torch.float).to(
             previous_dtype
         )  # b,s,n