ITADN

[Bug] Fix incorrect padding logic in both backends and wrong target_hidden construction in DFlash training

#471Openhukongyi 创建于 2026-02-13
H
hukongyicommented
### Checklist - [x] 1. I have searched related issues but cannot get the expected help. - [x] 2. The bug has not been fixed in the latest version. - [x] 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback. - [x] 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed. - [x] 5. Please use English, otherwise it will be closed. ### Describe the bug **Description:** I have identified critical issues in the data generation and training logic for DFlash. Specifically, the padding handling in both SGLang and HF backends is incorrect, and the construction of `target_hidden` in the training loop is fundamentally flawed. ### 1. Incorrect Padding Logic in Target Models **File:** `specforge/modeling/target/dflash_target_model.py` **Issue:** There is an inconsistency and error in how padding is handled across different backends: * **SGLang Backend (`SGLangDFlashTargetModel`):** It currently pads both `hidden_states` and `input_ids`. Padding `hidden_states` is incorrect and leads to dimension issues downstream. * **Hugging Face Backend (`HFDFlashTargetModel`):** It currently performs **no padding** at all. This is also incorrect because `input_ids` must be padded to handle batched training data properly. **Proposed Fix:** Both backends should follow the same correct logic: 1. **Pad `input_ids`**: This is necessary predict next tokens. 2. **Do NOT pad `hidden_states`**. ### 2. Wrong `block_hidden` Construction in `_build_blocks_from_anchors` **File:** `specforge/core/dflash.py` **Issue:** In `_build_blocks_from_anchors`, the code currently gathers `hidden_states` based on the block indices (`gather_idx`) to create `block_hidden`. **Problematic Code:** ```python # specforge/core/dflash.py def _build_blocks_from_anchors(...): # ... # This logic gathers hidden states of the tokens INSIDE the block block_hidden = torch.gather( hidden_states, 1, gather_idx.unsqueeze(-1).expand(-1, -1, hidden_states.size(-1)), ) # ... ``` **Why this is wrong:** This `block_hidden` is subsequently passed as `target_hidden` to the drafter model. The drafter uses this `target_hidden` as the **conditioning context** (via `self.fc(target_hidden)`). By gathering the hidden states of the tokens within the block (which are the targets to be predicted), we are destroying the necessary prefix/context information. The `target_hidden` passed to the drafter should **not undergo block-wise gathering**. It should represent the full context features from the target model to correctly condition the generation of the draft block. ### Status I am actively working on fixing these padding inconsistencies and the hidden state construction logic. A Pull Request will be submitted shortly. ### Reproduction same as https://github.com/sgl-project/SpecForge/issues/465 ### Environment same as https://github.com/sgl-project/SpecForge/issues/465
0 条评论