From d2caaebd396254607c4f7bc2cd921cdf050a53e3 Mon Sep 17 00:00:00 2001 From: chenyue1122 Date: Thu, 2 Jul 2026 12:17:53 +0800 Subject: [PATCH 1/3] feat: add chunked gather for anchor hidden state alignment Process anchors in batches of 64 to reduce memory footprint. --- deepspec/modeling/dspark/gemma4/modeling.py | 32 +++++++++++---------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/deepspec/modeling/dspark/gemma4/modeling.py b/deepspec/modeling/dspark/gemma4/modeling.py index c8d70c3..683d372 100644 --- a/deepspec/modeling/dspark/gemma4/modeling.py +++ b/deepspec/modeling/dspark/gemma4/modeling.py @@ -519,21 +519,23 @@ def forward( aligned_target_logits = None if target_last_hidden_states is not None: target_pred_indices = (safe_label_indices - 1).clamp(min=0) - aligned_target_hidden = torch.gather( - target_last_hidden_states.unsqueeze(1).expand( - -1, - anchor_positions.size(1), - -1, - -1, - ), - 2, - target_pred_indices.unsqueeze(-1).expand( - -1, - -1, - -1, - target_last_hidden_states.size(-1), - ), - ) + hidden_chunks = [] + anchor_chunk_size = 64 + num_anchors = anchor_position.size(1) + for c_start in range(0, num_anchors, anchor_chunk_size): + c_end = min(c_start + anchor_chunk_size, num_anchors) + chunk_indices = target_pred_indices[:, c_start:c_end] + chunk_gather = torch.gather( + target_last_hidden_states.unsqueeze(1).expand( + -1, c_end - c_start, -1, -1, + ), + 2, + chunk_indices.unsqueeze(-1).expand( + -1, -1, -1, target_last_hidden_states.size(-1), + ), + ) + hidden_chunks.append(chunk_gather) + aligned_target_hidden = torch.cat(hidden_chunks, dim=1) aligned_target_logits = self.compute_logits(aligned_target_hidden) eval_mask = build_eval_mask( seq_len=seq_len, From 6a6e7ba78ca4c49493d740f7e371371dfbf4db54 Mon Sep 17 00:00:00 2001 From: chenyue1122 Date: Thu, 2 Jul 2026 15:59:26 +0800 Subject: [PATCH 2/3] feat: add chunked gather for anchor hidden state alignment Process anchors in batches of 64 to reduce memory footprint. --- deepspec/modeling/dspark/gemma4/modeling.py | 2 +- deepspec/modeling/dspark/qwen3/modeling.py | 32 +++++++++++---------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/deepspec/modeling/dspark/gemma4/modeling.py b/deepspec/modeling/dspark/gemma4/modeling.py index 683d372..4682fbf 100644 --- a/deepspec/modeling/dspark/gemma4/modeling.py +++ b/deepspec/modeling/dspark/gemma4/modeling.py @@ -521,7 +521,7 @@ def forward( target_pred_indices = (safe_label_indices - 1).clamp(min=0) hidden_chunks = [] anchor_chunk_size = 64 - num_anchors = anchor_position.size(1) + num_anchors = anchor_positions.size(1) for c_start in range(0, num_anchors, anchor_chunk_size): c_end = min(c_start + anchor_chunk_size, num_anchors) chunk_indices = target_pred_indices[:, c_start:c_end] diff --git a/deepspec/modeling/dspark/qwen3/modeling.py b/deepspec/modeling/dspark/qwen3/modeling.py index 24d4441..c435ff2 100644 --- a/deepspec/modeling/dspark/qwen3/modeling.py +++ b/deepspec/modeling/dspark/qwen3/modeling.py @@ -447,21 +447,23 @@ def forward( aligned_target_logits = None if target_last_hidden_states is not None: target_pred_indices = (safe_label_indices - 1).clamp(min=0) - aligned_target_hidden = torch.gather( - target_last_hidden_states.unsqueeze(1).expand( - -1, - anchor_positions.size(1), - -1, - -1, - ), - 2, - target_pred_indices.unsqueeze(-1).expand( - -1, - -1, - -1, - target_last_hidden_states.size(-1), - ), - ) + hidden_chunks = [] + anchor_chunk_size = 64 + num_anchors = anchor_positions.size(1) + for c_start in range(0, num_anchors, anchor_chunk_size): + c_end = min(c_start + anchor_chunk_size, num_anchors) + chunk_indices = target_pred_indices[:, c_start:c_end] + chunk_gather = torch.gather( + target_last_hidden_states.unsqueeze(1).expand( + -1, c_end - c_start, -1, -1, + ), + 2, + chunk_indices.unsqueeze(-1).expand( + -1, -1, -1, target_last_hidden_states.size(-1), + ), + ) + hidden_chunks.append(chunk_gather) + aligned_target_hidden = torch.cat(hidden_chunks, dim=1) aligned_target_logits = self.compute_logits(aligned_target_hidden) eval_mask = build_eval_mask( seq_len=seq_len, From 461fbac18b717a51fb3db4d744e50b796b1c6d96 Mon Sep 17 00:00:00 2001 From: chenyue1122 Date: Thu, 2 Jul 2026 19:25:09 +0800 Subject: [PATCH 3/3] feat: add chunked gather for anchor hidden state alignment Process anchors in batches of 64 to reduce memory footprint. --- deepspec/modeling/dspark/common.py | 31 +++++++++++++++++++++ deepspec/modeling/dspark/gemma4/modeling.py | 22 ++++----------- deepspec/modeling/dspark/qwen3/modeling.py | 22 ++++----------- 3 files changed, 41 insertions(+), 34 deletions(-) diff --git a/deepspec/modeling/dspark/common.py b/deepspec/modeling/dspark/common.py index 76d908a..c70d54c 100644 --- a/deepspec/modeling/dspark/common.py +++ b/deepspec/modeling/dspark/common.py @@ -261,6 +261,36 @@ def create_position_ids( ) +@torch.compiler.disable(recursive=False) +def chunked_gather_target_hidden( + target_last_hidden_states: torch.Tensor, + target_pred_indices: torch.Tensor, + *, + anchor_chunk_size: int = 64, +) -> torch.Tensor: + """Gather target hidden states along the anchor axis in chunks. + + Splits the expand+gather across num_anchors to bound the peak activation + memory of the backward pass. Kept out of torch.compile because dynamic + range() loops trigger graph breaks or recompiles under dynamic=True. + """ + hidden_chunks = [] + num_anchors = target_pred_indices.size(1) + hidden_dim = target_last_hidden_states.size(-1) + for c_start in range(0, num_anchors, anchor_chunk_size): + c_end = min(c_start + anchor_chunk_size, num_anchors) + chunk_indices = target_pred_indices[:, c_start:c_end] + chunk_gather = torch.gather( + target_last_hidden_states.unsqueeze(1).expand( + -1, c_end - c_start, -1, -1, + ), + 2, + chunk_indices.unsqueeze(-1).expand(-1, -1, -1, hidden_dim), + ) + hidden_chunks.append(chunk_gather) + return torch.cat(hidden_chunks, dim=1) + + def create_noise_embed( embed_tokens: nn.Module, input_ids: torch.Tensor, @@ -306,4 +336,5 @@ def create_noise_embed( "log_sampler_stats", "create_position_ids", "create_noise_embed", + "chunked_gather_target_hidden", ] diff --git a/deepspec/modeling/dspark/gemma4/modeling.py b/deepspec/modeling/dspark/gemma4/modeling.py index 4682fbf..3ab93b3 100644 --- a/deepspec/modeling/dspark/gemma4/modeling.py +++ b/deepspec/modeling/dspark/gemma4/modeling.py @@ -21,6 +21,7 @@ AcceptRatePredictor, DSparkForwardOutput, build_eval_mask, + chunked_gather_target_hidden, create_dspark_attention_mask, create_noise_embed, create_position_ids, @@ -519,23 +520,10 @@ def forward( aligned_target_logits = None if target_last_hidden_states is not None: target_pred_indices = (safe_label_indices - 1).clamp(min=0) - hidden_chunks = [] - anchor_chunk_size = 64 - num_anchors = anchor_positions.size(1) - for c_start in range(0, num_anchors, anchor_chunk_size): - c_end = min(c_start + anchor_chunk_size, num_anchors) - chunk_indices = target_pred_indices[:, c_start:c_end] - chunk_gather = torch.gather( - target_last_hidden_states.unsqueeze(1).expand( - -1, c_end - c_start, -1, -1, - ), - 2, - chunk_indices.unsqueeze(-1).expand( - -1, -1, -1, target_last_hidden_states.size(-1), - ), - ) - hidden_chunks.append(chunk_gather) - aligned_target_hidden = torch.cat(hidden_chunks, dim=1) + aligned_target_hidden = chunked_gather_target_hidden( + target_last_hidden_states, + target_pred_indices, + ) aligned_target_logits = self.compute_logits(aligned_target_hidden) eval_mask = build_eval_mask( seq_len=seq_len, diff --git a/deepspec/modeling/dspark/qwen3/modeling.py b/deepspec/modeling/dspark/qwen3/modeling.py index c435ff2..c681748 100644 --- a/deepspec/modeling/dspark/qwen3/modeling.py +++ b/deepspec/modeling/dspark/qwen3/modeling.py @@ -21,6 +21,7 @@ AcceptRatePredictor, DSparkForwardOutput, build_eval_mask, + chunked_gather_target_hidden, create_dspark_attention_mask, create_noise_embed, create_position_ids, @@ -447,23 +448,10 @@ def forward( aligned_target_logits = None if target_last_hidden_states is not None: target_pred_indices = (safe_label_indices - 1).clamp(min=0) - hidden_chunks = [] - anchor_chunk_size = 64 - num_anchors = anchor_positions.size(1) - for c_start in range(0, num_anchors, anchor_chunk_size): - c_end = min(c_start + anchor_chunk_size, num_anchors) - chunk_indices = target_pred_indices[:, c_start:c_end] - chunk_gather = torch.gather( - target_last_hidden_states.unsqueeze(1).expand( - -1, c_end - c_start, -1, -1, - ), - 2, - chunk_indices.unsqueeze(-1).expand( - -1, -1, -1, target_last_hidden_states.size(-1), - ), - ) - hidden_chunks.append(chunk_gather) - aligned_target_hidden = torch.cat(hidden_chunks, dim=1) + aligned_target_hidden = chunked_gather_target_hidden( + target_last_hidden_states, + target_pred_indices, + ) aligned_target_logits = self.compute_logits(aligned_target_hidden) eval_mask = build_eval_mask( seq_len=seq_len,