fix: more specific warning

Signed-off-by: Will Johnson <mwjohnson728@gmail.com>

Author: willmj

tuning/sft_trainer.py

@@ -171,17 +171,19 @@ def train(
         fast_moe_config = None
     if fast_moe_config is not None:
         # If LoRA with ScatterMoE detected, raise warning
+        accepted_layers = ["all-linear"]
         if (
             peft_config is not None
             and hasattr(peft_config, "target_modules")
             and fast_moe_config.fast_moe is not None
+            and peft_config.target_modules != accepted_layers
         ):
             logger.warning(
                 "You are running lora with the ScatterMoE plugin, please note that "
                 "passing target modules that are part of the moe module can cause unexpected "
                 "behaviors and unsuccessful tuning while LoRA tuning with ScatterMoE. "
                 "For safe tuning, only pass linear modules such as those in the attn layer "
-                "(i.e. ['q_proj', 'v_proj', 'o_proj', 'k_proj']) or pass 'all-linear'"
+                "(i.e. ['q_proj', 'v_proj', 'o_proj', 'k_proj'])"
             )
 
     task_type = "CAUSAL_LM"