syxin/Specforge/scripts/train_dflash_lora_inject.py

#!/usr/bin/env python3
# coding=utf-8
"""DFlash LoRA Training Script with Direct Layer Injection.

Trains Qwen3-8B with LoRA adapters using direct layer-by-layer injection from target model.

Key features:
  - Target model extracts hidden states at each layer
  - Draft model (same structure + LoRA) receives these hidden states layer-by-layer
  - No feature extraction layers (fc + hidden_norm) - direct injection
  - Only LoRA parameters are trained; base model is frozen
  - Saves LoRA adapter weights only
"""

import argparse
import json
import logging
import math
import os
import time
import warnings
from typing import Optional, Tuple

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from accelerate.utils import set_seed

from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import AutoTokenizer

from datasets import load_dataset
from specforge.args import TrackerArgs
from specforge.core.dflash_lora_inject import OnlineDFlashLoRAInjectModel
from specforge.data import build_eagle3_dataset, prepare_dp_dataloaders
from specforge.distributed import destroy_distributed, get_dp_group, init_distributed
from specforge.modeling.draft.dflash_lora_inject import DFlashLoRAInjectDraftModel
from specforge.modeling.target.dflash_target_model import get_dflash_target_model
from specforge.optimizer import BF16Optimizer
from specforge.tracker import create_tracker
from specforge.utils import get_last_checkpoint, print_on_rank0, print_with_rank


def parse_args():
    parser = argparse.ArgumentParser(description="Train DFlash LoRA with Direct Layer Injection")

    model_group = parser.add_argument_group("model")
    model_group.add_argument("--target-model-path", type=str, required=True,
                             help="Path to target model (for extracting hidden states)")
    model_group.add_argument("--target-model-backend", type=str, default="hf",
                             choices=["hf", "sglang"],
                             help="Backend for target model")
    model_group.add_argument("--draft-model-path", type=str, default=None,
                             help="Path to draft model base (default: same as target)")
    model_group.add_argument("--block-size", type=int, default=16)
    model_group.add_argument("--mask-token-id", type=int, default=None,
                             help="MASK token ID. Auto-detected from tokenizer if not set.")
    model_group.add_argument("--context-len", type=int, default=0,
                             help="Fixed context length before blocks. 0 = treat whole seq as blocks.")
    model_group.add_argument("--trust-remote-code", action="store_true")
    model_group.add_argument("--attn-implementation", type=str, default="sdpa",
                             choices=["sdpa", "eager"],
                             help="Attention backend for additive mask path.")
    model_group.add_argument("--attention-backend", type=str, default="flex_attention",
                             choices=["flex_attention", "additive"],
                             help="flex_attention: use BlockMask. additive: use 4D mask.")

    lora_group = parser.add_argument_group("lora")
    lora_group.add_argument("--lora-rank", type=int, default=16)
    lora_group.add_argument("--lora-alpha", type=int, default=32)
    lora_group.add_argument("--lora-dropout", type=float, default=0.05)
    lora_group.add_argument("--lora-target-modules", type=str, nargs="+",
                            default=["q_proj", "k_proj", "v_proj", "o_proj"],
                            help="Which modules to apply LoRA to")
    lora_group.add_argument("--lora-config", type=str, default=None,
                            help="Path to JSON file with LoRA config")

    dataset_group = parser.add_argument_group("dataset")
    dataset_group.add_argument("--train-data-path", type=str, required=True)
    dataset_group.add_argument("--eval-data-path", type=str, default=None)
    dataset_group.add_argument("--chat-template", type=str, default="qwen")
    dataset_group.add_argument("--is-preformatted", action="store_true")
    dataset_group.add_argument("--dataloader-num-workers", type=int, default=8)
    dataset_group.add_argument("--build-dataset-num-proc", type=int,
                               default=int(os.environ.get("SPECFORGE_DATA_NUM_PROC", 8)))

    training_group = parser.add_argument_group("training")
    training_group.add_argument("--num-epochs", type=int, default=3)
    training_group.add_argument("--batch-size", type=int, default=1)
    training_group.add_argument("--learning-rate", type=float, default=2e-4)
    training_group.add_argument("--max-length", type=int, default=2048)
    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("--loss-decay-gamma", type=float, default=None)
    training_group.add_argument("--optimizer-type", type=str, default="adamw",
                                choices=["adamw", "adamw_8bit"])
    training_group.add_argument("--no-fp32-params", action="store_true")
    training_group.add_argument("--gradient-checkpointing", action="store_true")
    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)

    output_group = parser.add_argument_group("output")
    output_group.add_argument("--output-dir", type=str, required=True)
    output_group.add_argument("--cache-dir", type=str, default="./cache")
    output_group.add_argument("--log-interval", type=int, default=50)
    output_group.add_argument("--eval-interval", type=int, default=1000)
    output_group.add_argument("--save-interval", type=int, default=1000)

    early_stop_group = parser.add_argument_group("early stopping")
    early_stop_group.add_argument("--early-stop", action="store_true",
                                  help="Enable early stopping based on training accuracy")
    early_stop_group.add_argument("--early-stop-patience", type=int, default=5,
                                  help="Stop after N consecutive log intervals without improvement (default: 5)")
    early_stop_group.add_argument("--early-stop-min-delta", type=float, default=0.005,
                                  help="Minimum accuracy improvement to reset patience (default: 0.005)")
    early_stop_group.add_argument("--early-stop-acc-threshold", type=float, default=None,
                                  help="Hard accuracy threshold — stop immediately when reached (default: None)")
    early_stop_group.add_argument("--early-stop-warmup-steps", type=int, default=0,
                                  help="Number of log intervals to skip before enabling early stopping (default: 0)")
    early_stop_group.add_argument("--early-stop-relative-delta", action="store_true",
                                  help="Treat min_delta as a fraction of best_acc instead of absolute value")

    dist_group = parser.add_argument_group("distributed")
    dist_group.add_argument("--dist-timeout", type=int, default=30)

    tracker_group = parser.add_argument_group("tracker")
    TrackerArgs.add_args(tracker_group)

    return parser.parse_args()


def build_model(args) -> Tuple[DFlashLoRAInjectDraftModel, OnlineDFlashLoRAInjectModel, any]:
    """Load target model and draft model with LoRA."""
    print_on_rank0(f"Loading target model from {args.target_model_path}")

    # Load target model
    target_model = get_dflash_target_model(
        pretrained_model_name_or_path=args.target_model_path,
        backend=args.target_model_backend,
        torch_dtype=torch.bfloat16,
        device="cuda",
        cache_dir=args.cache_dir,
        trust_remote_code=args.trust_remote_code,
    )

    # Set target model to capture all layer hidden states
    # For HF backend, this will return all layers
    # For SGLang backend, might need specific configuration
    if hasattr(target_model, 'set_capture_layers'):
        # Capture all layers - will be determined by model's num_layers
        target_model.set_capture_layers(None)

    print_on_rank0(f"Loading draft model from {args.draft_model_path or args.target_model_path}")

    # Load LoRA config
    lora_rank = args.lora_rank
    lora_alpha = args.lora_alpha
    lora_dropout = args.lora_dropout
    lora_target_modules = args.lora_target_modules

    if args.lora_config is not None:
        with open(args.lora_config) as f:
            lora_cfg = json.load(f)
        lora_rank = lora_cfg.get("lora_rank", lora_rank)
        lora_alpha = lora_cfg.get("lora_alpha", lora_alpha)
        lora_dropout = lora_cfg.get("lora_dropout", lora_dropout)
        lora_target_modules = lora_cfg.get("lora_target_modules", lora_target_modules)
        print_on_rank0(f"Loaded LoRA config from {args.lora_config}")

    # Resolve attention implementation
    if args.attention_backend == "flex_attention":
        attn_impl = "flex_attention"
    else:
        attn_impl = args.attn_implementation

    # Load draft model with LoRA
    draft_model = DFlashLoRAInjectDraftModel.from_pretrained(
        pretrained_model_name_or_path=args.draft_model_path or args.target_model_path,
        lora_rank=lora_rank,
        lora_alpha=lora_alpha,
        lora_dropout=lora_dropout,
        lora_target_modules=lora_target_modules,
        block_size=args.block_size,
        mask_token_id=args.mask_token_id or 151669,
        torch_dtype=torch.bfloat16,
        device_map="cuda",
        trust_remote_code=args.trust_remote_code,
        attn_implementation=attn_impl,
    )

    # Wrap in training model
    online_model = OnlineDFlashLoRAInjectModel(
        draft_model=draft_model,
        target_model=target_model,
        block_size=args.block_size,
        mask_token_id=args.mask_token_id or 151669,
        loss_decay_gamma=args.loss_decay_gamma,
        attention_backend=args.attention_backend,
        lm_head_chunk_size=0,  # Not implemented for inject version yet
        random_anchor=False,  # Not implemented for inject version yet
        num_anchors=512,
    )

    trainable = sum(p.numel() for p in draft_model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in draft_model.parameters())
    print_on_rank0(f"Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")

    return draft_model, online_model, target_model


def build_dataloader(args, tokenizer) -> Tuple[DataLoader, Optional[DataLoader]]:
    """Build train and eval dataloaders."""
    import hashlib

    cache_params_string = (
        f"{args.train_data_path}-{args.max_length}-{args.chat_template}-{args.target_model_path}"
    )
    cache_key = hashlib.md5(cache_params_string.encode()).hexdigest()

    rank = dist.get_rank()

    # Load dataset
    if os.path.isdir(args.train_data_path):
        train_dataset = load_dataset(args.train_data_path, split="train", verification_mode="no_checks")
    else:
        train_dataset = load_dataset("json", data_files=args.train_data_path)["train"]

    dataset_kwargs = dict(
        dataset=train_dataset,
        tokenizer=tokenizer,
        chat_template=args.chat_template,
        max_length=args.max_length,
        is_preformatted=args.is_preformatted,
        cache_dir=os.path.join(args.cache_dir, "processed_dataset"),
        cache_key=cache_key,
        num_proc=args.build_dataset_num_proc,
    )

    # Only rank 0 preprocesses, others wait
    if rank == 0:
        train_eagle3_dataset = build_eagle3_dataset(**dataset_kwargs)
    dist.barrier()
    if rank != 0:
        train_eagle3_dataset = build_eagle3_dataset(**dataset_kwargs)

    min_loss_tokens = 2 * args.block_size
    original_size = len(train_eagle3_dataset)
    train_eagle3_dataset = train_eagle3_dataset.filter(
        lambda x: x["loss_mask"].sum() >= min_loss_tokens
    )
    print_on_rank0(f"Filtered train dataset: {original_size} -> {len(train_eagle3_dataset)} samples")

    train_dataloader = prepare_dp_dataloaders(
        train_eagle3_dataset,
        args.batch_size,
        num_workers=args.dataloader_num_workers,
        shuffle=True,
        process_group=get_dp_group(),
    )

    eval_dataloader = None
    if args.eval_data_path:
        if os.path.isdir(args.eval_data_path):
            eval_dataset = load_dataset(args.eval_data_path, split="train")
        else:
            eval_dataset = load_dataset("json", data_files=args.eval_data_path)["train"]
        eval_eagle3_dataset = build_eagle3_dataset(
            dataset=eval_dataset,
            tokenizer=tokenizer,
            chat_template=args.chat_template,
            max_length=args.max_length,
            is_preformatted=args.is_preformatted,
        )
        eval_dataloader = prepare_dp_dataloaders(
            eval_eagle3_dataset,
            args.batch_size,
            num_workers=args.dataloader_num_workers,
            shuffle=False,
            process_group=get_dp_group(),
        )

    return train_dataloader, eval_dataloader


def save_checkpoint(args, epoch, step, online_model, draft_model, optimizer):
    """Save LoRA adapter weights + training state."""
    save_dir = os.path.join(args.output_dir, f"epoch_{epoch}_step_{step}")
    if dist.get_rank() == 0:
        os.makedirs(save_dir, exist_ok=True)
    dist.barrier()

    if dist.get_rank() == 0:
        # Unwrap DDP
        module = online_model.module if isinstance(online_model, DDP) else online_model
        lora_state_dict = {
            k: v for k, v in module.draft_model.model.state_dict().items()
            if "lora_" in k
        }

        try:
            from safetensors.torch import save_file as safetensors_save
            safetensors_save(lora_state_dict, os.path.join(save_dir, "adapter_model.safetensors"))
        except (ImportError, Exception):
            torch.save(lora_state_dict, os.path.join(save_dir, "adapter_model.bin"))

        draft_model.model.peft_config["default"].save_pretrained(save_dir)

        torch.save(
            {
                "epoch": epoch,
                "global_step": step,
                "args": args,
                **optimizer.state_dict(),
            },
            os.path.join(save_dir, "training_state.pt"),
        )
        print_on_rank0(f"Saved LoRA checkpoint to {save_dir}")

    dist.barrier()


def record_metrics(args, loss, accuracy, global_step, tracker, optimizer, mode="train"):
    logdict = {}
    if mode == "train" and optimizer is not None:
        logdict["train/lr"] = optimizer.get_learning_rate()
    logdict[f"{mode}/loss"] = loss
    logdict[f"{mode}/accuracy"] = accuracy
    print_on_rank0(
        f"{mode.capitalize()} - Step {global_step}, Loss: {loss:.4f}, Acc: {accuracy:.4f}"
    )
    tracker.log(logdict, step=global_step)


class EarlyStopping:
    """Monitor accuracy and signal when training should stop."""

    def __init__(self, patience: int, min_delta: float, acc_threshold: float = None,
                 warmup_steps: int = 0, relative_delta: bool = False):
        self.patience = patience
        self.min_delta = min_delta
        self.acc_threshold = acc_threshold
        self.warmup_steps = warmup_steps
        self.relative_delta = relative_delta
        self.best_acc = -1.0
        self.counter = 0
        self.num_calls = 0

    def should_stop(self, acc: float) -> bool:
        self.num_calls += 1
        # Skip checking during warmup phase
        if self.num_calls <= self.warmup_steps:
            # Still track best_acc during warmup
            if acc > self.best_acc:
                self.best_acc = acc
            return False
        # Hard accuracy threshold check
        if self.acc_threshold is not None and acc >= self.acc_threshold:
            return True
        # Patience-based check with optional relative delta
        if self.relative_delta and self.best_acc > 0:
            delta = self.min_delta * self.best_acc
        else:
            delta = self.min_delta
        if acc > self.best_acc + delta:
            self.best_acc = acc
            self.counter = 0
        else:
            self.counter += 1
        return self.counter >= self.patience


def main():
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    warnings.filterwarnings(
        "ignore",
        "The .grad attribute of a Tensor that is not a leaf Tensor is being accessed",
    )

    mp.set_sharing_strategy('file_system')

    args = parse_args()
    set_seed(args.seed)

    # tp_size=1: LoRA training doesn't use tensor parallelism
    init_distributed(timeout=args.dist_timeout, tp_size=1)
    print_with_rank("Initialized distributed")

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.target_model_path)
    if args.mask_token_id is not None:
        mask_token_id = args.mask_token_id
    elif tokenizer.mask_token_id is not None:
        mask_token_id = tokenizer.mask_token_id
    else:
        tokenizer.add_special_tokens({"mask_token": "<|MASK|>"})
        mask_token_id = tokenizer.mask_token_id
    print_on_rank0(f"Using mask_token_id: {mask_token_id}")
    args.mask_token_id = mask_token_id

    draft_model, online_model, target_model = build_model(args)

    # Update mask_token_id
    draft_model.mask_token_id = mask_token_id
    online_model.mask_token_id = mask_token_id

    if args.gradient_checkpointing:
        draft_model.gradient_checkpointing_enable(
            gradient_checkpointing_kwargs={"use_reentrant": False}
        )
        print_on_rank0("Gradient checkpointing enabled")

    # Resume from checkpoint
    resume_state = None
    if args.ckpt_dir is not None:
        if os.path.isdir(args.ckpt_dir):
            print_on_rank0(f"Loading LoRA weights from {args.ckpt_dir}")
            from peft import PeftModel
            draft_model.model = PeftModel.from_pretrained(
                draft_model.model.base_model.model, args.ckpt_dir
            )
        else:
            raise ValueError(f"ckpt_dir {args.ckpt_dir} is not a valid directory")

    if args.resume and os.path.isdir(args.output_dir):
        last_ckpt = get_last_checkpoint(args.output_dir, prefix=r"epoch_\d+_step")
        if last_ckpt:
            print_on_rank0(f"Resuming from {last_ckpt}")
            from peft import PeftModel
            draft_model.model = PeftModel.from_pretrained(
                draft_model.model.base_model.model, last_ckpt
            )
            training_state_path = os.path.join(last_ckpt, "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']}, step {resume_state['global_step']}"
                )

    train_dataloader, eval_dataloader = build_dataloader(args, tokenizer)

    steps_per_epoch = math.ceil(len(train_dataloader) / args.accumulation_steps)
    total_steps = args.num_epochs * steps_per_epoch
    print_on_rank0(f"Total training steps: {total_steps}")

    # Use local rank for device_ids in multi-node training
    local_rank = int(os.environ.get("LOCAL_RANK", dist.get_rank() % torch.cuda.device_count()))
    online_model = DDP(online_model, device_ids=[local_rank], find_unused_parameters=False)
    print_with_rank("Initialized DDP")

    optimizer = BF16Optimizer(
        draft_model,
        lr=args.learning_rate,
        max_grad_norm=args.max_grad_norm,
        warmup_ratio=args.warmup_ratio,
        total_steps=total_steps,
        use_fp32_params=not args.no_fp32_params,
        optimizer_type=args.optimizer_type,
    )

    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)

    tracker = create_tracker(args, args.output_dir)
    last_time = time.time()
    print_on_rank0(f"Starting training from epoch {start_epoch}, step {global_step}")

    # Early stopping
    early_stopper = None
    if args.early_stop:
        early_stopper = EarlyStopping(
            patience=args.early_stop_patience,
            min_delta=args.early_stop_min_delta,
            acc_threshold=args.early_stop_acc_threshold,
            warmup_steps=args.early_stop_warmup_steps,
            relative_delta=args.early_stop_relative_delta,
        )
        print_on_rank0(
            f"Early stopping enabled: patience={args.early_stop_patience}, "
            f"min_delta={args.early_stop_min_delta}, "
            f"relative_delta={args.early_stop_relative_delta}, "
            f"warmup_steps={args.early_stop_warmup_steps}, "
            f"acc_threshold={args.early_stop_acc_threshold}"
        )
    should_stop = False

    for epoch in range(start_epoch, args.num_epochs):
        train_dataloader.sampler.set_epoch(epoch)
        draft_model.train()

        if dist.get_rank() == 0:
            progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch}", leave=True)
        else:
            progress_bar = train_dataloader

        for step_in_epoch, data in enumerate(progress_bar):
            global_step += 1
            if epoch == start_epoch and step_in_epoch < skip_steps:
                continue

            input_ids = data["input_ids"].cuda()
            attention_mask = data["attention_mask"].cuda()
            loss_mask = data["loss_mask"].cuda()

            loss, accuracy = online_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                loss_mask=loss_mask,
                context_len=args.context_len,
            )
            (loss / args.accumulation_steps).backward()

            if global_step % args.accumulation_steps == 0:
                optimizer.step()

            if global_step % args.log_interval == 0:
                loss_val = loss.item()
                acc_val = accuracy.item()
                loss_t = torch.tensor(loss_val, device="cuda")
                acc_t = torch.tensor(acc_val, device="cuda")
                dist.all_reduce(loss_t)
                dist.all_reduce(acc_t)
                avg_acc = acc_t.item() / dist.get_world_size()
                record_metrics(args, loss_t.item() / dist.get_world_size(),
                               avg_acc, global_step,
                               tracker, optimizer, mode="train")

                # Early stopping check
                if early_stopper is not None:
                    stop_flag = torch.tensor(0, device="cuda")
                    if dist.get_rank() == 0:
                        if early_stopper.should_stop(avg_acc):
                            stop_flag.fill_(1)
                            print_on_rank0(
                                f"Early stopping triggered at step {global_step}, "
                                f"best_acc={early_stopper.best_acc:.4f}, "
                                f"patience={early_stopper.counter}/{early_stopper.patience}"
                            )
                    dist.broadcast(stop_flag, src=0)
                    if stop_flag.item() == 1:
                        save_checkpoint(args, epoch, global_step, online_model, draft_model, optimizer)
                        should_stop = True
                        break

            if dist.get_rank() == 0:
                elapsed = time.time() - last_time
                last_time = time.time()
                progress_bar.set_postfix({
                    "loss": f"{loss.item():.4f}",
                    "acc": f"{accuracy.item():.4f}",
                    "iter_time": f"{elapsed:.2f}s",
                })

            if global_step % args.save_interval == 0:
                save_checkpoint(args, epoch, global_step, online_model, draft_model, optimizer)

        if should_stop:
            break

    save_checkpoint(args, args.num_epochs, global_step, online_model, draft_model, optimizer)
    tracker.close()
    destroy_distributed()


if __name__ == "__main__":
    main()