Skip to content

Chunked DFlash loss path flattens prev_token_ids, breaking DSpark subclass heads #3

Description

@xy200303

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions