only compute acc on tokens whose label is not IGNORE_TOKEN_ID

Signed-off-by: Ye Yu <[email protected]>

Author: yeyu-nvidia

modelopt/torch/speculative/plugins/megatron_eagle.py

@@ -68,6 +68,13 @@
 except ImportError:
     warnings.warn("Fail to import megatron.core.post_training! EAGLE feature will be disable!")
 
+try:
+    from transformers.trainer_pt_utils import LabelSmoother
+
+    IGNORE_TOKEN_ID = LabelSmoother.ignore_index
+except ImportError:
+    IGNORE_TOKEN_ID = -100
+
 
 def dict_to_config(
     architecture_config,
@@ -1145,7 +1152,7 @@ def forward(
                             eagle_top1 += self.eagle_module.d2t[eagle_top1]
                         top1_p = (
                             torch.eq(labels[:, i + ttt_step + 1 :], eagle_top1).sum()
-                            / eagle_top1.numel()
+                            / (labels != IGNORE_TOKEN_ID).sum().item()
                         )
                         acc.append(top1_p)