[Fix] dispatch internlm rote (#530)

dispatch internlm rote
InternLM · Mar 29, 2024 · 61b9451 · 61b9451
1 parent 7de581c
commit 61b9451
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/xtuner/model/modules/dispatch/__init__.py b/xtuner/model/modules/dispatch/__init__.py
@@ -305,12 +305,12 @@ def dispatch_modules(model, use_varlen_attn=False):
         dispatch_internlm2_attn_forward(model, use_varlen_attn)
         if USE_TRITON_KERNEL:
             dispatch_internlm2_rmsnorm_forward(model)
-        # replace_internlm2_rote(model)
+        replace_internlm2_rote(model)
     elif 'internlm' in model_name:
         dispatch_internlm_attn_forward(model, use_varlen_attn)
         if USE_TRITON_KERNEL:
             dispatch_internlm_rmsnorm_forward(model)
-        # replace_internlm_rote(model)
+        replace_internlm_rote(model)
     elif 'llama' in model_name:
         dispatch_llama_attn_forward(model, use_varlen_attn)
         if USE_TRITON_KERNEL: