diff --git a/configs/qwen3-8b-domino.json b/configs/qwen3-8b-domino.json new file mode 100644 index 000000000..f7546bb0f --- /dev/null +++ b/configs/qwen3-8b-domino.json @@ -0,0 +1,50 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dflash.DFlashDraftModel" + }, + "block_size": 16, + "bos_token_id": 151643, + "dflash_config": { + "mask_token_id": 151669, + "target_layer_ids": [1, 9, 17, 25, 33], + "projector_type": "domino", + "pure_draft_prefix_len": 1, + "emb_dim": 256, + "gru_hidden_dim": 1024, + "shift_label": true + }, + "dtype": "bfloat16", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 40960, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "num_target_layers": 36, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": false, + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/examples/run_qwen3_8b_domino_online.sh b/examples/run_qwen3_8b_domino_online.sh new file mode 100755 index 000000000..b000919e4 --- /dev/null +++ b/examples/run_qwen3_8b_domino_online.sh @@ -0,0 +1,47 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +ROOT_DIR=$(dirname "$SCRIPT_DIR") + +export CUDA_HOME=/usr/local/cuda +export CONDA_PREFIX=${CONDA_PREFIX:-$HOME/.conda/envs/specforge} + +export PATH=$CUDA_HOME/bin:$CONDA_PREFIX/bin:$PATH +export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$CONDA_PREFIX/lib:$LD_LIBRARY_PATH + +export NCCL_DEBUG=WARN +export SPECFORGE_DATA_NUM_PROC=32 + +NUM_GPUS=${1:-8} +ATTENTION_BACKEND=${2:-flex_attention} +TARGET_MODEL_PATH=${TARGET_MODEL_PATH:-/path/to/Qwen3-8B} +TRAIN_DATA_PATH=${TRAIN_DATA_PATH:-/path/to/sharegpt_train.jsonl} + +torchrun \ + --standalone \ + --nproc_per_node $NUM_GPUS \ + $ROOT_DIR/scripts/train_domino.py \ + --target-model-path $TARGET_MODEL_PATH \ + --draft-config-path $ROOT_DIR/configs/qwen3-8b-domino.json \ + --train-data-path $TRAIN_DATA_PATH \ + --output-dir $ROOT_DIR/outputs/same_data/qwen3-8b-domino_sharegpt \ + --num-epochs 6 \ + --batch-size 2 \ + --learning-rate 6e-4 \ + --warmup-ratio 0.04 \ + --max-grad-norm 1.0 \ + --max-length 3072 \ + --chat-template qwen \ + --attention-backend $ATTENTION_BACKEND \ + --num-anchors 256 \ + --loss-decay-gamma 7.0 \ + --log-interval 50 \ + --save-interval 2000 \ + --report-to wandb \ + --wandb-project specforge-qwen3-8b-domino \ + --target-model-backend sglang \ + --block-size 16 \ + --lambda-base-start 1.0 \ + --lambda-base-decay-ratio 1.0 \ + --wandb-name qwen3-8b-domino_sharegpt diff --git a/scripts/train_domino.py b/scripts/train_domino.py new file mode 100755 index 000000000..98beae896 --- /dev/null +++ b/scripts/train_domino.py @@ -0,0 +1,663 @@ +#!/usr/bin/env python3 +# coding=utf-8 +"""Domino Training Script.""" + +import argparse +import logging +import math +import os +import shutil +import time +import warnings +from typing import Optional, Tuple + +import torch +import torch.distributed as dist +from accelerate.utils import set_seed +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoConfig, AutoTokenizer + +from datasets import load_dataset +from specforge.args import SGLangBackendArgs, TrackerArgs +from specforge.core.domino import OnlineDominoModel +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 import DFlashDraftModel +from specforge.modeling.target.dflash_target_model import ( + DFlashTargetModel, + get_dflash_target_model, +) +from specforge.modeling.target.target_utils import TargetEmbeddingsAndHead +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 Domino Draft Model") + + model_group = parser.add_argument_group("model") + model_group.add_argument("--target-model-path", type=str, required=True) + model_group.add_argument( + "--target-model-backend", + type=str, + default="hf", + choices=["sglang", "hf"], + help="Backend for target model: 'sglang' (service) or 'hf' (local)", + ) + model_group.add_argument("--draft-config-path", type=str, default=None) + model_group.add_argument("--block-size", type=int, default=16) + model_group.add_argument("--num-draft-layers", type=int, default=1) + model_group.add_argument( + "--mask-token-id", + type=int, + default=None, + help="MASK token ID. If not provided, auto-detect from tokenizer.", + ) + model_group.add_argument( + "--attention-backend", + type=str, + default="flex_attention", + choices=["eager", "sdpa", "flex_attention"], + help="Attention backend for draft model.", + ) + model_group.add_argument( + "--trust-remote-code", action="store_true", help="Trust remote code" + ) + model_group.add_argument( + "--num-anchors", + type=int, + default=512, + help="Number of anchor positions per sequence", + ) + 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.", + ) + model_group.add_argument( + "--embedding-key", + type=str, + default=None, + help="Embedding weight key in the target model. " + "Default: 'model.embed_tokens.weight' for standard models, " + "'model.language_model.embed_tokens.weight' for multimodal models like Qwen3.5-A3B.", + ) + model_group.add_argument( + "--lm-head-key", + type=str, + default=None, + help="LM head weight key in the target model. Default: 'lm_head.weight'.", + ) + + 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=6) + training_group.add_argument("--batch-size", type=int, default=1) + 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) + training_group.add_argument("--resume", action="store_true") + training_group.add_argument( + "--lambda-base-start", + type=float, + default=1.0, + help="Initial weight of base loss.", + ) + training_group.add_argument( + "--lambda-base-decay-ratio", + type=float, + default=0.5, + help="Fraction of total steps used to decay lambda_base to 0.", + ) + 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) + + optimization_group = parser.add_argument_group("optimization") + optimization_group.add_argument( + "--tp-size", + type=int, + default=1, + help="The size of the tensor parallel for the target model", + ) + + tracker_group = parser.add_argument_group("tracker") + TrackerArgs.add_args(tracker_group) + + dist_group = parser.add_argument_group("distributed") + dist_group.add_argument("--dist-timeout", type=int, default=30) + + # SGLang specific args + sglang_group = parser.add_argument_group("sglang backend") + SGLangBackendArgs.add_args(sglang_group) + + return parser.parse_args() + + +def build_models(args) -> Tuple[DFlashTargetModel, DFlashDraftModel]: + """Build target model (backend wrapper) and draft model.""" + print_on_rank0( + f"Loading target model from {args.target_model_path} using {args.target_model_backend} backend" + ) + + target_model_kwargs = {} + if args.target_model_backend == "sglang": + target_model_kwargs = SGLangBackendArgs.from_args(args).to_kwargs() + + 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" if args.target_model_backend == "hf" else None, + trust_remote_code=args.trust_remote_code, + **target_model_kwargs, + ) + + if args.draft_config_path: + draft_config = AutoConfig.from_pretrained(args.draft_config_path) + print_on_rank0(f"Loaded draft config from {args.draft_config_path}") + # Warn if command-line args differ from config + if ( + hasattr(draft_config, "block_size") + and draft_config.block_size != args.block_size + ): + print_on_rank0( + f"Warning: checkpoint block_size ({draft_config.block_size}) differs from " + f"command-line arg ({args.block_size}). Using checkpoint value." + ) + else: + 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") + + if not hasattr(draft_config, "dflash_config") or draft_config.dflash_config is None: + draft_config.dflash_config = {} + + projector_type = draft_config.dflash_config.get("projector_type", None) + if projector_type != "domino": + raise ValueError( + "Domino training requires dflash_config.projector_type='domino'." + ) + + required_fields = [ + "emb_dim", + "gru_hidden_dim", + "pure_draft_prefix_len", + "shift_label", + ] + missing_fields = [ + field for field in required_fields if field not in draft_config.dflash_config + ] + if missing_fields: + raise ValueError( + f"Domino config missing dflash_config fields: {missing_fields}" + ) + if not hasattr(draft_config, "vocab_size"): + raise ValueError("Domino config missing draft config field: vocab_size") + + pure_prefix = draft_config.dflash_config["pure_draft_prefix_len"] + print_on_rank0( + f"Using Domino projector: pure_prefix={pure_prefix}, " + f"emb_dim={draft_config.dflash_config['emb_dim']}, " + f"gru_hidden_dim={draft_config.dflash_config['gru_hidden_dim']}" + ) + + 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) + + target_model.set_capture_layers(draft_model.target_layer_ids) + + print_on_rank0( + f"Draft config: block_size={draft_config.block_size}, " + f"num_hidden_layers={draft_config.num_hidden_layers}, " + f"num_target_layers={draft_config.num_target_layers}" + ) + print_on_rank0( + f"Draft model parameters: {sum(p.numel() for p in draft_model.parameters()):,}" + ) + + return target_model, draft_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}-" + f"{args.max_length}-" + f"{args.chat_template}-" + f"{args.target_model_path}" + ) + cache_key = hashlib.md5(cache_params_string.encode()).hexdigest() + + train_dataset = load_dataset("json", data_files=args.train_data_path)["train"] + train_eagle3_dataset = build_eagle3_dataset( + 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, + ) + + 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: + 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, domino_model, draft_model, optimizer): + """Save checkpoint.""" + 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() + + with FSDP.state_dict_type(domino_model, StateDictType.FULL_STATE_DICT): + state_dict = domino_model.state_dict() + draft_state_dict = { + k.replace("draft_model.", ""): v + for k, v in state_dict.items() + if "draft_model." in k + } + + if dist.get_rank() == 0: + torch.save( + { + "epoch": epoch, + "global_step": step, + "args": args, + **optimizer.state_dict(), + }, + os.path.join(save_dir, "training_state.pt"), + ) + + draft_model.save_pretrained(save_dir, state_dict=draft_state_dict) + + modeling_src = os.path.join( + os.path.dirname(__file__), + "..", + "specforge", + "modeling", + "draft", + "dflash.py", + ) + modeling_dst = os.path.join(save_dir, "dflash.py") + if os.path.exists(modeling_src): + shutil.copy(modeling_src, modeling_dst) + + print_on_rank0(f"Saved checkpoint to {save_dir}") + + dist.barrier() + + +def reduce_metrics_dict(metrics): + if not metrics: + return {} + + world_size = dist.get_world_size() + reduced = {} + + for k, v in metrics.items(): + if v is None: + continue + + if torch.is_tensor(v): + t = v.detach().clone() + dist.all_reduce(t) + reduced[k] = (t / world_size).item() + else: + reduced[k] = float(v) + + return reduced + + +def record_metrics( + args, + loss: float, + accuracy: float, + global_step: int, + tracker, + optimizer, + train_dataloader=None, + mode: str = "train", + extra_metrics: dict = None, +) -> None: + logdict = {} + + if mode == "train" and optimizer is not None: + logdict[f"{mode}/lr"] = optimizer.get_learning_rate() + + logdict[f"{mode}/loss"] = loss + logdict[f"{mode}/accuracy"] = accuracy + + if extra_metrics: + logdict.update( + {f"{mode}/{k}": float(v) for k, v in extra_metrics.items() if v is not None} + ) + + print_msg = ( + f"{mode.capitalize()} - Step {global_step} " + f"[{global_step}/{args.num_epochs * len(train_dataloader) // args.accumulation_steps}?], " + f"Loss: {loss:.4f}, Acc: {accuracy:.4f}" + ) + + if extra_metrics is not None: + if "base_loss" in extra_metrics: + print_msg += f", BaseLoss: {extra_metrics['base_loss']:.4f}" + if "base_accuracy" in extra_metrics: + print_msg += f", BaseAcc: {extra_metrics['base_accuracy']:.4f}" + + print_on_rank0(print_msg) + tracker.log(logdict, step=global_step) + + +def get_lambda_base( + global_step: int, + total_steps: int, + lambda_start: float = 1.0, + decay_ratio: float = 0.5, +) -> float: + decay_steps = max(1, int(total_steps * decay_ratio)) + progress = min(global_step / decay_steps, 1.0) + lambda_base = lambda_start * (1.0 - progress) + + # Clamp to [0, 1]. + lambda_base = max(0.0, min(1.0, lambda_base)) + return lambda_base + + +def main(): + + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logging.getLogger().setLevel(logging.INFO) + warnings.filterwarnings( + "ignore", + "The .grad attribute of a Tensor that is not a leaf Tensor is being accessed", + ) + + args = parse_args() + set_seed(args.seed) + + init_distributed(timeout=args.dist_timeout, tp_size=args.tp_size) + print_with_rank("Initialized distributed") + + draft_model_last_checkpoint = None + ckpt_info = (0, 0) + if args.resume and os.path.isdir(args.output_dir): + draft_model_last_checkpoint, ckpt_info = get_last_checkpoint(args.output_dir) + print(f"Last checkpoint detected: {draft_model_last_checkpoint}") + + # If resuming, load config from checkpoint to ensure consistency + if draft_model_last_checkpoint: + checkpoint_config_path = os.path.join( + draft_model_last_checkpoint, "config.json" + ) + if os.path.exists(checkpoint_config_path): + print(f"Loading draft config from checkpoint: {checkpoint_config_path}") + args.draft_config_path = checkpoint_config_path + + target_model, draft_model = build_models(args) + + 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("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( + f"Will resume from epoch {resume_state['epoch']}, " + f"step {resume_state['global_step']}" + ) + + 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}") + + 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) + total_steps = args.num_epochs * steps_per_epoch + print_on_rank0(f"Total training steps: {total_steps}") + + print_on_rank0("Loading target embeddings and head...") + target_components = TargetEmbeddingsAndHead.from_pretrained( + args.target_model_path, + embed_key=args.embedding_key, + lm_head_key=args.lm_head_key, + device="cuda", + trust_remote_code=args.trust_remote_code, + ) + + domino_model = OnlineDominoModel( + draft_model=draft_model, + target_lm_head=target_components.lm_head, + target_embed_tokens=target_components.embed_tokens, + block_size=draft_model.block_size, + mask_token_id=mask_token_id, + attention_backend=args.attention_backend, + num_anchors=args.num_anchors, + loss_decay_gamma=args.loss_decay_gamma, + shift_label=draft_model.shift_label, + ) + + domino_model = FSDP( + domino_model, + use_orig_params=True, + mixed_precision=MixedPrecision( + param_dtype=torch.bfloat16, + buffer_dtype=torch.bfloat16, + ), + sharding_strategy=ShardingStrategy.SHARD_GRAD_OP, + ) + print_with_rank("Initialized FSDP") + + start_epoch = ckpt_info[0] + global_step = ckpt_info[1] + + optimizer = BF16Optimizer( + draft_model, + lr=args.learning_rate, + max_grad_norm=args.max_grad_norm, + warmup_ratio=args.warmup_ratio, + total_steps=total_steps, + ) + + 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 optimizer/scheduler state: " + f"epoch={start_epoch}, step={global_step}, " + f"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.") + + last_time = time.time() + print_on_rank0(f"Starting training from epoch {start_epoch}, step {global_step}") + + 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"Training Epoch {epoch}", leave=True + ) + else: + progress_bar = train_dataloader + + 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() + target_output = target_model.generate_dflash_data( + input_ids, attention_mask, loss_mask + ) + hidden_states = target_output.hidden_states.cuda() # Ensure on GPU + lambda_base = get_lambda_base( + global_step=global_step, + total_steps=total_steps, + lambda_start=args.lambda_base_start, + decay_ratio=args.lambda_base_decay_ratio, + ) + + loss, accuracy, metrics = domino_model( + input_ids=input_ids, + hidden_states=hidden_states, + loss_mask=loss_mask, + lambda_base=lambda_base, + ) + + (loss / args.accumulation_steps).backward() + + if global_step % args.accumulation_steps == 0: + optimizer.step() + + if global_step % args.log_interval == 0: + loss_log = loss.clone() + acc_log = accuracy.clone() + dist.all_reduce(loss_log) + dist.all_reduce(acc_log) + loss_log = loss_log / dist.get_world_size() + acc_log = acc_log / dist.get_world_size() + metrics = reduce_metrics_dict(metrics) + + record_metrics( + args, + loss_log.item(), + acc_log.item(), + global_step, + tracker, + optimizer, + train_dataloader, + mode="train", + extra_metrics=metrics, + ) + + 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, domino_model, draft_model, optimizer + ) + + save_checkpoint( + args, args.num_epochs, global_step, domino_model, draft_model, optimizer + ) + + tracker.close() + destroy_distributed() + + +if __name__ == "__main__": + main() diff --git a/specforge/core/__init__.py b/specforge/core/__init__.py index 1b45f4f7a..8a18642ff 100644 --- a/specforge/core/__init__.py +++ b/specforge/core/__init__.py @@ -1,8 +1,10 @@ from .dflash import OnlineDFlashModel +from .domino import OnlineDominoModel from .eagle3 import OnlineEagle3Model, QwenVLOnlineEagle3Model __all__ = [ "OnlineDFlashModel", + "OnlineDominoModel", "OnlineEagle3Model", "QwenVLOnlineEagle3Model", ] diff --git a/specforge/core/domino.py b/specforge/core/domino.py new file mode 100644 index 000000000..467a8fbef --- /dev/null +++ b/specforge/core/domino.py @@ -0,0 +1,433 @@ +# coding=utf-8 +"""Domino Training Wrapper.""" + +from typing import Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from specforge.core.dflash import ( + BlockMask, + create_dflash_block_mask, + create_dflash_sdpa_mask, +) +from specforge.modeling.draft.dflash import DFlashDraftModel + + +def compute_accept_len( + pred_ids_4d: torch.Tensor, + target_ids_4d: torch.Tensor, + valid_mask_4d: torch.Tensor, +) -> torch.Tensor: + """Compute per-block acceptance length. + + For each block, returns the number of consecutive correct predictions + starting from position 0 (or the first valid position). + """ + correct = (pred_ids_4d == target_ids_4d) | (~valid_mask_4d) + accept_prefix = correct.long().cumprod(dim=2) * valid_mask_4d.long() + return accept_prefix.sum(dim=2).float() + + +class OnlineDominoModel(nn.Module): + """Domino online training wrapper with block-wise CE loss.""" + + def __init__( + self, + draft_model: DFlashDraftModel, + target_lm_head: nn.Module, + target_embed_tokens: nn.Module, + mask_token_id: int, + block_size: int = 16, + attention_backend: str = "flex_attention", + num_anchors: int = 512, + loss_decay_gamma: Optional[float] = None, + shift_label: bool = False, + ): + super().__init__() + self.draft_model = draft_model + self.lm_head = target_lm_head + self.embed_tokens = target_embed_tokens + self.block_size = block_size + self.mask_token_id = mask_token_id + self.attention_backend = attention_backend + self.num_anchors = num_anchors + self.loss_decay_gamma = loss_decay_gamma + self.shift_label = shift_label + + self._cached_block_mask: Optional[BlockMask] = None + self._cached_seq_len: Optional[int] = None + self._cached_bsz: Optional[int] = None + + def _sample_anchor_positions( + self, seq_len: int, loss_mask: torch.Tensor, device: torch.device + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Randomly sample anchor positions per sample; returns (anchors, keep_mask).""" + bs = self.block_size + bsz = loss_mask.shape[0] + max_anchor = max(seq_len - bs, 0) + + valid = loss_mask[:, : max_anchor + 1] > 0.5 + valid_counts = valid.sum(dim=1) + max_n = max(1, min(self.num_anchors, int(valid_counts.max().item()) - 1)) + + indices = ( + torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1) + ) + masked_indices = torch.where( + valid, indices, torch.tensor(seq_len + 1, device=device) + ) + + random_vals = torch.rand(bsz, max_anchor + 1, device=device) + random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device)) + + _, sorted_idx = random_vals.sort(dim=1) + gathered = torch.gather(masked_indices, 1, sorted_idx) + anchors = gathered[:, :max_n].sort(dim=1).values + + keep_mask = torch.arange(max_n, device=device).unsqueeze( + 0 + ) < valid_counts.unsqueeze(1).clamp(max=max_n) + anchors = torch.where( + keep_mask, anchors, torch.tensor(0, dtype=torch.long, device=device) + ) + + return anchors, keep_mask + + def prepare_noise_input( + self, input_ids: torch.Tensor, block_ids: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """Prepare noise input: first token of each block is real, rest are MASK.""" + bsz, seq_len = input_ids.shape + device = input_ids.device + + if block_ids is not None: + is_block_start = torch.ones(bsz, seq_len, dtype=torch.bool, device=device) + is_block_start[:, 1:] = block_ids[:, 1:] != block_ids[:, :-1] + else: + positions = torch.arange(seq_len, device=device) + is_block_start = (positions % self.block_size) == 0 + is_block_start = is_block_start.unsqueeze(0).expand(bsz, -1) + + noise_input_ids = torch.full_like(input_ids, self.mask_token_id) + noise_input_ids[is_block_start] = input_ids[is_block_start] + return noise_input_ids + + def _create_position_ids(self, anchor_positions: torch.Tensor) -> torch.Tensor: + """Create absolute position IDs for parallel draft blocks.""" + bsz, n_blocks = anchor_positions.shape + device = anchor_positions.device + offsets = torch.arange(self.block_size, device=device).view(1, 1, -1) + pos_ids = anchor_positions.unsqueeze(-1) + offsets + return pos_ids.view(bsz, -1) + + def _create_noise_embed(self, input_ids, anchor_positions, block_keep_mask): + bsz, seq_len = input_ids.shape + n = anchor_positions.shape[1] + bs = self.block_size + device = input_ids.device + + noise_ids = torch.full( + (bsz, n * bs), self.mask_token_id, dtype=torch.long, device=device + ) + + block_starts = torch.arange(n, device=device) * bs + block_starts = block_starts.unsqueeze(0).expand(bsz, -1) + + valid_anchor_positions = anchor_positions.clamp(0, seq_len - 1) + anchor_tokens = torch.gather(input_ids, 1, valid_anchor_positions) + + flat_batch_idx = torch.arange(bsz, device=device).unsqueeze(1).expand(bsz, n) + noise_ids[flat_batch_idx, block_starts] = torch.where( + block_keep_mask, + anchor_tokens, + torch.tensor(self.mask_token_id, dtype=torch.long, device=device), + ) + + return self.embed_tokens(noise_ids) + + @property + def _suffix_start(self) -> int: + """Return suffix_start index based on shift_label and pure_draft_prefix_len.""" + pure_prefix = getattr(self.draft_model, "pure_draft_prefix_len", 0) + return pure_prefix if self.shift_label else (1 + pure_prefix) + + def _build_domino_head_inputs( + self, + input_ids: torch.Tensor, + anchor_positions: torch.Tensor, + target_ids: torch.Tensor, + output_hidden: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + bsz, n, bs = target_ids.shape + hidden4d = output_hidden.reshape(bsz, n, bs, output_hidden.shape[-1]) + + prev_ids = target_ids + if self.shift_label: + prev_offsets = torch.arange( + 0, self.block_size, device=input_ids.device + ).view(1, 1, -1) + prev_indices = (anchor_positions.unsqueeze(-1) + prev_offsets).clamp( + max=input_ids.size(1) - 1 + ) + prev_ids = torch.gather( + input_ids.unsqueeze(1).expand(-1, anchor_positions.size(1), -1), + 2, + prev_indices, + ) + + return hidden4d, prev_ids + + def _apply_domino_head( + self, + base_logits4d: torch.Tensor, + hidden4d: torch.Tensor, + prev_ids: torch.Tensor, + target_ids: torch.Tensor, + ) -> torch.Tensor: + """Apply the Domino head: GRU causal state plus logit correction.""" + bsz, n, bs = target_ids.shape + if self.shift_label: + block_emb = self.embed_tokens(prev_ids) + gru_inputs = block_emb.reshape(bsz * n, bs, -1) + gru_out, _ = self.draft_model.prefix_gru(gru_inputs) + gru_out = gru_out.reshape(bsz, n, bs, -1) + prefix_states = gru_out[:, :, self._suffix_start :, :] + else: + block_emb = self.embed_tokens(target_ids) + gru_inputs = block_emb[:, :, : bs - 1, :].reshape(bsz * n, bs - 1, -1) + gru_out, _ = self.draft_model.prefix_gru(gru_inputs) + gru_out = gru_out.reshape(bsz, n, bs - 1, -1) + prefix_states = gru_out[:, :, self._suffix_start - 1 :, :] + z_n = hidden4d[:, :, self._suffix_start :, :] + concat_features = torch.cat([z_n, prefix_states], dim=-1) + logits_e = self.draft_model.embed_proj(concat_features) + + prefix_logits = base_logits4d[:, :, : self._suffix_start, :] + suffix_logits = base_logits4d[:, :, self._suffix_start :, :] + logits_e + return torch.cat([prefix_logits, suffix_logits], dim=2) + + def _compute_extra_metrics( + self, + pred_ids: torch.Tensor, + flat_base_logits: torch.Tensor, + flat_targets: torch.Tensor, + binary_eval_mask: torch.Tensor, + actual_token_count: torch.Tensor, + target_ids: torch.Tensor, + eval_weight_mask: torch.Tensor, + final_loss: torch.Tensor, + base_loss: torch.Tensor, + lambda_base: float, + ) -> Dict[str, torch.Tensor]: + """Compute auxiliary training metrics that do not affect gradients.""" + bsz, n, bs = target_ids.shape + + base_pred_ids = torch.argmax(flat_base_logits, dim=-1) + base_correct = (base_pred_ids == flat_targets) & (binary_eval_mask > 0.5) + base_accuracy = base_correct.sum().float() / actual_token_count + + valid_mask_4d = (eval_weight_mask > 0).bool() + pred_accept_len = compute_accept_len( + pred_ids.view(bsz, n, bs), target_ids, valid_mask_4d + ) + base_accept_len = compute_accept_len( + base_pred_ids.view(bsz, n, bs), target_ids, valid_mask_4d + ) + + valid_block_mask = valid_mask_4d.any(dim=2) + num_valid_blocks = valid_block_mask.sum().float() + 1e-6 + avg_accept_len = ( + (pred_accept_len + 1.0) * valid_block_mask.float() + ).sum() / num_valid_blocks + base_avg_accept_len = ( + (base_accept_len + 1.0) * valid_block_mask.float() + ).sum() / num_valid_blocks + + return { + "final_loss": final_loss.detach(), + "base_loss": base_loss.detach(), + "base_accuracy": base_accuracy.detach(), + "accept_len": avg_accept_len.detach(), + "base_accept_len": base_avg_accept_len.detach(), + "lambda_base": torch.tensor(lambda_base, device=final_loss.device), + } + + def _compute_weighted_losses( + self, + final_logits: torch.Tensor, + base_logits: torch.Tensor, + target_ids: torch.Tensor, + weight_mask: torch.Tensor, + lambda_base: float, + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + flat_logits = final_logits.reshape(-1, final_logits.size(-1)) + flat_base_logits = base_logits.reshape(-1, base_logits.size(-1)) + flat_targets = target_ids.reshape(-1) + flat_weights = weight_mask.reshape(-1) + + valid_token_count = flat_weights.sum() + 1e-6 + + final_loss_per_token = F.cross_entropy( + flat_logits, flat_targets, reduction="none" + ) + final_loss = (final_loss_per_token * flat_weights).sum() / valid_token_count + + base_loss_per_token = F.cross_entropy( + flat_base_logits, flat_targets, reduction="none" + ) + base_loss = (base_loss_per_token * flat_weights).sum() / valid_token_count + + loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss + + return loss, final_loss, base_loss, flat_logits, flat_base_logits, flat_targets + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + loss_mask: torch.Tensor, + lambda_base: float = 0.0, + ): + """Parallel block-wise training forward pass.""" + bsz, seq_len = input_ids.shape + device = input_ids.device + + anchor_positions, block_keep_mask = self._sample_anchor_positions( + seq_len, loss_mask, device + ) + + noise_embedding = self._create_noise_embed( + input_ids, anchor_positions, block_keep_mask + ) + + context_position_ids = ( + torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) + ) + draft_position_ids = self._create_position_ids(anchor_positions) + full_position_ids = torch.cat([context_position_ids, draft_position_ids], dim=1) + + if self.attention_backend == "flex_attention": + dflash_attn_mask = create_dflash_block_mask( + anchor_positions=anchor_positions, + block_keep_mask=block_keep_mask, + S=seq_len, + block_size=self.block_size, + device=device, + ) + else: + dflash_attn_mask = create_dflash_sdpa_mask( + anchor_positions=anchor_positions, + block_keep_mask=block_keep_mask, + S=seq_len, + block_size=self.block_size, + device=device, + ) + + output_hidden = self.draft_model( + position_ids=full_position_ids, + noise_embedding=noise_embedding, + target_hidden=hidden_states, + attention_mask=dflash_attn_mask, + ) + + # --- Labels --- + label_start = 1 if self.shift_label else 0 + label_offsets = torch.arange( + label_start, label_start + self.block_size, device=device + ).view(1, 1, -1) + label_indices = anchor_positions.unsqueeze(-1) + label_offsets + valid_label_mask = label_indices < seq_len + safe_target_indices = label_indices.clamp(max=seq_len - 1) + + target_ids = torch.gather( + input_ids.unsqueeze(1).expand(-1, anchor_positions.size(1), -1), + 2, + safe_target_indices, + ) + + bsz, n, bs = target_ids.shape + base_logits = self.lm_head(output_hidden) + hidden4d, prev_ids = self._build_domino_head_inputs( + input_ids=input_ids, + anchor_positions=anchor_positions, + target_ids=target_ids, + output_hidden=output_hidden, + ) + base_logits4d = base_logits.reshape(bsz, n, bs, -1) + final_logits = self._apply_domino_head( + base_logits4d=base_logits4d, + hidden4d=hidden4d, + prev_ids=prev_ids, + target_ids=target_ids, + ).reshape(bsz, n * bs, -1) + + # --- Weight mask: block validity * bounds * exclude anchor (pos 0) * loss_mask --- + weight_mask = ( + block_keep_mask.unsqueeze(-1).expand(-1, -1, self.block_size).float() + ) + weight_mask = weight_mask * valid_label_mask.float() + + if not self.shift_label: + pos_in_block = torch.arange(self.block_size, device=device).view(1, 1, -1) + weight_mask = weight_mask * (pos_in_block > 0).float() + + original_loss_mask_gathered = torch.gather( + loss_mask.unsqueeze(1).expand(-1, anchor_positions.size(1), -1), + 2, + safe_target_indices, + ) + weight_mask = weight_mask * original_loss_mask_gathered + + # Save eval mask before decay (for accept_len / accuracy stats) + eval_weight_mask = weight_mask.clone() + binary_eval_mask = weight_mask.view(-1) + + # --- Loss decay: first valid position gets weight 1.0 --- + if self.loss_decay_gamma is not None and self.loss_decay_gamma > 0: + k = torch.arange(self.block_size, device=device).view(1, 1, -1) + offset = 0 if self.shift_label else 1 + decay_weights = torch.exp( + -(k - offset).clamp(min=0).float() / self.loss_decay_gamma + ) + weight_mask = weight_mask * decay_weights + + loss, final_loss, base_loss, flat_logits, flat_base_logits, flat_targets = ( + self._compute_weighted_losses( + final_logits=final_logits, + base_logits=base_logits, + target_ids=target_ids, + weight_mask=weight_mask, + lambda_base=lambda_base, + ) + ) + + # --- Accuracy --- + with torch.no_grad(): + pred_ids = torch.argmax(flat_logits, dim=-1) + correct = (pred_ids == flat_targets) & (binary_eval_mask > 0.5) + actual_token_count = binary_eval_mask.sum() + 1e-6 + accuracy = correct.sum().float() / actual_token_count + + metrics = self._compute_extra_metrics( + pred_ids=pred_ids, + flat_base_logits=flat_base_logits, + flat_targets=flat_targets, + binary_eval_mask=binary_eval_mask, + actual_token_count=actual_token_count, + target_ids=target_ids, + eval_weight_mask=eval_weight_mask, + final_loss=final_loss, + base_loss=base_loss, + lambda_base=lambda_base, + ) + + return loss, accuracy, metrics diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 0aea03fe1..0f750920e 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -237,6 +237,28 @@ def __init__(self, config) -> None: self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.block_size = config.block_size self.mask_token_id = dflash_config.get("mask_token_id", None) + self.projector_type = dflash_config.get("projector_type", None) + self.pure_draft_prefix_len = dflash_config.get("pure_draft_prefix_len", 0) + self.shift_label = dflash_config.get("shift_label", False) + + if self.projector_type == "domino": + self.emb_dim = dflash_config["emb_dim"] + self.gru_hidden_dim = dflash_config["gru_hidden_dim"] + self.prefix_gru = nn.GRU( + input_size=config.hidden_size, + hidden_size=self.gru_hidden_dim, + num_layers=1, + batch_first=True, + bias=False, + ) + in_dim = config.hidden_size + self.gru_hidden_dim + self.embed_proj = nn.Sequential( + nn.Linear(in_dim, self.emb_dim, bias=False), + nn.SiLU(), + nn.Linear(self.emb_dim, config.vocab_size, bias=False), + ) + elif self.projector_type is not None: + raise ValueError(f"Unknown draft projector_type: {self.projector_type}") self.post_init() def forward(