Chunked DFlash loss path flattens prev_token_ids, breaking DSpark subclass heads
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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