refactor

.gitignore

@@ -132,7 +132,6 @@ runs_bak/
 LLM_ALERT
 small_demo/
 7b_llama_nopp/
-test/
 
 # Pytorch
 *.pth

.pre-commit-config.yaml

@@ -50,4 +50,4 @@ repos:
             [
                 '--rcfile=.pylintrc',
                 '--disable=C0114,C0415,W0212,W0235,W0238,W0621,C0103,R1735,C2801,E0402,C0412,W0719,R1728,W1514,W0718,W0105,W0707,C0209,W0703,W1203'
-            ]
+            ]

tests/test_training/test_loss.py

@@ -186,11 +186,11 @@ def train(
         # do forward and backward
         timer("fwd-bwd").start()
 
-        _, _, loss, _ = trainer.execute_schedule(batch, forward_only=False, return_loss=True, return_output_label=False)
+        _, _, loss, moe_loss = trainer.execute_schedule(batch, forward_only=False, return_loss=True, return_output_label=False)
         if gpc.is_rank_for_log():
             assert loss is not None and not math.isnan(loss.item())
             global cur_loss_list
-            cur_loss_list.append(loss.item())
+            cur_loss_list.append(loss.item() - moe_loss.item())
         timer("fwd-bwd").stop()
 
         # update parameters, and returns (success_update, grad_norm)