Merge pull request sgl-project#731 from modal-projects/dcw02/dflash-correctness

Correctness fixes for DFlash training path

Author: jiapingW

docs/basic_usage/training.md

@@ -154,6 +154,31 @@ model:
   draft_block_size: 8         # DFlash only
 ```
 
+For DFlash, configure the attention layout in the referenced draft JSON. Each
+entry corresponds to one draft layer; sliding layers share one positive window:
+
+```json
+{
+  "num_hidden_layers": 5,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "use_sliding_window": true,
+  "sliding_window": 2048
+}
+```
+
+Use `"full_attention"` for every entry, `"use_sliding_window": false`, and
+`"sliding_window": null` for a full-only draft. The layout length must equal
+`num_hidden_layers`. A layer-count override may resize a uniform layout, but a
+mixed layout must be edited explicitly in the draft JSON.
+
+The `eager`, `sdpa`, and `flex_attention` backends support both layouts.
+
 Domino and DSpark need their projector/head metadata, so they require an
 explicit draft config (or a pretrained warm-start source that contains
 `config.json`). The old Domino parser exposed an optional config flag, but its

scripts/prepare_hidden_states.py

@@ -50,7 +50,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Mapping, Optional
+from typing import Callable, Dict, List, Mapping, Optional
 
 import torch
 import torch.distributed as dist
@@ -92,6 +92,7 @@ class OfflineCapturePlan:
     capture_method: str
     capture_layers: tuple[int, ...]
     layout: OfflineCaptureLayout
+    loss_mask_filter: Optional[Callable[[object], bool]]
 
 
 def parse_args():
@@ -341,6 +342,7 @@ def resolve_offline_capture_plan(
         capture_method=resolved.capture_method,
         capture_layers=resolved.capture_layers,
         layout=resolved.layout,
+        loss_mask_filter=resolved.loss_mask_filter,
     )
 
 
@@ -826,7 +828,7 @@ def main():
             ),
             num_proc=min(args.build_dataset_num_proc, 32),
         )
-    if args.num_samples is not None:
+    if args.num_samples is not None and capture_plan.loss_mask_filter is None:
         dataset = dataset.select(range(args.num_samples))
     # Tokenizer and cache key
     tokenizer = load_tokenizer(
@@ -847,6 +849,15 @@ def main():
             cache_key=cache_key,
             is_preformatted=args.is_preformatted,
             num_proc=args.build_dataset_num_proc,
+            loss_mask_filter=capture_plan.loss_mask_filter,
+        )
+    if capture_plan.loss_mask_filter is not None and args.num_samples is not None:
+        eagle3_dataset = eagle3_dataset.select(
+            range(min(args.num_samples, len(eagle3_dataset)))
+        )
+    if not len(eagle3_dataset):
+        raise ValueError(
+            f"no samples satisfy {capture_plan.strategy} training eligibility"
         )
     print_with_rank(f"Dataset prepared with {len(eagle3_dataset)} samples.")
 

specforge/algorithms/common/dflash_family_data.py

@@ -5,6 +5,7 @@
 from functools import partial
 
 from specforge.algorithms.common.collation import pad_and_concatenate_features
+from specforge.data.loss_mask import has_consecutive_supervised_tokens
 
 NORMALIZER_ID = "dflash_family_offline_v1"
 DSPARK_NORMALIZER_ID = "dspark_offline_v1"
@@ -56,6 +57,10 @@ def normalize_offline_sample(raw, max_len: int):
             f"loss_mask={loss_mask.shape[1]}, "
             f"hidden_states={hidden_states.shape[1]}"
         )
+    if not has_consecutive_supervised_tokens(loss_mask[0]):
+        raise ValueError(
+            "offline DFlash-family samples require two consecutive supervised tokens"
+        )
     return {
         "input_ids": input_ids,
         "loss_mask": loss_mask,