Merge pull request facebookresearch#2 from ustcwhy/patch-1

Update fairseq_criterion.py
sunyt32 · Nov 14, 2023 · 7b370e1 · 7b370e1
2 parents 8069d60 + 87e6d91
commit 7b370e1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/fairseq/criterions/fairseq_criterion.py b/fairseq/criterions/fairseq_criterion.py
@@ -164,6 +164,10 @@ class MoECriterion(FairseqCriterion):
         "unused_expert2_count",    # average number of 2nd experts which process no tokens
         "all_to_all_cpu_time_ms",  # CPU time spent in all to all calls in milliseconds
         "all_to_all_cuda_time_ms", # CUDA ttime spent in all to all calls in milliseconds
+        "moe_second_expert_threshold",
+        "expert1_token",
+        "expert2_token",
+        "moe_expert_noise_threshold",
     ]
     def __init__(self, task, moe_gate_loss_wt, moe_gate_loss_combine_method, moe_gate_loss_transform, sentence_avg):
         super().__init__(task)