Tencent / Tencent/AngelSpec

Chunked DFlash loss path flattens prev_token_ids, breaking DSpark subclass heads

Open
#3 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
270
Forks
28
PR merge metrics
No merged PRs in 30d

Description

In _forward_chunked_decay, prev_token_ids is flattened to 2-D before being passed to the _compute_draft_logits hook:

https://github.com/Tencent/AngelSpec/blob/main/angelspec/models/dflash.py#L1089-L1094

prev_chunk = prev_token_ids[:, start:end].reshape(bsz, nb * bs)   # 2-D
logits_chunk = self._compute_draft_logits(dh_chunk, lm_head_weight, prev_chunk, nb)

But DSparkModel._compute_draft_logits (angelspec/models/dspark.py) forwards prev_token_ids to markov_head.apply_block_logits(token_ids=...), which expects [B, n_blocks, block_size] (angelspec/models/draft/dspark.py). With a 2-D input, compute_step_bias produces a [bsz, nb*bs, V] bias that is added to [bsz, nb, bs, V] logits — a RuntimeError whenever nb > 1 (and a silently wrong broadcast followed by a reshape crash when nb == 1, bsz > 1).

The gate that is supposed to keep subclass heads off the chunked path only checks _extra_distill_needed():

https://github.com/Tencent/AngelSpec/blob/main/angelspec/models/dflash.py#L754-L761

For DSpark that returns True only when a confidence head is active, so a DSpark model with only a markov_head (or confidence_head_alpha=0) plus ANGELSPEC_DFLASH_LOSS_CHUNK > 0 and the decay objective takes the chunked path and crashes. This contradicts the stated intent in _dflash_loss_chunk's docstring ("dpace / distillation / subclass heads keep the single full-vocab projection").

A secondary issue on the same path: _forward_chunked_decay never calls _compute_extra_loss, so the position-adaptive-alpha smoothness regularizer (pos_alpha.smooth_loss()) is silently dropped for any DSpark config that reaches the chunked path without crashing.

Possible fixes: extend the chunk gate to exclude models that override _compute_draft_logits / _compute_extra_loss, or pass prev_token_ids as [B, nb, bs] and make the hooks shape-consistent.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading _forward_chunked_decay and the chunk gate in angelspec/models/dflash.py, then compare the hook implementations in angelspec/models/dspark.py and angelspec/models/draft/dspark.py. Reproduce the decay objective with ANGELSPEC_DFLASH_LOSS_CHUNK > 0 using a markov-only DSpark configuration and with confidence_head_alpha=0. Done means the chunked path no longer crashes and the pos_alpha.smooth_loss() regularizer is retained when applicable.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.