Eagle3: compact supervised positions before vocabulary projection

README.md

@@ -38,6 +38,23 @@ bash scripts/train/train.sh
 
 Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce `CUDA_VISIBLE_DEVICES`.
 
+### Compact Eagle3 loss positions
+
+Eagle3 can optionally compact supervised positions before the draft norm / LM
+head and before teacher-probability conversion:
+
+```bash
+python train.py --config config/eagle3/eagle3_qwen3_8b.py \
+  --opts model.trim_loss_positions=true \
+  --opts data.target_cache_path=/path/to/target_cache
+```
+
+This option is disabled by default. It preserves the original full-length
+teacher LM-head projection and every full-length TTT backbone step, including
+Q/K/V projection, attention, cache updates, output projection, and MLP. Only
+the vocabulary-side metric and loss path is compacted, and the loss keeps the
+original `local_mean` denominator (`batch_size * sequence_length`).
+
 
 ## Evaluation
 

config/eagle3/eagle3_gemma4_12b.py

@@ -12,6 +12,7 @@
     target_layer_ids=[5, 17, 29, 41, 46],
     ttt_length=7,
     step_loss_decay=0.8,
+    trim_loss_positions=False,
     draft_num_hidden_layers=1,
 )
 

config/eagle3/eagle3_qwen3_14b.py

@@ -13,6 +13,7 @@
     target_layer_ids=[1, 10, 19, 28, 37],
     ttt_length=7,
     step_loss_decay=0.8,
+    trim_loss_positions=False,
     draft_num_hidden_layers=1,
 )
 

config/eagle3/eagle3_qwen3_4b.py

@@ -12,6 +12,7 @@
     target_layer_ids=[1, 9, 17, 25, 33],
     ttt_length=7,
     step_loss_decay=0.8,
+    trim_loss_positions=False,
     draft_num_hidden_layers=1,
 )
 

config/eagle3/eagle3_qwen3_8b.py

@@ -13,6 +13,7 @@
     target_layer_ids=[1, 9, 17, 25, 33],
     ttt_length=7,
     step_loss_decay=0.8,
+    trim_loss_positions=False,
     draft_num_hidden_layers=1,
 )
 

deepspec/modeling/eagle3/gemma4/config.py

@@ -82,6 +82,9 @@ def build_draft_config(*, target_config, model_args):
     draft_config.target_layer_ids = target_layer_ids
     draft_config.ttt_length = ttt_length
     draft_config.step_loss_decay = step_loss_decay
+    draft_config.trim_loss_positions = bool(
+        getattr(model_args, "trim_loss_positions", False)
+    )
     draft_config.draft_num_hidden_layers = draft_num_hidden_layers
     draft_config.tie_word_embeddings = False
     draft_config._attn_implementation = TRAIN_ATTN_IMPLEMENTATION

deepspec/modeling/eagle3/gemma4/modeling.py

@@ -408,6 +408,7 @@ def forward(
         target_logits_only: bool = False,
         return_logits: bool = False,
         rope_cache_step_offset: bool = False,
+        logit_indices: Optional[torch.LongTensor] = None,
         **kwargs,
     ) -> torch.Tensor | Eagle3ForwardOutput:
         if target_logits_only:
@@ -473,7 +474,14 @@ def forward(
                 **kwargs,
             )
         if return_logits:
-            draft_logits = self.compute_logits(hidden_states)
+            logit_hidden_states = hidden_states
+            if logit_indices is not None:
+                hidden_size = int(hidden_states.shape[-1])
+                logit_hidden_states = hidden_states.reshape(
+                    -1, hidden_size
+                ).index_select(0, logit_indices)
+                logit_hidden_states = logit_hidden_states.unsqueeze(0)
+            draft_logits = self.compute_logits(logit_hidden_states)
             target_logits = None
             if target_last_hidden_states is not None:
                 with torch.no_grad():