From c7fc37d089302baed6b733b281a091f93e716884 Mon Sep 17 00:00:00 2001 From: julyanghar Date: Sat, 29 Aug 2026 21:16:57 -0500 Subject: [PATCH 1/2] Add supervised-position compaction for Eagle3 --- README.md | 17 + config/eagle3/eagle3_gemma4_12b.py | 1 + config/eagle3/eagle3_qwen3_14b.py | 1 + config/eagle3/eagle3_qwen3_4b.py | 1 + config/eagle3/eagle3_qwen3_8b.py | 1 + deepspec/modeling/eagle3/gemma4/config.py | 3 + deepspec/modeling/eagle3/gemma4/modeling.py | 10 +- deepspec/modeling/eagle3/loss.py | 426 +++++++++++++++--- deepspec/modeling/eagle3/qwen3/config.py | 3 + deepspec/modeling/eagle3/qwen3/modeling.py | 10 +- deepspec/trainer/eagle3_trainer.py | 3 + .../benchmark_eagle3_trim_loss_positions.py | 279 ++++++++++++ .../run_eagle3_trim_loss_positions_fsdp.py | 211 +++++++++ tests/test_eagle3_trim_loss_positions.py | 376 ++++++++++++++++ 14 files changed, 1269 insertions(+), 73 deletions(-) create mode 100644 scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py create mode 100644 tests/distributed/run_eagle3_trim_loss_positions_fsdp.py create mode 100644 tests/test_eagle3_trim_loss_positions.py diff --git a/README.md b/README.md index 183501d8..7a3b2104 100644 --- a/README.md +++ b/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 diff --git a/config/eagle3/eagle3_gemma4_12b.py b/config/eagle3/eagle3_gemma4_12b.py index 1bc636ca..be4760e8 100644 --- a/config/eagle3/eagle3_gemma4_12b.py +++ b/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, ) diff --git a/config/eagle3/eagle3_qwen3_14b.py b/config/eagle3/eagle3_qwen3_14b.py index ce6b6dff..f406c1c6 100644 --- a/config/eagle3/eagle3_qwen3_14b.py +++ b/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, ) diff --git a/config/eagle3/eagle3_qwen3_4b.py b/config/eagle3/eagle3_qwen3_4b.py index 03b78687..322cebec 100644 --- a/config/eagle3/eagle3_qwen3_4b.py +++ b/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, ) diff --git a/config/eagle3/eagle3_qwen3_8b.py b/config/eagle3/eagle3_qwen3_8b.py index fb577b98..23feb836 100644 --- a/config/eagle3/eagle3_qwen3_8b.py +++ b/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, ) diff --git a/deepspec/modeling/eagle3/gemma4/config.py b/deepspec/modeling/eagle3/gemma4/config.py index 5b0b1f64..7cdc0a36 100644 --- a/deepspec/modeling/eagle3/gemma4/config.py +++ b/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 diff --git a/deepspec/modeling/eagle3/gemma4/modeling.py b/deepspec/modeling/eagle3/gemma4/modeling.py index 5dff1f36..3d4da5d4 100644 --- a/deepspec/modeling/eagle3/gemma4/modeling.py +++ b/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(): diff --git a/deepspec/modeling/eagle3/loss.py b/deepspec/modeling/eagle3/loss.py index 163081d1..dd126a47 100644 --- a/deepspec/modeling/eagle3/loss.py +++ b/deepspec/modeling/eagle3/loss.py @@ -12,6 +12,8 @@ PyTorch path retains across TTT steps. """ +from dataclasses import dataclass + import torch import triton import triton.language as tl @@ -20,6 +22,34 @@ from deepspec.utils.metrics import add_metric +@dataclass(frozen=True) +class Eagle3CompactStep: + """Row mappings for one A-level-compacted TTT step. + + ``draft_indices`` indexes the full ``[B, T, H]`` draft hidden states after + the backbone has run. ``teacher_table_indices`` indexes a table containing + teacher probabilities for unique supervised target rows. Empty steps carry + one zero-masked dummy row so every rank still executes the same modules. + """ + + draft_indices: torch.Tensor + teacher_indices: torch.Tensor + teacher_table_indices: torch.Tensor + position_mask: torch.Tensor + scatter_indices: torch.Tensor + num_valid_rows: int + + +@dataclass(frozen=True) +class Eagle3CompactPlan: + """A-level row plan that leaves teacher projection and backbone full-length.""" + + teacher_unique_indices: torch.Tensor + steps: tuple[Eagle3CompactStep, ...] + batch_size: int + seq_len: int + + def _calculate_settings(n: int): # BLOCK_SIZE is the per-iteration chunk; the kernel loops over V in # chunks of BLOCK_SIZE, so large vocabularies (e.g. Qwen3 151936) do @@ -73,6 +103,103 @@ def _compute_loss_normalizers( ] +@torch.no_grad() +def _build_compact_plan( + *, + position_masks: list[torch.Tensor], +) -> Eagle3CompactPlan: + """Build exact row mappings for A-level supervised-position compaction. + + At TTT step ``j``, a live draft row ``r`` is supervised by target-model + logits at row ``r + j + 1``. The full path expresses this with a padded, + shifted target-probability tensor. This plan represents the same mapping + directly, without changing the full-length recurrent backbone. + """ + + assert position_masks, "position_masks must not be empty." + batch_size, seq_len = position_masks[0].shape[:2] + device = position_masks[0].device + raw_steps = [] + all_teacher_indices = [] + + for step_idx, position_mask in enumerate(position_masks): + assert tuple(position_mask.shape[:2]) == (batch_size, seq_len) + flat_mask = position_mask.reshape(-1) + draft_indices = (flat_mask > 0).nonzero(as_tuple=False).squeeze(-1) + batch_indices = torch.div(draft_indices, seq_len, rounding_mode="floor") + row_indices = draft_indices.remainder(seq_len) + teacher_rows = row_indices + int(step_idx) + 1 + # `_build_next_token_position_mask` clears every sequence's final valid + # row, and each later mask is shifted left once. Therefore every live + # row satisfies row + step_idx + 1 < seq_len. Keep this on-device; + # converting an `.all()` result to bool would synchronize every step. + teacher_indices = batch_indices * seq_len + teacher_rows + raw_steps.append((draft_indices, teacher_indices, flat_mask)) + if teacher_indices.numel() > 0: + all_teacher_indices.append(teacher_indices) + + if all_teacher_indices: + teacher_unique_indices = torch.unique( + torch.cat(all_teacher_indices), sorted=True + ) + else: + # Keep the full teacher projection unchanged, and select one dummy row + # after it so empty-loss batches still have a valid compact tensor. + teacher_unique_indices = torch.zeros(1, dtype=torch.long, device=device) + + steps = [] + dummy_index = torch.zeros(1, dtype=torch.long, device=device) + for draft_indices, teacher_indices, flat_mask in raw_steps: + num_valid_rows = int(draft_indices.numel()) + if num_valid_rows == 0: + padded_draft_indices = dummy_index + padded_teacher_indices = teacher_unique_indices[:1] + teacher_table_indices = dummy_index + compact_position_mask = flat_mask.new_zeros((1, 1, 1)) + else: + padded_draft_indices = draft_indices + padded_teacher_indices = teacher_indices + teacher_table_indices = torch.searchsorted( + teacher_unique_indices, teacher_indices + ) + compact_position_mask = flat_mask.index_select( + 0, draft_indices + ).reshape(1, -1, 1) + steps.append( + Eagle3CompactStep( + draft_indices=padded_draft_indices, + teacher_indices=padded_teacher_indices, + teacher_table_indices=teacher_table_indices, + position_mask=compact_position_mask, + scatter_indices=draft_indices, + num_valid_rows=num_valid_rows, + ) + ) + + return Eagle3CompactPlan( + teacher_unique_indices=teacher_unique_indices, + steps=tuple(steps), + batch_size=int(batch_size), + seq_len=int(seq_len), + ) + + +@torch.no_grad() +def _build_compact_teacher_probs( + *, + target_logits: torch.Tensor, + compact_plan: Eagle3CompactPlan, +) -> torch.Tensor: + """Convert only unique supervised teacher-logit rows to fp32 probabilities.""" + + vocab_size = int(target_logits.shape[-1]) + target_logits_flat = target_logits.reshape(-1, vocab_size) + compact_logits = target_logits_flat.index_select( + 0, compact_plan.teacher_unique_indices + ) + return torch.softmax(compact_logits.float(), dim=-1).detach() + + @torch.no_grad() def _build_next_token_position_mask( *, @@ -110,6 +237,8 @@ def _log_eagle3_step_metrics( draft_logits: torch.Tensor, target_probs: torch.Tensor, position_mask: torch.Tensor, + full_shape: tuple[int, int] | None = None, + scatter_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: valid_mask = position_mask.squeeze(-1) > 0 correct_mask = (draft_logits.argmax(-1) == target_probs.argmax(-1)) & valid_mask @@ -131,6 +260,24 @@ def _log_eagle3_step_metrics( reduction="dp_sum", tag="train", ) + if scatter_indices is not None: + assert full_shape is not None + num_valid_rows = int(scatter_indices.numel()) + full_numel = int(full_shape[0]) * int(full_shape[1]) + + def scatter(values: torch.Tensor) -> torch.Tensor: + output = values.new_zeros(full_numel) + if num_valid_rows > 0: + output.index_copy_( + 0, + scatter_indices, + values.reshape(-1)[:num_valid_rows], + ) + return output.reshape(full_shape) + + correct_mask = scatter(correct_mask) + accept_rate_mask = scatter(accept_rate_mask) + valid_mask = scatter(valid_mask) return correct_mask, accept_rate_mask, valid_mask @@ -279,6 +426,71 @@ def _log_softmax_backward_kernel( tl.store(logits_ptr + offsets, grad_block.to(tl.float32), mask=mask) +def _fused_log_softmax_forward(logits, target_p, position_mask, normalizer): + assert logits.is_cuda, "FusedLogSoftmaxLoss requires CUDA tensors." + assert logits.shape == target_p.shape + assert position_mask.shape[:2] == logits.shape[:2] + B, T, V = logits.shape + loss = torch.zeros((B * T, 1), device=logits.device, dtype=torch.float32) + logits_flat = logits.contiguous().view(B * T, V) + target_flat = target_p.contiguous().view(B * T, V) + position_mask_flat = position_mask.contiguous().view(B * T, 1).bool() + grid = (B * T,) + m = torch.zeros((B * T,), device=logits.device, dtype=torch.float32) + d = torch.zeros((B * T,), device=logits.device, dtype=torch.float32) + BLOCK_SIZE, num_warps = _calculate_settings(V) + _log_softmax_forward_kernel[grid]( + logits_flat, + logits_flat.stride(0), + target_flat, + target_flat.stride(0), + position_mask_flat, + position_mask_flat.stride(0), + loss, + loss.stride(0), + m, + d, + V, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + return loss.squeeze(1).sum() / float(normalizer), m, d + + +def _fused_log_softmax_backward( + *, + logits, + target_p, + position_mask, + m, + d, + normalizer, + grad_output, +): + B, T, V = logits.shape + scaling_factor = 1.0 / float(normalizer) + logits_flat = logits.contiguous().view(B * T, V) + target_flat = target_p.contiguous().view(B * T, V) + position_mask_flat = position_mask.contiguous().view(B * T, 1).bool() + grid = (B * T,) + BLOCK_SIZE, num_warps = _calculate_settings(V) + _log_softmax_backward_kernel[grid]( + logits_flat, + logits_flat.stride(0), + target_flat, + target_flat.stride(0), + position_mask_flat, + grad_output, + scaling_factor, + m, + d, + V, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + return logits_flat.view(B, T, V) + + class FusedLogSoftmaxLoss(torch.autograd.Function): """Soft cross-entropy with fused Triton forward/backward. @@ -293,62 +505,88 @@ class FusedLogSoftmaxLoss(torch.autograd.Function): @staticmethod def forward(ctx, logits, target_p, position_mask, normalizer): - assert logits.is_cuda, "FusedLogSoftmaxLoss requires CUDA tensors." - assert logits.shape == target_p.shape - assert position_mask.shape[:2] == logits.shape[:2] - B, T, V = logits.shape - loss = torch.zeros((B * T, 1), device=logits.device, dtype=torch.float32) - logits_flat = logits.contiguous().view(B * T, V) - target_flat = target_p.contiguous().view(B * T, V) - position_mask_flat = position_mask.contiguous().view(B * T, 1).bool() - grid = (B * T,) - m = torch.zeros((B * T,), device=logits.device, dtype=torch.float32) - d = torch.zeros((B * T,), device=logits.device, dtype=torch.float32) - BLOCK_SIZE, num_warps = _calculate_settings(V) - _log_softmax_forward_kernel[grid]( - logits_flat, - logits_flat.stride(0), - target_flat, - target_flat.stride(0), - position_mask_flat, - position_mask_flat.stride(0), - loss, - loss.stride(0), - m, - d, - V, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, + loss, m, d = _fused_log_softmax_forward( + logits, + target_p, + position_mask, + normalizer, ) ctx.save_for_backward(logits.detach(), target_p, position_mask, m, d) ctx.normalizer = float(normalizer) - return loss.squeeze(1).sum() / ctx.normalizer + return loss @staticmethod def backward(ctx, grad_output): logits, target, position_mask, m, d = ctx.saved_tensors - B, T, V = logits.shape - scaling_factor = 1.0 / ctx.normalizer - logits_flat = logits.contiguous().view(B * T, V) - target_flat = target.contiguous().view(B * T, V) - position_mask_flat = position_mask.contiguous().view(B * T, 1).bool() - grid = (B * T,) - BLOCK_SIZE, num_warps = _calculate_settings(V) - _log_softmax_backward_kernel[grid]( - logits_flat, - logits_flat.stride(0), - target_flat, - target_flat.stride(0), - position_mask_flat, - grad_output, - scaling_factor, + grad_logits = _fused_log_softmax_backward( + logits=logits, + target_p=target, + position_mask=position_mask, + m=m, + d=d, + normalizer=ctx.normalizer, + grad_output=grad_output, + ) + return grad_logits, None, None, None + + +class FusedIndexedLogSoftmaxLoss(torch.autograd.Function): + """Fused loss whose target rows are views of one shared compact table. + + Forward and backward materialize only the current step's indexed rows. The + autograd context saves the shared table plus row indices, rather than one + ``[N_step, V]`` target copy for every TTT step. + """ + + @staticmethod + def forward( + ctx, + logits, + target_table, + target_indices, + position_mask, + normalizer, + ): + assert logits.shape[0] == 1 + target_p = target_table.index_select(0, target_indices).unsqueeze(0) + loss, m, d = _fused_log_softmax_forward( + logits, + target_p, + position_mask, + normalizer, + ) + ctx.save_for_backward( + logits.detach(), + target_table, + target_indices, + position_mask, m, d, - V, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, ) - return logits_flat.view(B, T, V), None, None, None + ctx.normalizer = float(normalizer) + return loss + + @staticmethod + def backward(ctx, grad_output): + ( + logits, + target_table, + target_indices, + position_mask, + m, + d, + ) = ctx.saved_tensors + target_p = target_table.index_select(0, target_indices).unsqueeze(0) + grad_logits = _fused_log_softmax_backward( + logits=logits, + target_p=target_p, + position_mask=position_mask, + m=m, + d=d, + normalizer=ctx.normalizer, + grad_output=grad_output, + ) + return grad_logits, None, None, None, None def compute_eagle3_loss( @@ -357,6 +595,7 @@ def compute_eagle3_loss( batch: dict[str, torch.Tensor], ttt_length: int, step_loss_decay: float, + trim_loss_positions: bool = False, ) -> torch.Tensor: input_ids = batch["input_ids"].long() attention_mask = batch["attention_mask"].long() @@ -373,18 +612,6 @@ def compute_eagle3_loss( loss_mask=batch["loss_mask"], attention_mask=attention_mask, ) - past_key_values = DynamicCache() - total_loss = hidden_states.new_zeros((), dtype=torch.float32) - with torch.no_grad(): - target_logits = model( - target_last_hidden_states=target_last_hidden_states, - target_logits_only=True, - ) - target_probs = _build_padded_next_token_target_probs( - target_logits=target_logits, - ttt_length=int(ttt_length), - ) - del target_logits position_masks = [] for _ in range(int(ttt_length)): position_masks.append(shifted_position_mask) @@ -395,18 +622,59 @@ def compute_eagle3_loss( loss_normalizers = _compute_loss_normalizers( position_masks=position_masks, ) + compact_plan = ( + _build_compact_plan(position_masks=position_masks) + if bool(trim_loss_positions) + else None + ) + + past_key_values = DynamicCache() + total_loss = hidden_states.new_zeros((), dtype=torch.float32) + with torch.no_grad(): + # A-level deliberately leaves this full-length teacher LM-head + # projection unchanged. Compaction starts only after target_logits. + target_logits = model( + target_last_hidden_states=target_last_hidden_states, + target_logits_only=True, + ) + if compact_plan is None: + target_probs = _build_padded_next_token_target_probs( + target_logits=target_logits, + ttt_length=int(ttt_length), + ) + compact_teacher_probs = None + else: + target_probs = None + compact_teacher_probs = _build_compact_teacher_probs( + target_logits=target_logits, + compact_plan=compact_plan, + ) + del target_logits correct_masks = [] accept_rate_masks = [] valid_masks = [] for step_idx in range(int(ttt_length)): - # Keep this slice alignment in sync with the Eagle3 reference. - target_step_probs = target_probs[ - :, - step_idx : step_idx + seq_len, - :, - ].contiguous() - position_mask_step = position_masks[step_idx] + if compact_plan is None: + # Keep this slice alignment in sync with the Eagle3 reference. + target_step_probs = target_probs[ + :, + step_idx : step_idx + seq_len, + :, + ].contiguous() + position_mask_step = position_masks[step_idx] + logit_indices = None + metric_scatter_indices = None + metric_full_shape = None + else: + compact_step = compact_plan.steps[step_idx] + target_step_probs = compact_teacher_probs.index_select( + 0, compact_step.teacher_table_indices + ).unsqueeze(0) + position_mask_step = compact_step.position_mask + logit_indices = compact_step.draft_indices + metric_scatter_indices = compact_step.scatter_indices + metric_full_shape = (compact_plan.batch_size, compact_plan.seq_len) output = model( hidden_states=hidden_states, input_ids=current_input_ids, @@ -416,6 +684,7 @@ def compute_eagle3_loss( use_cache=True, return_logits=True, rope_cache_step_offset=True, + logit_indices=logit_indices, ) hidden_states = output.hidden_states correct_mask, accept_rate_mask, valid_mask = _log_eagle3_step_metrics( @@ -423,16 +692,27 @@ def compute_eagle3_loss( draft_logits=output.draft_logits, target_probs=target_step_probs, position_mask=position_mask_step, + full_shape=metric_full_shape, + scatter_indices=metric_scatter_indices, ) correct_masks.append(correct_mask) accept_rate_masks.append(accept_rate_mask) valid_masks.append(valid_mask) - step_loss = FusedLogSoftmaxLoss.apply( - output.draft_logits, - target_step_probs, - position_mask_step, - loss_normalizers[step_idx], - ) + if compact_plan is None: + step_loss = FusedLogSoftmaxLoss.apply( + output.draft_logits, + target_step_probs, + position_mask_step, + loss_normalizers[step_idx], + ) + else: + step_loss = FusedIndexedLogSoftmaxLoss.apply( + output.draft_logits, + compact_teacher_probs, + compact_step.teacher_table_indices, + position_mask_step, + loss_normalizers[step_idx], + ) add_metric( f"ploss_{step_idx}", step_loss.detach(), @@ -453,6 +733,10 @@ def compute_eagle3_loss( __all__ = [ + "Eagle3CompactPlan", + "Eagle3CompactStep", + "FusedIndexedLogSoftmaxLoss", "FusedLogSoftmaxLoss", + "_build_compact_plan", "compute_eagle3_loss", ] diff --git a/deepspec/modeling/eagle3/qwen3/config.py b/deepspec/modeling/eagle3/qwen3/config.py index 2ac00ef5..ab571ada 100644 --- a/deepspec/modeling/eagle3/qwen3/config.py +++ b/deepspec/modeling/eagle3/qwen3/config.py @@ -33,6 +33,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 diff --git a/deepspec/modeling/eagle3/qwen3/modeling.py b/deepspec/modeling/eagle3/qwen3/modeling.py index 7190c838..785d0f0e 100644 --- a/deepspec/modeling/eagle3/qwen3/modeling.py +++ b/deepspec/modeling/eagle3/qwen3/modeling.py @@ -334,6 +334,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: @@ -394,7 +395,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(): diff --git a/deepspec/trainer/eagle3_trainer.py b/deepspec/trainer/eagle3_trainer.py index 70343a3b..17ee60d3 100644 --- a/deepspec/trainer/eagle3_trainer.py +++ b/deepspec/trainer/eagle3_trainer.py @@ -64,6 +64,9 @@ def run_batch(self, batch): batch=batch, ttt_length=int(self.draft_model.ttt_length), step_loss_decay=float(self.draft_model.step_loss_decay), + trim_loss_positions=bool( + getattr(self.draft_model.config, "trim_loss_positions", False) + ), ) diff --git a/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py b/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py new file mode 100644 index 00000000..b0cbfd19 --- /dev/null +++ b/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import gc +import json +import math +import os +import statistics +from types import SimpleNamespace + +import torch +from transformers import AutoConfig + +import deepspec.modeling.eagle3.loss as loss_module +from deepspec.modeling.eagle3.loss import compute_eagle3_loss +from deepspec.modeling.eagle3.qwen3.config import build_draft_config +from deepspec.modeling.eagle3.qwen3.modeling import Qwen3Eagle3Model + + +GIB = 1024**3 + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Benchmark Eagle3 A-level supervised-position compaction." + ) + parser.add_argument("--config-path", required=True) + parser.add_argument("--seq-len", type=int, required=True) + parser.add_argument("--ttt-length", type=int, required=True) + parser.add_argument("--densities", type=float, nargs="+", required=True) + parser.add_argument("--layout", choices=("tail",), required=True) + parser.add_argument("--warmup-pairs", type=int, required=True) + parser.add_argument("--measure-pairs", type=int, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--output", required=True) + return parser.parse_args() + + +def build_model(args, device): + target_config = AutoConfig.from_pretrained( + args.config_path, + local_files_only=True, + ) + model_args = SimpleNamespace( + target_model_name_or_path=args.config_path, + target_layer_ids=[1, 9, 17, 25, 33], + ttt_length=args.ttt_length, + step_loss_decay=0.8, + draft_num_hidden_layers=1, + trim_loss_positions=False, + ) + draft_config = build_draft_config( + target_config=target_config, + model_args=model_args, + ) + model = Qwen3Eagle3Model(draft_config).to( + device=device, + dtype=torch.bfloat16, + ) + model.embed_tokens.requires_grad_(False) + model.lm_head.requires_grad_(False) + model.train() + return model + + +def build_inputs(args, model, device): + generator = torch.Generator(device=device).manual_seed(args.seed) + hidden_size = int(model.config.hidden_size) + num_target_layers = len(model.target_layer_ids) + vocab_size = int(model.config.vocab_size) + return { + "input_ids": torch.randint( + 0, + vocab_size, + (1, args.seq_len), + generator=generator, + device=device, + ), + "attention_mask": torch.ones( + 1, args.seq_len, dtype=torch.long, device=device + ), + "target_hidden_states": torch.randn( + 1, + args.seq_len, + num_target_layers * hidden_size, + generator=generator, + device=device, + dtype=torch.bfloat16, + ), + "target_last_hidden_states": torch.randn( + 1, + args.seq_len, + hidden_size, + generator=generator, + device=device, + dtype=torch.bfloat16, + ), + } + + +def build_loss_mask(seq_len: int, density: float, device): + assert 0.0 < density <= 1.0 + # The final valid token has no in-sequence next-token teacher. Construct an + # exact contiguous tail over the remaining L-1 eligible positions. + eligible = seq_len - 1 + supervised = max(1, int(round(eligible * density))) + start = eligible - supervised + mask = torch.zeros(1, seq_len, dtype=torch.long, device=device) + mask[:, start:eligible] = 1 + return mask, supervised + + +def run_once(*, model, batch, ttt_length: int, trim: bool): + model.zero_grad(set_to_none=True) + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + loss = compute_eagle3_loss( + model=model, + batch=batch, + ttt_length=ttt_length, + step_loss_decay=0.8, + trim_loss_positions=trim, + ) + loss.backward() + end.record() + torch.cuda.synchronize() + elapsed_ms = float(start.elapsed_time(end)) + peak_bytes = int(torch.cuda.max_memory_allocated()) + loss_value = float(loss.detach().float().item()) + del loss + model.zero_grad(set_to_none=True) + return { + "time_ms": elapsed_ms, + "peak_allocated_bytes": peak_bytes, + "loss": loss_value, + } + + +def run_pairs(*, model, batch, ttt_length: int, count: int, offset: int): + pairs = [] + for pair_idx in range(count): + full_first = (pair_idx + offset) % 2 == 0 + order = (False, True) if full_first else (True, False) + observations = {} + for trim in order: + observations["trim" if trim else "full"] = run_once( + model=model, + batch=batch, + ttt_length=ttt_length, + trim=trim, + ) + pairs.append( + { + "order": ["full", "trim"] if full_first else ["trim", "full"], + "full": observations["full"], + "trim": observations["trim"], + "speedup": observations["full"]["time_ms"] + / observations["trim"]["time_ms"], + } + ) + return pairs + + +def summarize(pairs): + full_times = [pair["full"]["time_ms"] for pair in pairs] + trim_times = [pair["trim"]["time_ms"] for pair in pairs] + speedups = [pair["speedup"] for pair in pairs] + full_peaks = [pair["full"]["peak_allocated_bytes"] for pair in pairs] + trim_peaks = [pair["trim"]["peak_allocated_bytes"] for pair in pairs] + return { + "paired_median_speedup": statistics.median(speedups), + "median_full_time_ms": statistics.median(full_times), + "median_trim_time_ms": statistics.median(trim_times), + "median_full_peak_allocated_bytes": int(statistics.median(full_peaks)), + "median_trim_peak_allocated_bytes": int(statistics.median(trim_peaks)), + "median_peak_reduction_bytes": int( + statistics.median(full_peaks) - statistics.median(trim_peaks) + ), + } + + +def validate_results(results): + by_density = {entry["density"]: entry for entry in results} + required = (0.3, 0.6, 0.9) + assert all(density in by_density for density in required) + for entry in results: + for pair in entry["measured_pairs"]: + assert math.isfinite(pair["full"]["loss"]) + assert math.isfinite(pair["trim"]["loss"]) + + d30 = by_density[0.3]["summary"] + d60 = by_density[0.6]["summary"] + d90 = by_density[0.9]["summary"] + assert d30["paired_median_speedup"] >= 1.05 + assert d30["median_peak_reduction_bytes"] >= GIB + assert d60["paired_median_speedup"] >= 1.02 + assert d60["median_peak_reduction_bytes"] > 0 + assert d90["paired_median_speedup"] >= 0.98 + assert ( + d30["paired_median_speedup"] + >= d60["paired_median_speedup"] + >= d90["paired_median_speedup"] + ) + + +def main(): + args = parse_args() + assert torch.cuda.is_available(), "CUDA is required." + device = torch.device("cuda") + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + loss_module.add_metric = lambda *unused_args, **unused_kwargs: None + + model = build_model(args, device) + base_batch = build_inputs(args, model, device) + results = [] + for density_idx, density in enumerate(args.densities): + loss_mask, supervised = build_loss_mask(args.seq_len, density, device) + batch = dict(base_batch) + batch["loss_mask"] = loss_mask + run_pairs( + model=model, + batch=batch, + ttt_length=args.ttt_length, + count=args.warmup_pairs, + offset=density_idx, + ) + measured_pairs = run_pairs( + model=model, + batch=batch, + ttt_length=args.ttt_length, + count=args.measure_pairs, + offset=density_idx + args.warmup_pairs, + ) + entry = { + "density": float(density), + "effective_supervised_rows": supervised, + "measured_pairs": measured_pairs, + "summary": summarize(measured_pairs), + } + results.append(entry) + print( + f"density={density:.2f} " + f"speedup={entry['summary']['paired_median_speedup']:.4f}x " + "peak_reduction_gib=" + f"{entry['summary']['median_peak_reduction_bytes'] / GIB:.3f}", + flush=True, + ) + + payload = { + "config": { + "config_path": args.config_path, + "seq_len": args.seq_len, + "ttt_length": args.ttt_length, + "densities": args.densities, + "layout": args.layout, + "warmup_pairs": args.warmup_pairs, + "measure_pairs": args.measure_pairs, + "seed": args.seed, + "dtype": "bfloat16", + "device_name": torch.cuda.get_device_name(device), + }, + "results": results, + } + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + validate_results(results) + print(f"BENCHMARK_GATES_OK output={args.output}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/run_eagle3_trim_loss_positions_fsdp.py b/tests/distributed/run_eagle3_trim_loss_positions_fsdp.py new file mode 100644 index 00000000..2fddbe8b --- /dev/null +++ b/tests/distributed/run_eagle3_trim_loss_positions_fsdp.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import copy +import os + +import torch +import torch.distributed as dist +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config + +import deepspec.modeling.eagle3.loss as loss_module +from deepspec.modeling.eagle3.loss import compute_eagle3_loss +from deepspec.modeling.eagle3.qwen3.modeling import Qwen3Eagle3Model +from deepspec.trainer.base_trainer import _build_fsdp_kwargs + + +def tiny_config() -> Qwen3Config: + config = Qwen3Config( + vocab_size=97, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=64, + attention_dropout=0.0, + attention_bias=False, + ) + config.target_layer_ids = [0, 1, 2, 3, 4] + config.ttt_length = 3 + config.step_loss_decay = 0.8 + config.trim_loss_positions = False + config._attn_implementation = "flex_attention" + return config + + +def build_model(device: torch.device) -> Qwen3Eagle3Model: + model = Qwen3Eagle3Model(tiny_config()).to( + device=device, dtype=torch.bfloat16 + ) + model.embed_tokens.requires_grad_(False) + model.lm_head.requires_grad_(False) + model.train() + return model + + +def build_batch(device: torch.device, rank: int): + generator = torch.Generator(device=device).manual_seed(100 + rank) + batch_size, seq_len, hidden_size = 1, 16, 32 + loss_mask = torch.zeros( + batch_size, seq_len, dtype=torch.long, device=device + ) + if rank == 0: + loss_mask[0, [0, 2, 5, 7, 13]] = 1 + else: + # Steps 1 and 2 have no valid rows on this rank. The compact path must + # still execute a zero-masked dummy LM-head row without diverging in FSDP. + loss_mask[0, 0] = 1 + return { + "input_ids": torch.randint( + 0, + 97, + (batch_size, seq_len), + generator=generator, + device=device, + ), + "attention_mask": torch.ones( + batch_size, seq_len, dtype=torch.long, device=device + ), + "loss_mask": loss_mask, + "target_hidden_states": torch.randn( + batch_size, + seq_len, + 5 * hidden_size, + generator=generator, + device=device, + dtype=torch.bfloat16, + ), + "target_last_hidden_states": torch.randn( + batch_size, + seq_len, + hidden_size, + generator=generator, + device=device, + dtype=torch.bfloat16, + ), + } + + +def run_arm( + *, + initial_state: dict[str, torch.Tensor], + batch: dict[str, torch.Tensor], + device: torch.device, + world_size: int, + trim_loss_positions: bool, +): + model = build_model(device) + model.load_state_dict(initial_state) + fsdp_model = FSDP( + model, + **_build_fsdp_kwargs( + sharding_strategy_name="no_shard", + precision_dtype=torch.bfloat16, + world_size=world_size, + ), + ) + optimizer = torch.optim.SGD( + [parameter for parameter in fsdp_model.parameters() if parameter.requires_grad], + lr=1e-2, + ) + optimizer.zero_grad(set_to_none=True) + loss = compute_eagle3_loss( + model=fsdp_model, + batch=batch, + ttt_length=3, + step_loss_decay=0.8, + trim_loss_positions=trim_loss_positions, + ) + loss.backward() + optimizer.step() + dist.barrier() + + with FSDP.summon_full_params(fsdp_model, recurse=True, writeback=False): + parameters = { + name: parameter.detach().float().cpu().clone() + for name, parameter in fsdp_model.module.named_parameters() + if parameter.requires_grad + } + for parameter in fsdp_model.module.parameters(): + if not parameter.requires_grad: + continue + reference = parameter.detach().clone() + dist.broadcast(reference, src=0) + torch.testing.assert_close( + parameter, + reference, + rtol=5e-3, + atol=5e-3, + ) + + loss_value = float(loss.detach().float().item()) + del optimizer, fsdp_model, model + torch.cuda.empty_cache() + dist.barrier() + return loss_value, parameters + + +def main(): + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group(backend="nccl", device_id=device) + loss_module.add_metric = lambda *args, **kwargs: None + + torch.manual_seed(0) + template = build_model(device) + with torch.no_grad(): + template.lm_head.weight.mul_(16.0) + initial_state = copy.deepcopy(template.state_dict()) + del template + batch = build_batch(device, rank) + + full_loss, full_parameters = run_arm( + initial_state=initial_state, + batch=batch, + device=device, + world_size=world_size, + trim_loss_positions=False, + ) + trim_loss, trim_parameters = run_arm( + initial_state=initial_state, + batch=batch, + device=device, + world_size=world_size, + trim_loss_positions=True, + ) + + loss_tolerance = 5e-3 * max(abs(full_loss), abs(trim_loss)) + 1e-4 + assert abs(full_loss - trim_loss) <= loss_tolerance + assert full_parameters.keys() == trim_parameters.keys() + max_parameter_diff = 0.0 + for name in full_parameters: + max_parameter_diff = max( + max_parameter_diff, + float((full_parameters[name] - trim_parameters[name]).abs().max().item()), + ) + torch.testing.assert_close( + full_parameters[name], + trim_parameters[name], + rtol=5e-3, + atol=5e-3, + msg=lambda message, name=name: f"parameter {name}: {message}", + ) + + if rank == 0: + print( + "FSDP_EQUIV_OK " + f"max_parameter_diff={max_parameter_diff:.8g} " + "rank_consistent=true", + flush=True, + ) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/test_eagle3_trim_loss_positions.py b/tests/test_eagle3_trim_loss_positions.py new file mode 100644 index 00000000..2b6d9e75 --- /dev/null +++ b/tests/test_eagle3_trim_loss_positions.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +import copy +from dataclasses import replace +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from transformers import DynamicCache +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config + +import deepspec.modeling.eagle3.loss as loss_module +from deepspec.modeling.eagle3.loss import ( + _build_compact_plan, + _build_next_token_position_mask, + _shift_with_zero_padding, + compute_eagle3_loss, +) +from deepspec.modeling.eagle3.qwen3.modeling import Qwen3Eagle3Model +from deepspec.modeling.eagle3.qwen3.config import build_draft_config +from deepspec.utils.config import load_config, parse_opts_to_config + + +def _position_masks(loss_mask: torch.Tensor, attention_mask: torch.Tensor, k: int): + shifted = _build_next_token_position_mask( + loss_mask=loss_mask, + attention_mask=attention_mask, + ) + masks = [] + for _ in range(k): + masks.append(shifted) + shifted = _shift_with_zero_padding( + shifted.squeeze(-1), left=False + ).unsqueeze(-1) + return masks + + +def test_compact_plan_exact_indices(): + loss_mask = torch.tensor( + [ + [1, 0, 1, 0, 0, 1, 0], + [0, 1, 0, 1, 0, 0, 0], + ], + dtype=torch.long, + ) + attention_mask = torch.tensor( + [ + [1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 0, 0], + ], + dtype=torch.long, + ) + plan = _build_compact_plan( + position_masks=_position_masks(loss_mask, attention_mask, k=4) + ) + + expected_draft = ( + [0, 2, 5, 8, 10], + [1, 4, 7, 9], + [0, 3, 8], + [2, 7], + ) + expected_teacher = ( + [1, 3, 6, 9, 11], + [3, 6, 9, 11], + [3, 6, 11], + [6, 11], + ) + assert [step.num_valid_rows for step in plan.steps] == [5, 4, 3, 2] + for step, draft_indices, teacher_indices in zip( + plan.steps, expected_draft, expected_teacher, strict=True + ): + assert step.draft_indices.tolist() == draft_indices + assert step.teacher_indices.tolist() == teacher_indices + + +def test_trim_config_is_opt_in(): + target_config = _tiny_config() + target_config.num_hidden_layers = 5 + base_args = dict( + target_model_name_or_path="tiny-local-fixture", + target_layer_ids=[0, 1, 2, 3, 4], + ttt_length=3, + step_loss_decay=0.8, + draft_num_hidden_layers=1, + ) + default_config = build_draft_config( + target_config=target_config, + model_args=SimpleNamespace(**base_args), + ) + enabled_config = build_draft_config( + target_config=target_config, + model_args=SimpleNamespace(**base_args, trim_loss_positions=True), + ) + assert default_config.trim_loss_positions is False + assert enabled_config.trim_loss_positions is True + + +def test_public_config_accepts_trim_override(): + config = load_config("config/eagle3/eagle3_qwen3_8b.py") + assert config.model.trim_loss_positions is False + overridden = parse_opts_to_config( + ["model.trim_loss_positions=true"], + config, + ) + assert overridden.model.trim_loss_positions is True + + +def _tiny_config() -> Qwen3Config: + config = Qwen3Config( + vocab_size=97, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=64, + attention_dropout=0.0, + attention_bias=False, + ) + config.target_layer_ids = [0, 1, 2, 3, 4] + config.ttt_length = 3 + config.step_loss_decay = 0.8 + config.trim_loss_positions = False + config._attn_implementation = "flex_attention" + return config + + +def _tiny_model(device: torch.device) -> Qwen3Eagle3Model: + model = Qwen3Eagle3Model(_tiny_config()).to( + device=device, dtype=torch.bfloat16 + ) + model.embed_tokens.requires_grad_(False) + model.lm_head.requires_grad_(False) + model.train() + return model + + +def _tiny_batch(device: torch.device) -> dict[str, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(0) + batch_size, seq_len, hidden_size = 1, 16, 32 + loss_mask = torch.zeros( + batch_size, seq_len, dtype=torch.long, device=device + ) + loss_mask[0, [0, 2, 5, 7, 13]] = 1 + return { + "input_ids": torch.randint( + 0, + 97, + (batch_size, seq_len), + generator=generator, + device=device, + ), + "attention_mask": torch.ones( + batch_size, seq_len, dtype=torch.long, device=device + ), + "loss_mask": loss_mask, + "target_hidden_states": torch.randn( + batch_size, + seq_len, + 5 * hidden_size, + generator=generator, + device=device, + dtype=torch.bfloat16, + ), + "target_last_hidden_states": torch.randn( + batch_size, + seq_len, + hidden_size, + generator=generator, + device=device, + dtype=torch.bfloat16, + ), + } + + +def _capture_metrics(): + metrics = {} + + def capture(name, value, *, den=None, reduction="dp_sum", tag="train"): + del reduction + numerator = float(torch.as_tensor(value).detach().float().item()) + if den is None: + resolved = numerator + else: + denominator = float(torch.as_tensor(den).detach().float().item()) + resolved = 0.0 if denominator == 0.0 else numerator / denominator + metrics[f"{tag}/{name}"] = resolved + + return metrics, capture + + +def _run_tiny_arm( + *, + state_dict: dict[str, torch.Tensor], + batch: dict[str, torch.Tensor], + trim_loss_positions: bool, + wrong_compact_indices: bool = False, +): + device = batch["input_ids"].device + model = _tiny_model(device) + model.load_state_dict(state_dict) + metrics, capture_metric = _capture_metrics() + + original_build_plan = loss_module._build_compact_plan + + def build_wrong_plan(*args, **kwargs): + plan = original_build_plan(*args, **kwargs) + max_index = plan.batch_size * plan.seq_len + wrong_steps = tuple( + replace( + step, + draft_indices=(step.draft_indices + 1).remainder(max_index), + ) + for step in plan.steps + ) + return replace(plan, steps=wrong_steps) + + plan_patch = ( + mock.patch.object(loss_module, "_build_compact_plan", build_wrong_plan) + if wrong_compact_indices + else mock.patch.object( + loss_module, "_build_compact_plan", original_build_plan + ) + ) + with mock.patch.object(loss_module, "add_metric", capture_metric), plan_patch: + loss = compute_eagle3_loss( + model=model, + batch=batch, + ttt_length=3, + step_loss_decay=0.8, + trim_loss_positions=trim_loss_positions, + ) + loss.backward() + + gradients = { + name: parameter.grad.detach().float().cpu().clone() + for name, parameter in model.named_parameters() + if parameter.requires_grad and parameter.grad is not None + } + return float(loss.detach().float().item()), gradients, metrics + + +def _within_loss_tolerance(a: float, b: float) -> bool: + return abs(a - b) <= 5e-3 * max(abs(a), abs(b)) + 1e-4 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_trim_matches_full_loss_grad_and_metrics_cuda(): + torch.manual_seed(0) + device = torch.device("cuda") + template = _tiny_model(device) + # Make the row-shift negative control observably wrong under the same + # bf16 tolerance used for the equivalence arm. The production head is + # frozen too; scaling this tiny fixture changes sensitivity, not semantics. + with torch.no_grad(): + template.lm_head.weight.mul_(16.0) + state_dict = copy.deepcopy(template.state_dict()) + del template + batch = _tiny_batch(device) + + full_loss, full_gradients, full_metrics = _run_tiny_arm( + state_dict=state_dict, + batch=batch, + trim_loss_positions=False, + ) + trim_loss, trim_gradients, trim_metrics = _run_tiny_arm( + state_dict=state_dict, + batch=batch, + trim_loss_positions=True, + ) + + assert _within_loss_tolerance(full_loss, trim_loss) + assert full_gradients.keys() == trim_gradients.keys() + for name in full_gradients: + torch.testing.assert_close( + full_gradients[name], + trim_gradients[name], + rtol=5e-3, + atol=5e-3, + msg=lambda message, name=name: f"gradient {name}: {message}", + ) + assert full_metrics.keys() == trim_metrics.keys() + for name in full_metrics: + assert _within_loss_tolerance(full_metrics[name], trim_metrics[name]), ( + name, + full_metrics[name], + trim_metrics[name], + ) + + wrong_loss, wrong_gradients, _ = _run_tiny_arm( + state_dict=state_dict, + batch=batch, + trim_loss_positions=True, + wrong_compact_indices=True, + ) + negative_control_detected = not _within_loss_tolerance(full_loss, wrong_loss) + if not negative_control_detected: + for name in full_gradients: + try: + torch.testing.assert_close( + full_gradients[name], + wrong_gradients[name], + rtol=5e-3, + atol=5e-3, + ) + except AssertionError: + negative_control_detected = True + break + assert negative_control_detected + + +def _run_scope_trace(*, trim_loss_positions: bool): + device = torch.device("cuda") + torch.manual_seed(0) + model = _tiny_model(device) + batch = _tiny_batch(device) + q_lengths = [] + lm_head_lengths = [] + cache_lengths = [] + + q_hook = model.layers[0].self_attn.q_proj.register_forward_pre_hook( + lambda _module, args: q_lengths.append(int(args[0].shape[1])) + ) + lm_head_hook = model.lm_head.register_forward_pre_hook( + lambda _module, args: lm_head_lengths.append(int(args[0].shape[1])) + ) + + cache = DynamicCache() + original_update = cache.update + + def recording_update(key_states, value_states, layer_idx, cache_kwargs=None): + output = original_update( + key_states, value_states, layer_idx, cache_kwargs + ) + cache_lengths.append(int(cache.get_seq_length(layer_idx))) + return output + + cache.update = recording_update + try: + with mock.patch.object(loss_module, "DynamicCache", return_value=cache), mock.patch.object( + loss_module, "add_metric", lambda *args, **kwargs: None + ): + compute_eagle3_loss( + model=model, + batch=batch, + ttt_length=3, + step_loss_decay=0.8, + trim_loss_positions=trim_loss_positions, + ) + finally: + q_hook.remove() + lm_head_hook.remove() + return q_lengths, lm_head_lengths, cache_lengths + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_trim_scope_keeps_backbone_and_teacher_projection_full_cuda(): + q_lengths, lm_head_lengths, cache_lengths = _run_scope_trace( + trim_loss_positions=True + ) + assert q_lengths == [16, 16, 16] + assert lm_head_lengths == [16, 5, 4, 4] + assert cache_lengths == [16, 32, 48] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_trim_default_false_keeps_full_logits(): + q_lengths, lm_head_lengths, cache_lengths = _run_scope_trace( + trim_loss_positions=False + ) + assert q_lengths == [16, 16, 16] + assert lm_head_lengths == [16, 16, 16, 16] + assert cache_lengths == [16, 32, 48] From 52995d24e24ae060a962593b8685513d2327c3b6 Mon Sep 17 00:00:00 2001 From: julyanghar Date: Sat, 29 Aug 2026 22:00:03 -0500 Subject: [PATCH 2/2] Avoid materializing indexed teacher rows --- deepspec/modeling/eagle3/loss.py | 232 ++++++++++++++++-- .../benchmark_eagle3_trim_loss_positions.py | 9 +- tests/test_eagle3_trim_loss_positions.py | 141 +++++++++++ 3 files changed, 353 insertions(+), 29 deletions(-) diff --git a/deepspec/modeling/eagle3/loss.py b/deepspec/modeling/eagle3/loss.py index dd126a47..9c4f34ee 100644 --- a/deepspec/modeling/eagle3/loss.py +++ b/deepspec/modeling/eagle3/loss.py @@ -281,6 +281,136 @@ def scatter(values: torch.Tensor) -> torch.Tensor: return correct_mask, accept_rate_mask, valid_mask +@triton.jit +def _indexed_accept_rate_kernel( + logits_ptr, + logits_stride, + target_table_ptr, + target_stride, + target_indices_ptr, + position_mask_ptr, + accept_rate_ptr, + n_cols, + BLOCK_SIZE: tl.constexpr, +): + program_id = tl.program_id(0).to(tl.int64) + logits_ptr += program_id * logits_stride + target_row = tl.load(target_indices_ptr + program_id).to(tl.int64) + target_table_ptr += target_row * target_stride + position_mask = tl.load(position_mask_ptr + program_id) + if position_mask == 0: + tl.store(accept_rate_ptr + program_id, 0.0) + return + + m = float("-inf") + d = 0.0 + for i in range(0, n_cols, BLOCK_SIZE): + offsets = i + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + logits_block = tl.load( + logits_ptr + offsets, mask=mask, other=float("-inf") + ).cast(tl.float32) + block_max = tl.max(tl.where(mask, logits_block, float("-inf"))) + m_new = tl.maximum(m, block_max) + d = d * tl.exp(m - m_new) + tl.sum( + tl.where(mask, tl.exp(logits_block - m_new), 0.0) + ) + m = m_new + + l1_distance = 0.0 + for i in range(0, n_cols, BLOCK_SIZE): + offsets = i + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + logits_block = tl.load(logits_ptr + offsets, mask=mask, other=0.0).cast( + tl.float32 + ) + target_block = tl.load( + target_table_ptr + offsets, mask=mask, other=0.0 + ).cast(tl.float32) + draft_prob = tl.exp(logits_block - m) / d + l1_distance += tl.sum( + tl.where(mask, tl.abs(draft_prob - target_block), 0.0) + ) + + accept_rate = tl.maximum(0.0, tl.minimum(1.0, 1.0 - 0.5 * l1_distance)) + tl.store(accept_rate_ptr + program_id, accept_rate.to(tl.float32)) + + +@torch.no_grad() +def _log_eagle3_indexed_step_metrics( + *, + step_idx: int, + draft_logits: torch.Tensor, + target_table: torch.Tensor, + target_token_ids_table: torch.Tensor, + target_indices: torch.Tensor, + position_mask: torch.Tensor, + full_shape: tuple[int, int], + scatter_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute compact metrics by reading shared teacher rows through indices.""" + + assert draft_logits.is_cuda + assert draft_logits.shape[0] == 1 + _, num_rows, vocab_size = draft_logits.shape + assert target_table.ndim == 2 and target_table.shape[1] == vocab_size + assert target_indices.numel() == num_rows + valid_mask = position_mask.squeeze(-1) > 0 + target_token_ids = target_token_ids_table.index_select(0, target_indices).view( + 1, num_rows + ) + correct_mask = (draft_logits.argmax(-1) == target_token_ids) & valid_mask + + accept_rate_mask = torch.empty( + (num_rows,), device=draft_logits.device, dtype=torch.float32 + ) + logits_flat = draft_logits.contiguous().view(num_rows, vocab_size) + target_table = target_table.contiguous() + target_indices = target_indices.contiguous() + position_mask_flat = position_mask.contiguous().view(num_rows) + BLOCK_SIZE, num_warps = _calculate_settings(vocab_size) + _indexed_accept_rate_kernel[(num_rows,)]( + logits_flat, + logits_flat.stride(0), + target_table, + target_table.stride(0), + target_indices, + position_mask_flat, + accept_rate_mask, + vocab_size, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + accept_rate_mask = accept_rate_mask.view(1, num_rows) + + correct = correct_mask.to(torch.float32).sum() + valid_count = valid_mask.to(torch.float32).sum() + accept_sum = accept_rate_mask.sum() + add_metric(f"accuracy@{step_idx}", correct, den=valid_count, tag="train") + add_metric(f"accept_rate@{step_idx}", accept_sum, den=valid_count, tag="train") + add_metric( + f"valid_tokens@{step_idx}", + valid_count, + reduction="dp_sum", + tag="train", + ) + + num_valid_rows = int(scatter_indices.numel()) + full_numel = int(full_shape[0]) * int(full_shape[1]) + + def scatter(values: torch.Tensor) -> torch.Tensor: + output = values.new_zeros(full_numel) + if num_valid_rows > 0: + output.index_copy_( + 0, + scatter_indices, + values.reshape(-1)[:num_valid_rows], + ) + return output.reshape(full_shape) + + return scatter(correct_mask), scatter(accept_rate_mask), scatter(valid_mask) + + @torch.no_grad() def _log_eagle3_prefix_metrics( *, @@ -311,6 +441,7 @@ def _log_softmax_forward_kernel( logits_stride, target_ptr, target_stride, + target_indices_ptr, position_mask_ptr, position_mask_stride, loss_ptr, @@ -318,11 +449,15 @@ def _log_softmax_forward_kernel( m_ptr, d_ptr, n_cols, + HAS_TARGET_INDICES: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): program_id = tl.program_id(0).to(tl.int64) logits_ptr += program_id * logits_stride - target_ptr += program_id * target_stride + target_row = program_id + if HAS_TARGET_INDICES: + target_row = tl.load(target_indices_ptr + program_id).to(tl.int64) + target_ptr += target_row * target_stride position_mask_ptr += program_id * position_mask_stride position_mask = tl.load(position_mask_ptr) if position_mask == 0: @@ -374,17 +509,22 @@ def _log_softmax_backward_kernel( logits_stride, target_ptr, target_stride, + target_indices_ptr, position_mask_ptr, grad_output_ptr, scaling_factor, m_ptr, d_ptr, n_cols, + HAS_TARGET_INDICES: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): program_id = tl.program_id(0).to(tl.int64) logits_ptr += program_id * logits_stride - target_ptr += program_id * target_stride + target_row = program_id + if HAS_TARGET_INDICES: + target_row = tl.load(target_indices_ptr + program_id).to(tl.int64) + target_ptr += target_row * target_stride position_mask_ptr += program_id position_mask = tl.load(position_mask_ptr) @@ -426,14 +566,29 @@ def _log_softmax_backward_kernel( tl.store(logits_ptr + offsets, grad_block.to(tl.float32), mask=mask) -def _fused_log_softmax_forward(logits, target_p, position_mask, normalizer): +def _fused_log_softmax_forward( + logits, + target_p, + position_mask, + normalizer, + target_indices=None, +): assert logits.is_cuda, "FusedLogSoftmaxLoss requires CUDA tensors." - assert logits.shape == target_p.shape assert position_mask.shape[:2] == logits.shape[:2] B, T, V = logits.shape + if target_indices is None: + assert logits.shape == target_p.shape + target_indices_ptr = target_p + has_target_indices = False + else: + assert B == 1 + assert target_p.ndim == 2 and target_p.shape[1] == V + assert target_indices.numel() == B * T + target_indices_ptr = target_indices.contiguous() + has_target_indices = True loss = torch.zeros((B * T, 1), device=logits.device, dtype=torch.float32) logits_flat = logits.contiguous().view(B * T, V) - target_flat = target_p.contiguous().view(B * T, V) + target_flat = target_p.contiguous().view(-1, V) position_mask_flat = position_mask.contiguous().view(B * T, 1).bool() grid = (B * T,) m = torch.zeros((B * T,), device=logits.device, dtype=torch.float32) @@ -444,6 +599,7 @@ def _fused_log_softmax_forward(logits, target_p, position_mask, normalizer): logits_flat.stride(0), target_flat, target_flat.stride(0), + target_indices_ptr, position_mask_flat, position_mask_flat.stride(0), loss, @@ -451,6 +607,7 @@ def _fused_log_softmax_forward(logits, target_p, position_mask, normalizer): m, d, V, + HAS_TARGET_INDICES=has_target_indices, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, ) @@ -466,11 +623,22 @@ def _fused_log_softmax_backward( d, normalizer, grad_output, + target_indices=None, ): B, T, V = logits.shape + if target_indices is None: + assert logits.shape == target_p.shape + target_indices_ptr = target_p + has_target_indices = False + else: + assert B == 1 + assert target_p.ndim == 2 and target_p.shape[1] == V + assert target_indices.numel() == B * T + target_indices_ptr = target_indices.contiguous() + has_target_indices = True scaling_factor = 1.0 / float(normalizer) logits_flat = logits.contiguous().view(B * T, V) - target_flat = target_p.contiguous().view(B * T, V) + target_flat = target_p.contiguous().view(-1, V) position_mask_flat = position_mask.contiguous().view(B * T, 1).bool() grid = (B * T,) BLOCK_SIZE, num_warps = _calculate_settings(V) @@ -479,12 +647,14 @@ def _fused_log_softmax_backward( logits_flat.stride(0), target_flat, target_flat.stride(0), + target_indices_ptr, position_mask_flat, grad_output, scaling_factor, m, d, V, + HAS_TARGET_INDICES=has_target_indices, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, ) @@ -531,11 +701,11 @@ def backward(ctx, grad_output): class FusedIndexedLogSoftmaxLoss(torch.autograd.Function): - """Fused loss whose target rows are views of one shared compact table. + """Fused loss that reads target rows directly from one shared compact table. - Forward and backward materialize only the current step's indexed rows. The - autograd context saves the shared table plus row indices, rather than one - ``[N_step, V]`` target copy for every TTT step. + Forward and backward pass row indices to the Triton kernels, so neither + phase materializes ``[N_step, V]`` target copies. The autograd context saves + only the shared table plus row indices. """ @staticmethod @@ -548,12 +718,12 @@ def forward( normalizer, ): assert logits.shape[0] == 1 - target_p = target_table.index_select(0, target_indices).unsqueeze(0) loss, m, d = _fused_log_softmax_forward( logits, - target_p, + target_table, position_mask, normalizer, + target_indices=target_indices, ) ctx.save_for_backward( logits.detach(), @@ -576,15 +746,15 @@ def backward(ctx, grad_output): m, d, ) = ctx.saved_tensors - target_p = target_table.index_select(0, target_indices).unsqueeze(0) grad_logits = _fused_log_softmax_backward( logits=logits, - target_p=target_p, + target_p=target_table, position_mask=position_mask, m=m, d=d, normalizer=ctx.normalizer, grad_output=grad_output, + target_indices=target_indices, ) return grad_logits, None, None, None, None @@ -643,12 +813,14 @@ def compute_eagle3_loss( ttt_length=int(ttt_length), ) compact_teacher_probs = None + compact_teacher_token_ids = None else: target_probs = None compact_teacher_probs = _build_compact_teacher_probs( target_logits=target_logits, compact_plan=compact_plan, ) + compact_teacher_token_ids = compact_teacher_probs.argmax(dim=-1) del target_logits correct_masks = [] @@ -668,9 +840,7 @@ def compute_eagle3_loss( metric_full_shape = None else: compact_step = compact_plan.steps[step_idx] - target_step_probs = compact_teacher_probs.index_select( - 0, compact_step.teacher_table_indices - ).unsqueeze(0) + target_step_probs = None position_mask_step = compact_step.position_mask logit_indices = compact_step.draft_indices metric_scatter_indices = compact_step.scatter_indices @@ -687,14 +857,26 @@ def compute_eagle3_loss( logit_indices=logit_indices, ) hidden_states = output.hidden_states - correct_mask, accept_rate_mask, valid_mask = _log_eagle3_step_metrics( - step_idx=step_idx, - draft_logits=output.draft_logits, - target_probs=target_step_probs, - position_mask=position_mask_step, - full_shape=metric_full_shape, - scatter_indices=metric_scatter_indices, - ) + if compact_plan is None: + correct_mask, accept_rate_mask, valid_mask = _log_eagle3_step_metrics( + step_idx=step_idx, + draft_logits=output.draft_logits, + target_probs=target_step_probs, + position_mask=position_mask_step, + ) + else: + correct_mask, accept_rate_mask, valid_mask = ( + _log_eagle3_indexed_step_metrics( + step_idx=step_idx, + draft_logits=output.draft_logits, + target_table=compact_teacher_probs, + target_token_ids_table=compact_teacher_token_ids, + target_indices=compact_step.teacher_table_indices, + position_mask=position_mask_step, + full_shape=metric_full_shape, + scatter_indices=metric_scatter_indices, + ) + ) correct_masks.append(correct_mask) accept_rate_masks.append(accept_rate_mask) valid_masks.append(valid_mask) diff --git a/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py b/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py index b0cbfd19..975e1fe4 100644 --- a/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py +++ b/scripts/benchmarks/benchmark_eagle3_trim_loss_positions.py @@ -196,11 +196,12 @@ def validate_results(results): d30 = by_density[0.3]["summary"] d60 = by_density[0.6]["summary"] d90 = by_density[0.9]["summary"] - assert d30["paired_median_speedup"] >= 1.05 - assert d30["median_peak_reduction_bytes"] >= GIB - assert d60["paired_median_speedup"] >= 1.02 - assert d60["median_peak_reduction_bytes"] > 0 + assert d30["paired_median_speedup"] >= 1.50 + assert d30["median_peak_reduction_bytes"] >= 11.0 * GIB + assert d60["paired_median_speedup"] >= 1.15 + assert d60["median_peak_reduction_bytes"] >= 5.0 * GIB assert d90["paired_median_speedup"] >= 0.98 + assert d90["median_peak_reduction_bytes"] >= 0 assert ( d30["paired_median_speedup"] >= d60["paired_median_speedup"] diff --git a/tests/test_eagle3_trim_loss_positions.py b/tests/test_eagle3_trim_loss_positions.py index 2b6d9e75..7b3bc865 100644 --- a/tests/test_eagle3_trim_loss_positions.py +++ b/tests/test_eagle3_trim_loss_positions.py @@ -7,6 +7,7 @@ import pytest import torch +from torch.utils._python_dispatch import TorchDispatchMode from transformers import DynamicCache from transformers.models.qwen3.configuration_qwen3 import Qwen3Config @@ -22,6 +23,25 @@ from deepspec.utils.config import load_config, parse_opts_to_config +class _TeacherIndexSelectCounter(TorchDispatchMode): + def __init__(self, teacher_table: torch.Tensor): + super().__init__() + self.teacher_storage_ptr = teacher_table.untyped_storage().data_ptr() + self.count = 0 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + del types + kwargs = kwargs or {} + if func == torch.ops.aten.index_select.default: + source = args[0] + if ( + source.ndim == 2 + and source.untyped_storage().data_ptr() == self.teacher_storage_ptr + ): + self.count += 1 + return func(*args, **kwargs) + + def _position_masks(loss_mask: torch.Tensor, attention_mask: torch.Tensor, k: int): shifted = _build_next_token_position_mask( loss_mask=loss_mask, @@ -247,6 +267,127 @@ def _within_loss_tolerance(a: float, b: float) -> bool: return abs(a - b) <= 5e-3 * max(abs(a), abs(b)) + 1e-4 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_indexed_metrics_and_loss_match_materialized_reference_cuda(): + torch.manual_seed(0) + device = torch.device("cuda") + unique_rows, num_rows, vocab_size = 7, 5, 97 + target_table = torch.softmax( + torch.randn(unique_rows, vocab_size, device=device), dim=-1 + ) + target_token_ids_table = target_table.argmax(dim=-1) + target_indices = torch.tensor([0, 3, 3, 6, 1], device=device) + wrong_target_indices = (target_indices + 1).remainder(unique_rows) + position_mask = torch.ones(1, num_rows, 1, device=device) + scatter_indices = torch.arange(num_rows, device=device) + base_logits = torch.randn( + 1, num_rows, vocab_size, device=device, dtype=torch.bfloat16 + ) + + reference_logits = base_logits.detach().clone().requires_grad_(True) + reference_counter = _TeacherIndexSelectCounter(target_table) + with reference_counter, mock.patch.object( + loss_module, "add_metric", lambda *args, **kwargs: None + ): + materialized_target = target_table.index_select( + 0, target_indices + ).unsqueeze(0) + reference_metrics = loss_module._log_eagle3_step_metrics( + step_idx=0, + draft_logits=reference_logits, + target_probs=materialized_target, + position_mask=position_mask, + ) + reference_loss = loss_module.FusedLogSoftmaxLoss.apply( + reference_logits, + materialized_target, + position_mask, + float(num_rows), + ) + reference_loss.backward() + reference_gradient = reference_logits.grad.detach().float().clone() + assert reference_counter.count >= 1 + + indexed_logits = base_logits.detach().clone().requires_grad_(True) + indexed_counter = _TeacherIndexSelectCounter(target_table) + with indexed_counter, mock.patch.object( + loss_module, "add_metric", lambda *args, **kwargs: None + ): + indexed_metrics = loss_module._log_eagle3_indexed_step_metrics( + step_idx=0, + draft_logits=indexed_logits, + target_table=target_table, + target_token_ids_table=target_token_ids_table, + target_indices=target_indices, + position_mask=position_mask, + full_shape=(1, num_rows), + scatter_indices=scatter_indices, + ) + indexed_loss = loss_module.FusedIndexedLogSoftmaxLoss.apply( + indexed_logits, + target_table, + target_indices, + position_mask, + float(num_rows), + ) + indexed_loss.backward() + indexed_gradient = indexed_logits.grad.detach().float().clone() + assert indexed_counter.count == 0 + + torch.testing.assert_close(indexed_metrics[0], reference_metrics[0]) + torch.testing.assert_close( + indexed_metrics[1], reference_metrics[1], rtol=5e-3, atol=5e-3 + ) + torch.testing.assert_close(indexed_metrics[2], reference_metrics[2]) + torch.testing.assert_close( + indexed_loss.detach().float(), + reference_loss.detach().float(), + rtol=5e-3, + atol=5e-3, + ) + torch.testing.assert_close( + indexed_gradient, + reference_gradient, + rtol=5e-3, + atol=5e-3, + ) + + wrong_logits = base_logits.detach().clone().requires_grad_(True) + with mock.patch.object( + loss_module, "add_metric", lambda *args, **kwargs: None + ): + wrong_metrics = loss_module._log_eagle3_indexed_step_metrics( + step_idx=0, + draft_logits=wrong_logits, + target_table=target_table, + target_token_ids_table=target_token_ids_table, + target_indices=wrong_target_indices, + position_mask=position_mask, + full_shape=(1, num_rows), + scatter_indices=scatter_indices, + ) + wrong_loss = loss_module.FusedIndexedLogSoftmaxLoss.apply( + wrong_logits, + target_table, + wrong_target_indices, + position_mask, + float(num_rows), + ) + wrong_loss.backward() + negative_control_detected = False + for reference, wrong in ( + (reference_metrics[1], wrong_metrics[1]), + (reference_loss.detach().float(), wrong_loss.detach().float()), + (reference_gradient, wrong_logits.grad.detach().float()), + ): + try: + torch.testing.assert_close(reference, wrong, rtol=5e-3, atol=5e-3) + except AssertionError: + negative_control_detected = True + break + assert negative_control_detected + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_trim_matches_full_loss_grad_and_metrics_cuda(): torch.manual_seed(0)