feat(dflash): add checkpoint resume support and clean up comments

Author: xiaomin-D

examples/run_qwen3_8b_dflash_online.sh

@@ -4,7 +4,7 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 ROOT_DIR=$(dirname $SCRIPT_DIR)
 export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
 export SPECFORGE_DATA_NUM_PROC=32
-NUM_GPUS=${1:-1}
+NUM_GPUS=${1:-8}
 
 ATTENTION_BACKEND=${2:-flex_attention}
 
@@ -14,8 +14,8 @@ torchrun \
     $ROOT_DIR/scripts/train_dflash.py \
     --target-model-path Qwen/Qwen3-8B \
     --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 \
+    --train-data-path $ROOT_DIR/cache/dataset/perfectblend_qwen3-8b_regen.jsonl \
+    --output-dir $ROOT_DIR/outputs/qwen3-8b-perfectblend \
     --num-epochs 6 \
     --batch-size 4 \
     --learning-rate 6e-4 \
@@ -31,4 +31,4 @@ torchrun \
     --save-interval 1000 \
     --report-to wandb \
     --wandb-project specforge-qwen3-8b-dflash \
-    --wandb-name qwen3-8b-dflash-sharegpt
+    --wandb-name qwen3-8b-dflash-perfectblend

scripts/train_dflash.py

@@ -33,7 +33,7 @@
 from specforge.modeling.target.target_utils import TargetEmbeddingsAndHead
 from specforge.optimizer import BF16Optimizer
 from specforge.tracker import create_tracker
-from specforge.utils import print_on_rank0, print_with_rank
+from specforge.utils import get_last_checkpoint, print_on_rank0, print_with_rank
 
 
 def parse_args():
@@ -108,6 +108,12 @@ def parse_args():
     training_group.add_argument("--accumulation-steps", type=int, default=1)
     training_group.add_argument("--seed", type=int, default=42)
     training_group.add_argument("--resume", action="store_true")
+    training_group.add_argument(
+        "--ckpt-dir",
+        type=str,
+        default=None,
+        help="Directory of the checkpoint to resume training from",
+    )
 
     output_group = parser.add_argument_group("output")
     output_group.add_argument("--output-dir", type=str, required=True)
@@ -162,25 +168,21 @@ def build_models(args) -> Tuple[DFlashTargetModel, DFlashDraftModel]:
         draft_config = AutoConfig.from_pretrained(args.draft_config_path)
         print_on_rank0(f"Loaded draft config from {args.draft_config_path}")
     else:
-        # Load config from HF (needed for structure info even if backend is sglang)
         target_config = AutoConfig.from_pretrained(args.target_model_path)
         draft_config = AutoConfig.from_pretrained(args.target_model_path)
         draft_config.num_hidden_layers = args.num_draft_layers
         draft_config.block_size = args.block_size
         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}")
 
     draft_model = DFlashDraftModel(draft_config).cuda().to(torch.bfloat16)
 
-    # Set capture layers for target model based on draft model config
     target_model.set_capture_layers(draft_model.target_layer_ids)
 
     print_on_rank0(
@@ -199,7 +201,6 @@ def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]
     """Build train and eval dataloaders."""
     import hashlib
 
-    # convert to dataloader
     cache_params_string = (
         f"{args.train_data_path}-"
         f"{args.max_length}-"
@@ -220,7 +221,6 @@ def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]
         num_proc=args.build_dataset_num_proc,
     )
 
-    # Filter out samples with too few loss tokens (DFlash requires >= 2 * block_size)
     min_loss_tokens = 2 * args.block_size
     original_size = len(train_eagle3_dataset)
     train_eagle3_dataset = train_eagle3_dataset.filter(
@@ -287,7 +287,6 @@ def save_checkpoint(args, epoch, step, dflash_model, draft_model, optimizer):
 
             draft_model.save_pretrained(save_dir, state_dict=draft_state_dict)
 
-            # Copy dflash.py for inference compatibility (matches auto_map in config)
             modeling_src = os.path.join(
                 os.path.dirname(__file__),
                 "..",
@@ -331,16 +330,13 @@ def record_metrics(
 
 
 def main():
-    # Configure logging to ensure we see INFO logs
+
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt="%m/%d/%Y %H:%M:%S",
         level=logging.INFO,
     )
-    # Force the root logger to INFO as well, just in case
     logging.getLogger().setLevel(logging.INFO)
-
-    # Filter annoying FSDP warnings
     warnings.filterwarnings(
         "ignore",
         "The .grad attribute of a Tensor that is not a leaf Tensor is being accessed",
@@ -354,9 +350,45 @@ def main():
 
     target_model, draft_model = build_models(args)
 
+    draft_model_last_checkpoint = None
+    if args.ckpt_dir is not None:
+        if os.path.isdir(args.ckpt_dir):
+            draft_model_last_checkpoint = args.ckpt_dir
+            print_on_rank0(f"Using checkpoint: {draft_model_last_checkpoint}")
+        else:
+            raise ValueError(
+                f"Provided ckpt dir {args.ckpt_dir} is not a valid directory."
+            )
+
+    if args.resume and os.path.isdir(args.output_dir):
+        draft_model_last_checkpoint = get_last_checkpoint(
+            args.output_dir, prefix=r"epoch_\d+_step"
+        )
+        print_on_rank0(f"Last checkpoint detected: {draft_model_last_checkpoint}")
+
+    resume_state = None
+    if draft_model_last_checkpoint:
+        loaded_model = DFlashDraftModel.from_pretrained(
+            draft_model_last_checkpoint, torch_dtype=torch.bfloat16
+        )
+        draft_model.load_state_dict(loaded_model.state_dict())
+        del loaded_model
+        print_on_rank0("Loaded draft model weights from checkpoint")
+
+        training_state_path = os.path.join(
+            draft_model_last_checkpoint, "training_state.pt"
+        )
+        if os.path.exists(training_state_path):
+            resume_state = torch.load(
+                training_state_path, map_location="cpu", weights_only=False
+            )
+            print_on_rank0(
+                f"Will resume from epoch {resume_state['epoch']}, "
+                f"step {resume_state['global_step']}"
+            )
+
     tokenizer = AutoTokenizer.from_pretrained(args.target_model_path)
 
-    # Get mask_token_id
     if args.mask_token_id is not None:
         mask_token_id = args.mask_token_id
     elif tokenizer.mask_token_id is not None:
@@ -366,9 +398,6 @@ 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
@@ -380,10 +409,7 @@ def main():
     total_steps = args.num_epochs * steps_per_epoch
     print_on_rank0(f"Total training steps: {total_steps}")
 
-    # Note: We need embedding layer for DFlash wrapper.
-    # For SGLang backend, we can't easily get the embedding layer object.
-    # We use TargetEmbeddingsAndHead to efficiently load only needed weights.
-    print_on_rank0("Loading target embeddings and head efficiently...")
+    print_on_rank0("Loading target embeddings and head...")
     target_components = TargetEmbeddingsAndHead.from_pretrained(
         args.target_model_path,
         embed_key="model.embed_tokens.weight",  # Adjust if Qwen/Llama differs
@@ -423,14 +449,25 @@ def main():
         total_steps=total_steps,
     )
 
+    start_epoch = 0
+    global_step = 0
+    if resume_state is not None:
+        optimizer.scheduler.load_state_dict(resume_state["scheduler_state_dict"])
+        start_epoch = resume_state["epoch"]
+        global_step = resume_state["global_step"]
+        del resume_state
+        print_on_rank0(f"Restored scheduler, lr={optimizer.get_learning_rate():.6f}")
+
+    skip_steps = global_step - start_epoch * len(train_dataloader)
+
     print_on_rank0(f"Initializing tracker (report_to={args.report_to})...")
     tracker = create_tracker(args, args.output_dir)
     print_on_rank0("Tracker initialized successfully.")
 
-    global_step = 0
     last_time = time.time()
+    print_on_rank0(f"Starting training from epoch {start_epoch}, step {global_step}")
 
-    for epoch in range(args.num_epochs):
+    for epoch in range(start_epoch, args.num_epochs):
         train_dataloader.sampler.set_epoch(epoch)
         draft_model.train()
 
@@ -441,21 +478,20 @@ def main():
         else:
             progress_bar = train_dataloader
 
-        for data in progress_bar:
+        for step_in_epoch, data in enumerate(progress_bar):
+            if epoch == start_epoch and step_in_epoch < skip_steps:
+                continue
             global_step += 1
 
             input_ids = data["input_ids"].cuda()
             attention_mask = data["attention_mask"].cuda()
             loss_mask = data["loss_mask"].cuda()
 
-            # Generate context from Target Model (SGLang or HF)
-            # This calls the backend to get hidden states
             target_output = target_model.generate_dflash_data(
                 input_ids, attention_mask, loss_mask
             )
             hidden_states = target_output.hidden_states.cuda()  # Ensure on GPU
 
-            # Forward pass (Parallel Training)
             loss, accuracy = dflash_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,