From f37a06e7a82b8a503e7e8b55fd9abd003f5be3c0 Mon Sep 17 00:00:00 2001 From: liusong1222 Date: Thu, 4 Dec 2025 21:51:13 +0800 Subject: [PATCH] fix: remove the wrong attention mask updating --- .../compressor/speculative/train/trainer/eagle3_trainer.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py index d17c2b98..88b85e25 100644 --- a/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py +++ b/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py @@ -208,12 +208,6 @@ def draft_model_training_time_test( target_logits = padding(target_logits, left=False) loss_mask = padding(loss_mask, left=False) - # Update attention mask to prevent attending to future positions - ind = torch.arange(seq_length, device=attention_mask.device) - attention_mask[:, :, ind[idx:], ind[: seq_length - idx]] = torch.finfo( - attention_mask.dtype - ).min - # Step 8: Compute weighted loss ploss_weight = [0.8**i for i in range(len(plosses))] ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))])