Feat/dflash training improvements (#463)

* feat(dflash): add random anchor sampling, loss decay, and sync with upstream

- Add random anchor sampling for block construction (paper Sec 4.2)
- Add exponential loss decay weighting (paper Sec 4.2, Eq.4, Appendix A.3.1)
- Sync with upstream: dflash_config in config.json, mask_token_id from config, decoupled target_layer_ids
- Align training hyperparams with paper (lr=6e-4, warmup=0.04, epochs=6, max_length=3072)
- Fix auto_map and saved model file name for HuggingFace compatibility

* fix(dflash): per-sample anchor sampling with padding block isolation

- Sample anchors independently per batch sample (max strategy)
- Mark padding blocks with block_id=-1 for attention isolation
- Padding blocks excluded from both attention and loss computation
- Use absolute positions from gather_idx for position encoding
- Per-sample block_ids throughout: attention mask, loss mask, noise input

* fix(dflash): align acceptance rate metric with inference and trust_remote_code kwarg

- Use loss_mask to exclude prompt tokens from acc calculation, only
  measure on completion/assistant blocks to match inference behavior
- Replace token-level accuracy with cumprod-based acceptance length
- Clean up debug prints and redundant comments
- Add trust_remote_code kwarg

Author: xiaomin-D

configs/longcat-flash-dflash.json

@@ -5,10 +5,14 @@
     "attention_bias": false,
     "attention_dropout": 0.0,
     "auto_map": {
-      "AutoModel": "modeling_dflash.DFlashDraftModel"
+      "AutoModel": "dflash.DFlashDraftModel"
     },
     "block_size": 16,
     "bos_token_id": 1,
+    "dflash_config": {
+      "mask_token_id": 2,
+      "target_layer_ids": [1, 7, 13, 19, 25]
+    },
     "dtype": "bfloat16",
     "eos_token_id": 2,
     "head_dim": 128,

configs/qwen3-8b-dflash.json

@@ -5,10 +5,14 @@
   "attention_bias": false,
   "attention_dropout": 0.0,
   "auto_map": {
-    "AutoModel": "modeling_dflash.DFlashDraftModel"
+    "AutoModel": "dflash.DFlashDraftModel"
   },
   "block_size": 16,
   "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [1, 9, 17, 25, 33]
+  },
   "dtype": "bfloat16",
   "eos_token_id": 151645,
   "head_dim": 128,

examples/run_longcat_flash_dflash_online.sh

@@ -26,11 +26,16 @@ torchrun \
     --train-data-path $ROOT_DIR/cache/dataset/sharegpt_train.jsonl \
     --build-dataset-num-proc $BUILD_DATASET_NUM_PROC \
     --output-dir $ROOT_DIR/outputs/longcat-flash-dflash-sharegpt \
-    --num-epochs 20 \
+    --num-epochs 6 \
     --batch-size 2 \
-    --learning-rate 1e-4 \
-    --max-length 2048 \
+    --learning-rate 6e-4 \
+    --warmup-ratio 0.04 \
+    --max-grad-norm 1.0 \
+    --max-length 3072 \
     --chat-template longcat \
+    --random-anchor \
+    --num-anchors 512 \
+    --loss-decay-gamma 7.0 \
     --log-interval 50 \
     --save-interval 1000 \
     --report-to wandb \

examples/run_qwen3_8b_dflash_online.sh

@@ -16,12 +16,17 @@ torchrun \
     --draft-config-path $ROOT_DIR/configs/qwen3-8b-dflash.json \
     --train-data-path $ROOT_DIR/cache/dataset/sharegpt_train.jsonl \
     --output-dir $ROOT_DIR/outputs/qwen3-8b-dflash-sharegpt \
-    --num-epochs 20 \
+    --num-epochs 6 \
     --batch-size 4 \
-    --learning-rate 1e-4 \
-    --max-length 2048 \
+    --learning-rate 6e-4 \
+    --warmup-ratio 0.04 \
+    --max-grad-norm 1.0 \
+    --max-length 3072 \
     --chat-template qwen \
     --attention-backend $ATTENTION_BACKEND \
+    --random-anchor \
+    --num-anchors 512 \
+    --loss-decay-gamma 7.0 \
     --log-interval 50 \
     --save-interval 1000 \
     --report-to wandb \

scripts/train_dflash.py

@@ -67,6 +67,24 @@ def parse_args():
     model_group.add_argument(
         "--trust-remote-code", action="store_true", help="Trust remote code"
     )
+    model_group.add_argument(
+        "--random-anchor",
+        action="store_true",
+        help="Enable random anchor sampling for block construction (paper Sec 4.2).",
+    )
+    model_group.add_argument(
+        "--num-anchors",
+        type=int,
+        default=512,
+        help="Number of anchor positions per sequence when --random-anchor is set.",
+    )
+    model_group.add_argument(
+        "--loss-decay-gamma",
+        type=float,
+        default=None,
+        help="Gamma for exponential loss decay weighting (paper Eq.4). "
+        "Suggested: 7 for block_size=16, 5 for 10, 4 for 8. None disables.",
+    )
 
     dataset_group = parser.add_argument_group("dataset")
     dataset_group.add_argument("--train-data-path", type=str, required=True)
@@ -81,11 +99,11 @@ def parse_args():
     )
 
     training_group = parser.add_argument_group("training")
-    training_group.add_argument("--num-epochs", type=int, default=3)
+    training_group.add_argument("--num-epochs", type=int, default=6)
     training_group.add_argument("--batch-size", type=int, default=1)
-    training_group.add_argument("--learning-rate", type=float, default=1e-4)
-    training_group.add_argument("--max-length", type=int, default=2048)
-    training_group.add_argument("--warmup-ratio", type=float, default=0.01)
+    training_group.add_argument("--learning-rate", type=float, default=6e-4)
+    training_group.add_argument("--max-length", type=int, default=3072)
+    training_group.add_argument("--warmup-ratio", type=float, default=0.04)
     training_group.add_argument("--max-grad-norm", type=float, default=1.0)
     training_group.add_argument("--accumulation-steps", type=int, default=1)
     training_group.add_argument("--seed", type=int, default=42)
@@ -152,6 +170,10 @@ def build_models(args) -> Tuple[DFlashTargetModel, DFlashDraftModel]:
         draft_config.num_target_layers = target_config.num_hidden_layers
         print_on_rank0("Auto-generated draft config from target model")
 
+    # Ensure dflash_config exists in config (for target_layer_ids / mask_token_id)
+    if not hasattr(draft_config, "dflash_config") or draft_config.dflash_config is None:
+        draft_config.dflash_config = {}
+
     # Set attention implementation based on backend
     draft_config._attn_implementation = args.attention_backend
     print_on_rank0(f"Using attention backend: {args.attention_backend}")
@@ -265,7 +287,7 @@ def save_checkpoint(args, epoch, step, dflash_model, draft_model, optimizer):
 
             draft_model.save_pretrained(save_dir, state_dict=draft_state_dict)
 
-            # Copy modeling_dflash.py for inference compatibility
+            # Copy dflash.py for inference compatibility (matches auto_map in config)
             modeling_src = os.path.join(
                 os.path.dirname(__file__),
                 "..",
@@ -274,7 +296,7 @@ def save_checkpoint(args, epoch, step, dflash_model, draft_model, optimizer):
                 "draft",
                 "dflash.py",
             )
-            modeling_dst = os.path.join(save_dir, "modeling_dflash.py")
+            modeling_dst = os.path.join(save_dir, "dflash.py")
             if os.path.exists(modeling_src):
                 shutil.copy(modeling_src, modeling_dst)
 
@@ -344,6 +366,14 @@ def main():
         mask_token_id = tokenizer.mask_token_id
     print_on_rank0(f"Using mask_token_id: {mask_token_id}")
 
+    # Write mask_token_id and target_layer_ids into draft config so that
+    # save_pretrained produces a config.json compatible with the official
+    # dflash inference code (which reads from config.dflash_config).
+    draft_model.mask_token_id = mask_token_id
+    draft_model.config.dflash_config["mask_token_id"] = mask_token_id
+    draft_model.config.dflash_config["target_layer_ids"] = draft_model.target_layer_ids
+    print_on_rank0(f"dflash_config: {draft_model.config.dflash_config}")
+
     train_dataloader, eval_dataloader = build_dataloader(args, tokenizer)
 
     steps_per_epoch = math.ceil(len(train_dataloader) / args.accumulation_steps)
@@ -369,6 +399,9 @@ def main():
         block_size=draft_model.block_size,
         mask_token_id=mask_token_id,
         attention_backend=args.attention_backend,
+        random_anchor=args.random_anchor,
+        num_anchors=args.num_anchors,
+        loss_decay_gamma=args.loss_decay_gamma,
     )
 
     dflash_model = FSDP(