[Bug] Fix incorrect padding logic in both backends and wrong target_hidden construction in DFlash training
### 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 条评论