diff --git a/configs/qwen3-8b-peagle.json b/configs/qwen3-8b-peagle.json new file mode 100644 index 000000000..b3a11dd56 --- /dev/null +++ b/configs/qwen3-8b-peagle.json @@ -0,0 +1,31 @@ +{ + "architectures": [ + "PEagleDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "max_position_embeddings": 40960, + "max_window_layers": 36, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 4, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936, + "draft_vocab_size": 32000 +} diff --git a/examples/run_qwen3_8b_peagle_online.sh b/examples/run_qwen3_8b_peagle_online.sh new file mode 100755 index 000000000..63f8e00a9 --- /dev/null +++ b/examples/run_qwen3_8b_peagle_online.sh @@ -0,0 +1,42 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +ROOT_DIR=$(dirname $SCRIPT_DIR) +# peagle.py is not in the installed specforge package yet; prefer the repo source +export PYTHONPATH=$ROOT_DIR:${PYTHONPATH:-} +export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels + +# support tp8 train P-EAGLE for Qwen3-4B/8B/32B up to tp_size = 8 +NUM_GPUS=${1:-1} +TP_SIZE=${2:-1} +BUILD_DATASET_NUM_PROC=${BUILD_DATASET_NUM_PROC:-32} + +torchrun \ + --standalone \ + --nproc_per_node $NUM_GPUS \ + $ROOT_DIR/scripts/train_peagle.py \ + --target-model-path Qwen/Qwen3-8B \ + --draft-model-config $ROOT_DIR/configs/qwen3-8b-peagle.json \ + --train-data-path $ROOT_DIR/cache/dataset/perfectblend-qwen3-8b-regen.jsonl \ + --build-dataset-num-proc $BUILD_DATASET_NUM_PROC \ + --output-dir $ROOT_DIR/outputs/peagle_qwen3_8b \ + --num-epochs 20 \ + --batch-size 1 \ + --learning-rate 1e-4 \ + --max-length 4096 \ + --warmup-ratio 0.0025 \ + --max-grad-norm 1 \ + --chat-template qwen \ + --cache-dir $ROOT_DIR/cache \ + --tp-size $TP_SIZE \ + --num-depths 5 \ + --down-sample-ratio 0.8 \ + --down-sample-ratio-min 0.2 \ + --num-draft-layers 4 \ + --no-norm-before-residual \ + --target-model-backend sglang \ + --save-interval 50000 \ + --eval-interval 50000 \ + --log-interval 50 \ + --report-to wandb \ + --dist-timeout 120 diff --git a/scripts/train_peagle.py b/scripts/train_peagle.py new file mode 100644 index 000000000..5fde2e126 --- /dev/null +++ b/scripts/train_peagle.py @@ -0,0 +1,821 @@ +"""P-EAGLE (Parallel EAGLE) training script. + +Based on train_eagle3.py but replaces TTT with COD parallel sampling. +""" + +import argparse +import hashlib +import json +import math +import os +import time +from argparse import ArgumentParser, Namespace +from typing import Dict, Optional, Tuple + +import torch +import torch.distributed as dist +import torch.nn as nn +from accelerate.utils import set_seed +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType +from torch.optim import Optimizer +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoTokenizer + +from datasets import DatasetDict, load_dataset +from specforge import AutoDraftModelConfig, get_eagle3_target_model +from specforge.args import SGLangBackendArgs, TrackerArgs +from specforge.core.peagle import OnlinePEagleModel +from specforge.data import ( + build_eagle3_dataset, + generate_vocab_mapping_file, + prepare_dp_dataloaders, +) +from specforge.distributed import ( + destroy_distributed, + get_dp_group, + get_tp_group, + init_distributed, +) +from specforge.modeling.draft.peagle import PEagleDraftModel +from specforge.modeling.target import Eagle3TargetModel +from specforge.optimizer import BF16Optimizer +from specforge.tracker import Tracker, create_tracker, get_tracker_class +from specforge.utils import ( + get_last_checkpoint, + print_args_with_dots, + print_on_rank0, + print_with_rank, + rank_0_priority, +) + + +def parse_args() -> Tuple[ArgumentParser, Namespace]: + parser = argparse.ArgumentParser(description="Train P-EAGLE with online data") + + model_group = parser.add_argument_group("model") + model_group.add_argument("--target-model-path", type=str, required=True) + model_group.add_argument( + "--trust-remote-code", action="store_true", help="Trust remote code" + ) + model_group.add_argument("--draft-model-config", type=str, required=False) + model_group.add_argument( + "--embedding-key", + type=str, + default="model.embed_tokens.weight", + ) + model_group.add_argument( + "--target-model-backend", + type=str, + default="sglang", + choices=["sglang", "hf", "custom"], + ) + + # P-EAGLE specific args + peagle_group = parser.add_argument_group("peagle") + peagle_group.add_argument( + "--num-depths", + type=int, + default=8, + help="Number of parallel prediction depths for P-EAGLE COD sampling", + ) + peagle_group.add_argument( + "--down-sample-ratio", + type=float, + default=0.8, + help="Geometric decay ratio for COD sampling", + ) + peagle_group.add_argument( + "--down-sample-ratio-min", + type=float, + default=0.2, + help="Minimum retention ratio for COD sampling", + ) + peagle_group.add_argument( + "--mask-token-id", + type=int, + default=None, + help="Token ID for masking. If None, uses tokenizer.pad_token_id or 0", + ) + peagle_group.add_argument( + "--num-draft-layers", + type=int, + default=4, + help="Number of decoder layers in the P-EAGLE draft model", + ) + peagle_group.add_argument( + "--norm-before-residual", + action="store_true", + help="Whether to use normalized hidden as residual in the first layer", + ) + peagle_group.add_argument( + "--no-norm-before-residual", + action="store_true", + help="Explicitly disable norm-before-residual", + ) + + 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="llama3") + dataset_group.add_argument("--is-preformatted", action="store_true") + dataset_group.add_argument("--train-only-last-turn", action="store_true") + dataset_group.add_argument("--build-dataset-num-proc", type=int, default=8) + dataset_group.add_argument("--dataloader-num-workers", type=int, default=4) + + training_group = parser.add_argument_group("training") + training_group.add_argument("--num-epochs", type=int, default=10) + training_group.add_argument("--max-num-steps", type=int, default=None) + 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=2048) + training_group.add_argument("--warmup-ratio", type=float, default=0.015) + training_group.add_argument("--total-steps", type=int, default=None) + training_group.add_argument("--max-grad-norm", type=float, default=0.5) + training_group.add_argument("--resume", action="store_true") + training_group.add_argument("--ckpt-dir", type=str, default=None) + training_group.add_argument("--eval-interval", type=int, default=5000) + training_group.add_argument("--save-interval", type=int, default=5000) + training_group.add_argument("--log-interval", type=int, default=50) + training_group.add_argument("--seed", type=int, default=0) + training_group.add_argument("--draft-accumulation-steps", type=int, default=1) + + optimization_group = parser.add_argument_group("optimization") + optimization_group.add_argument("--tp-size", type=int, default=1) + + other_group = parser.add_argument_group("others") + other_group.add_argument("--cache-key", type=str, default=None) + other_group.add_argument("--cache-dir", type=str, default="./cache") + other_group.add_argument("--output-dir", type=str, required=True) + other_group.add_argument("--verbose", action="store_true") + other_group.add_argument("--dist-timeout", type=int, default=20) + other_group.add_argument("--model-download-dir", type=str, default=None) + + profiling_group = parser.add_argument_group("profiling") + profiling_group.add_argument("--profile", action="store_true") + profiling_group.add_argument("--profile-start-step", type=int, default=30) + profiling_group.add_argument("--profile-num-steps", type=int, default=4) + profiling_group.add_argument("--profile-record-shapes", action="store_true") + + sglang_group = parser.add_argument_group("sglang target model backend") + SGLangBackendArgs.add_args(sglang_group) + + tracker_group = parser.add_argument_group("tracker") + TrackerArgs.add_args(tracker_group) + + args = parser.parse_args() + return parser, args + + +def build_tracker(args: Namespace, parser: ArgumentParser) -> Tracker: + tracker_class = get_tracker_class(args.report_to) + if tracker_class: + tracker_class.validate_args(parser, args) + else: + parser.error(f"Unknown tracker: {args.report_to}") + return create_tracker(args, args.output_dir) + + +def build_target_model( + args: Namespace, draft_model_config: AutoDraftModelConfig +) -> Eagle3TargetModel: + if args.target_model_backend == "sglang": + target_model_kwargs = SGLangBackendArgs.from_args(args).to_kwargs() + else: + target_model_kwargs = {} + target_model = get_eagle3_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.model_download_dir, + **target_model_kwargs, + trust_remote_code=args.trust_remote_code, + ) + if ( + hasattr(draft_model_config, "eagle_config") + and draft_model_config.eagle_config is not None + and "eagle_aux_hidden_state_layer_ids" in draft_model_config.eagle_config + ): + target_model.set_aux_hidden_states_layers( + draft_model_config.eagle_config["eagle_aux_hidden_state_layer_ids"] + ) + else: + target_model.set_aux_hidden_states_layers() + return target_model + + +def build_draft_model(args: Namespace) -> Tuple: + ckpt_info = (0, 0) + resume_state = None + should_load_target_embedding = True + + if args.draft_model_config is not None: + draft_model_config = AutoDraftModelConfig.from_file(args.draft_model_config) + else: + from specforge.utils import create_draft_config_from_target + + auto_config_path = create_draft_config_from_target( + target_model_path=args.target_model_path, + cache_dir=args.model_download_dir, + ) + draft_model_config = AutoDraftModelConfig.from_file(auto_config_path) + + # Override num_hidden_layers for P-EAGLE multi-layer + draft_model_config.num_hidden_layers = args.num_draft_layers + + draft_model_last_checkpoint = None + is_resume_checkpoint = False + if args.ckpt_dir is not None: + if os.path.isdir(args.ckpt_dir): + draft_model_config = AutoDraftModelConfig.from_file( + os.path.join(args.ckpt_dir, "config.json") + ) + draft_model_config.num_hidden_layers = args.num_draft_layers + draft_model_last_checkpoint = args.ckpt_dir + should_load_target_embedding = False + print_on_rank0(f"Finetuning from base model: {draft_model_last_checkpoint}") + else: + raise ValueError( + f"Provided base model dir {args.ckpt_dir} is not a valid directory." + ) + + 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}") + is_resume_checkpoint = True + should_load_target_embedding = False + + norm_before_residual = ( + args.norm_before_residual and not args.no_norm_before_residual + ) + + if draft_model_last_checkpoint: + draft_model = PEagleDraftModel( + config=draft_model_config, + norm_before_residual=norm_before_residual, + ).to(dtype=torch.bfloat16, device="cuda") + safetensors_path = os.path.join( + draft_model_last_checkpoint, "model.safetensors" + ) + if os.path.exists(safetensors_path): + from safetensors.torch import load_file + + state_dict = load_file(safetensors_path, device="cuda") + draft_model.load_state_dict(state_dict, strict=False) + if "embed_tokens.weight" not in state_dict: + should_load_target_embedding = True + print_on_rank0( + "Checkpoint does not contain trainable P-EAGLE embeddings; " + "loading embeddings from the target model." + ) + else: + should_load_target_embedding = True + print_on_rank0( + f"No model.safetensors found in {draft_model_last_checkpoint}; " + "loading embeddings from the target model." + ) + else: + draft_model = PEagleDraftModel( + config=draft_model_config, + norm_before_residual=norm_before_residual, + ).to(dtype=torch.bfloat16, device="cuda") + + if is_resume_checkpoint and draft_model_last_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"Loaded training state from {training_state_path}: " + f"epoch={resume_state['epoch']}, step={resume_state['global_step']}" + ) + + if should_load_target_embedding: + draft_model.load_embedding( + args.target_model_path, embedding_key=args.embedding_key + ) + else: + print_on_rank0("Using embeddings from the P-EAGLE checkpoint.") + return draft_model_config, draft_model, ckpt_info, resume_state + + +def load_conversation_dataset(data_path: str): + """Load local JSON/JSONL data like DFlash, or an HF dataset id.""" + if os.path.isfile(data_path) and os.path.splitext(data_path)[1].lower() in ( + ".json", + ".jsonl", + ): + return load_dataset("json", data_files=data_path)["train"] + + dataset = load_dataset(data_path, split="train") + if isinstance(dataset, DatasetDict): + if "train" not in dataset: + raise ValueError( + f"Expected a 'train' split, but found splits: {list(dataset.keys())}" + ) + return dataset["train"] + return dataset + + +def build_dataloaders( + args: Namespace, + draft_model_config, +) -> Tuple[DataLoader, str, Optional[DataLoader]]: + tokenizer = AutoTokenizer.from_pretrained( + args.target_model_path, trust_remote_code=args.trust_remote_code + ) + + draft_vocab_size = getattr( + draft_model_config, "draft_vocab_size", draft_model_config.vocab_size + ) + cache_params_string = ( + f"{args.train_data_path}-" + f"{args.max_length}-" + f"{args.chat_template}-" + f"{args.target_model_path}-" + f"{draft_vocab_size}" + ) + cache_key = hashlib.md5(cache_params_string.encode()).hexdigest() + train_dataset = load_conversation_dataset(args.train_data_path) + with rank_0_priority(): + train_eagle3_dataset = build_eagle3_dataset( + dataset=train_dataset, + tokenizer=tokenizer, + chat_template=args.chat_template, + max_length=args.max_length, + cache_dir=os.path.join(args.cache_dir, "processed_dataset"), + cache_key=cache_key, + is_vlm=False, + is_preformatted=args.is_preformatted, + processor=None, + num_proc=args.build_dataset_num_proc, + train_only_last_turn=args.train_only_last_turn, + minimum_valid_tokens=1, + ) + vocab_mapping_path = generate_vocab_mapping_file( + dataset=train_eagle3_dataset, + target_vocab_size=draft_model_config.vocab_size, + draft_vocab_size=draft_vocab_size, + cache_dir=os.path.join(args.cache_dir, "vocab_mapping"), + cache_key=cache_key, + ) + + train_dataloader = prepare_dp_dataloaders( + train_eagle3_dataset, + args.target_batch_size, + num_workers=args.dataloader_num_workers, + shuffle=True, + process_group=get_dp_group(), + is_vlm=False, + ) + + eval_dataloader = None + if args.eval_data_path is not None: + eval_dataset = load_conversation_dataset(args.eval_data_path) + eval_eagle3_dataset = build_eagle3_dataset( + eval_dataset, + tokenizer, + args.chat_template, + args.max_length, + is_vlm=False, + processor=None, + num_proc=args.build_dataset_num_proc, + is_preformatted=args.is_preformatted, + train_only_last_turn=args.train_only_last_turn, + ) + eval_dataloader = prepare_dp_dataloaders( + eval_eagle3_dataset, + args.target_batch_size, + num_workers=args.dataloader_num_workers, + shuffle=False, + process_group=get_dp_group(), + is_vlm=False, + ) + print_with_rank("Initialized eval dataloader") + + return train_dataloader, vocab_mapping_path, eval_dataloader + + +def save_checkpoints( + args: Namespace, + epoch: int, + step: int, + peagle_model: nn.Module, + optimizer: Optimizer, +): + epoch_output_dir = os.path.join(args.output_dir, f"epoch_{epoch}_step_{step}") + if dist.get_rank() == 0: + os.makedirs(epoch_output_dir, exist_ok=True) + dist.barrier() + + with FSDP.state_dict_type(peagle_model, StateDictType.FULL_STATE_DICT): + model_state_dict = peagle_model.state_dict() + state_to_save = { + "epoch": epoch, + "global_step": step, + "args": args, + } + state_to_save.update(optimizer.state_dict()) + draft_model_state_dict = { + k.replace("draft_model.", ""): v + for k, v in model_state_dict.items() + if "draft_model." in k + } + + if dist.get_rank() == 0: + torch.save( + state_to_save, + os.path.join(epoch_output_dir, "training_state.pt"), + ) + peagle_model.draft_model.save_pretrained( + epoch_output_dir, + state_dict=draft_model_state_dict, + ) + peagle_config = { + "num_depths": args.num_depths, + "down_sample_ratio": args.down_sample_ratio, + "down_sample_ratio_min": args.down_sample_ratio_min, + "mask_token_id": args.mask_token_id, + "num_draft_layers": args.num_draft_layers, + "norm_before_residual": args.norm_before_residual, + } + with open(os.path.join(epoch_output_dir, "peagle_config.json"), "w") as f: + json.dump(peagle_config, f, indent=2) + + print_on_rank0(f"Saved model to {epoch_output_dir}") + dist.barrier() + + +def get_dp_data_shard_from_tp(tensor: torch.Tensor) -> torch.Tensor: + tp_size = dist.get_world_size(get_tp_group()) + tp_rank = dist.get_rank(get_tp_group()) + return tensor.chunk(tp_size, dim=0)[tp_rank] + + +def run_forward( + args: Namespace, + peagle_model: nn.Module, + data: dict, + target_model: Eagle3TargetModel, +) -> Tuple[torch.Tensor, Dict]: + eagle3_data = target_model.generate_eagle3_data( + input_ids=data["input_ids"].cuda(), + attention_mask=data["attention_mask"].cuda(), + loss_mask=data["loss_mask"].cuda(), + ) + + input_ids = get_dp_data_shard_from_tp(eagle3_data.input_ids) + attention_mask = get_dp_data_shard_from_tp(eagle3_data.attention_mask) + loss_mask = get_dp_data_shard_from_tp(eagle3_data.loss_mask) + target = get_dp_data_shard_from_tp(eagle3_data.target) + hidden_states = get_dp_data_shard_from_tp(eagle3_data.hidden_states) + + loss, metrics = peagle_model( + input_ids=input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + target=target, + hidden_states=hidden_states, + ) + return loss, metrics + + +def record_metrics( + args: Namespace, + metrics: Dict, + global_step: int, + tracker: Tracker, + optimizer: Optional[Optimizer] = None, + mode: str = "train", +) -> None: + logdict = {} + + if mode == "train" and optimizer is not None: + logdict["train/lr"] = optimizer.get_learning_rate() + + loss = metrics.get("loss_sum", torch.tensor(0.0)) + dist.all_reduce(loss, op=dist.ReduceOp.AVG) + logdict[f"{mode}/loss"] = loss.item() + print_on_rank0(f"{mode} - Step {global_step}, Loss: {loss.item():.4f}") + + full_acc_sum = metrics.get("full_acc_sum", torch.tensor(0.0)) + full_acc_total = metrics.get("full_acc_total", torch.tensor(1.0)) + dist.all_reduce(full_acc_sum, op=dist.ReduceOp.SUM) + dist.all_reduce(full_acc_total, op=dist.ReduceOp.SUM) + full_acc = (full_acc_sum / full_acc_total.clamp_min(1)).item() + logdict[f"{mode}/acc"] = full_acc + print_on_rank0(f"{mode} - Step {global_step}, Acc: {full_acc:.4f}") + + for d in range(args.num_depths): + key_sum = f"position_{d}_acc_sum" + key_total = f"position_{d}_acc_total" + if key_sum in metrics: + d_sum = metrics[key_sum] + d_total = metrics[key_total] + dist.all_reduce(d_sum, op=dist.ReduceOp.SUM) + dist.all_reduce(d_total, op=dist.ReduceOp.SUM) + d_acc = (d_sum / d_total.clamp_min(1)).item() + logdict[f"{mode}/acc_depth_{d}"] = d_acc + print_on_rank0(f"{mode} - Step {global_step}, Depth {d} Acc: {d_acc:.4f}") + + tracker.log(logdict, step=global_step) + + +def _print_on_rank0_or_local(message: str) -> None: + if dist.is_available() and dist.is_initialized(): + print_on_rank0(message) + else: + print_with_rank(message) + + +def _validate_mask_token_id(mask_token_id: int, embedding_vocab_size: int) -> int: + if not 0 <= mask_token_id < embedding_vocab_size: + raise ValueError( + f"mask_token_id {mask_token_id} is outside embedding vocab " + f"size {embedding_vocab_size}." + ) + return mask_token_id + + +def resolve_mask_token_id(args: Namespace, embedding_vocab_size: int) -> int: + if args.mask_token_id is not None: + return _validate_mask_token_id(args.mask_token_id, embedding_vocab_size) + + tokenizer = AutoTokenizer.from_pretrained( + args.target_model_path, trust_remote_code=args.trust_remote_code + ) + if getattr(tokenizer, "mask_token_id", None) is not None: + mask_token_id = _validate_mask_token_id( + tokenizer.mask_token_id, embedding_vocab_size + ) + _print_on_rank0_or_local( + f"Auto-set mask_token_id to tokenizer mask token {mask_token_id}" + ) + return mask_token_id + + if len(tokenizer) < embedding_vocab_size: + mask_token_id = len(tokenizer) + _print_on_rank0_or_local( + f"Auto-set mask_token_id to unused embedding slot {mask_token_id}" + ) + return mask_token_id + + for token_name in ("pad_token_id", "eos_token_id", "unk_token_id"): + token_id = getattr(tokenizer, token_name, None) + if token_id is not None: + mask_token_id = _validate_mask_token_id(token_id, embedding_vocab_size) + _print_on_rank0_or_local( + "Tokenizer has no mask token or unused draft embedding slot; " + f"falling back to {token_name}={mask_token_id}. " + "Pass --mask-token-id to use a dedicated trainable mask token." + ) + return mask_token_id + + raise ValueError( + "Could not resolve mask_token_id. Pass --mask-token-id or use a tokenizer " + "with mask/pad/eos/unk token." + ) + + +def main(): + # ================================================ + # 1. Initialize + # ================================================ + parser, args = parse_args() + set_seed(args.seed) + init_distributed(timeout=args.dist_timeout, tp_size=args.tp_size) + + args.dp_size = dist.get_world_size() // args.tp_size + args.target_batch_size = args.tp_size * args.batch_size + + print_args_with_dots(args) + print_with_rank("Initialized distributed environment") + + # ================================================ + # 2. Build models + # ================================================ + draft_model_config, draft_model, ckpt_info, resume_state = build_draft_model(args) + target_model = build_target_model(args, draft_model_config) + + # ================================================ + # 3. Build dataloader + # ================================================ + train_dataloader, vocab_mapping_path, eval_dataloader = build_dataloaders( + args, draft_model_config + ) + draft_model.load_vocab_mapping(vocab_mapping_path) + print_with_rank("Loaded vocab mapping") + + # Resolve mask_token_id + args.mask_token_id = resolve_mask_token_id( + args, + draft_model_config.vocab_size, + ) + + # Calculate total steps + if args.total_steps is None: + steps_per_epoch = math.ceil( + len(train_dataloader) / args.draft_accumulation_steps + ) + args.total_steps = args.num_epochs * steps_per_epoch + print_with_rank(f"Auto-calculated total_steps: {args.total_steps}") + + # ================================================ + # 4. Build P-EAGLE model + # ================================================ + peagle_model = OnlinePEagleModel( + draft_model=draft_model, + mask_token_id=args.mask_token_id, + num_depths=args.num_depths, + down_sample_ratio=args.down_sample_ratio, + down_sample_ratio_min=args.down_sample_ratio_min, + ) + + # ================================================ + # 5. Wrap with FSDP, then build optimizer and scheduler + # ================================================ + peagle_model = FSDP( + peagle_model, + use_orig_params=True, + mixed_precision=MixedPrecision( + param_dtype=torch.bfloat16, + buffer_dtype=torch.bfloat16, + ), + sharding_strategy=ShardingStrategy.SHARD_GRAD_OP, + process_group=dist.group.WORLD, + device_id=torch.cuda.current_device(), + ) + + # Build optimizer after FSDP so fp32 param copies match sharded shapes + optimizer = BF16Optimizer( + peagle_model, + lr=args.learning_rate, + max_grad_norm=args.max_grad_norm, + warmup_ratio=args.warmup_ratio, + total_steps=args.total_steps, + ) + print_with_rank("Initialized optimizer and scheduler") + + if resume_state is not None: + optimizer.load_state_dict(resume_state) + start_epoch = resume_state["epoch"] + global_step = resume_state["global_step"] + print_on_rank0( + f"Restored optimizer/scheduler state: " + f"epoch={start_epoch}, step={global_step}, " + f"lr={optimizer.get_learning_rate():.6f}" + ) + del resume_state + else: + start_epoch = ckpt_info[0] + global_step = ckpt_info[1] + + skip_steps = global_step - start_epoch * len(train_dataloader) + + # ================================================ + # 6. Build tracker + # ================================================ + tracker = build_tracker(args, parser) + dist.barrier() + + last_time = time.time() + + # ================================================ + # 7. Start training + # ================================================ + print_on_rank0( + f"Starting P-EAGLE training from epoch:{start_epoch} step:{global_step}" + ) + + for epoch in range(start_epoch, args.num_epochs): + train_dataloader.sampler.set_epoch(epoch + 1) + 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 + + # Profiling + if args.profile: + if global_step == args.profile_start_step + 1: + print("Start profile") + torch_profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, + record_shapes=args.profile_record_shapes, + ) + torch_profiler.start() + if global_step == args.profile_start_step + args.profile_num_steps + 1: + output_path = os.path.join( + args.output_dir, + f"profile_rank{dist.get_rank()}_{time.time()}.trace.json.gz", + ) + print(f"End profile {output_path=}") + torch_profiler.stop() + torch_profiler.export_chrome_trace(output_path) + + # Training Step + loss, metrics = run_forward(args, peagle_model, data, target_model) + scaled_loss = loss / args.draft_accumulation_steps + scaled_loss.backward() + + if global_step % args.draft_accumulation_steps == 0: + optimizer.step() + + # Logging + if global_step % (args.log_interval * args.draft_accumulation_steps) == 0: + record_metrics( + args, + metrics, + global_step // args.draft_accumulation_steps, + tracker, + optimizer, + mode="train", + ) + + if dist.get_rank() == 0: + time_per_step = time.time() - last_time + last_time = time.time() + acc = metrics.get("full_acc_sum", torch.tensor(0.0)) + acc_total = metrics.get("full_acc_total", torch.tensor(1.0)) + progress_bar.set_postfix( + { + "loss": f"{loss.item():.4f}", + "acc": f"{(acc / acc_total.clamp_min(1)).item():.4f}", + "time": f"{time_per_step:.2f}s", + } + ) + + # Evaluation + if ( + args.eval_data_path is not None + and eval_dataloader is not None + and global_step % (args.eval_interval * args.draft_accumulation_steps) + == 0 + ): + draft_model.eval() + eval_metrics_accum = {} + + for eval_data in tqdm( + eval_dataloader, desc=f"Evaluating Epoch {epoch}" + ): + with torch.no_grad(): + _, eval_m = run_forward( + args, peagle_model, eval_data, target_model + ) + for k, v in eval_m.items(): + if k not in eval_metrics_accum: + eval_metrics_accum[k] = [] + eval_metrics_accum[k].append(v) + + avg_metrics = { + k: torch.stack(v).mean() for k, v in eval_metrics_accum.items() + } + record_metrics( + args, + avg_metrics, + global_step // args.draft_accumulation_steps, + tracker, + mode="eval", + ) + draft_model.train() + + # Save Checkpoints + if global_step % args.save_interval == 0: + save_checkpoints(args, epoch, global_step, peagle_model, optimizer) + + if args.max_num_steps is not None and global_step >= args.max_num_steps: + break + + if args.max_num_steps is not None and global_step >= args.max_num_steps: + break + + if global_step % args.save_interval != 0: + print_on_rank0( + f"Training completed at step {global_step}, saving final checkpoint..." + ) + save_checkpoints(args, epoch, global_step, peagle_model, optimizer) + + tracker.close() + destroy_distributed() + + +if __name__ == "__main__": + main() diff --git a/specforge/core/__init__.py b/specforge/core/__init__.py index 8a18642ff..4d5dcc644 100644 --- a/specforge/core/__init__.py +++ b/specforge/core/__init__.py @@ -1,10 +1,12 @@ from .dflash import OnlineDFlashModel from .domino import OnlineDominoModel from .eagle3 import OnlineEagle3Model, QwenVLOnlineEagle3Model +from .peagle import OnlinePEagleModel __all__ = [ "OnlineDFlashModel", "OnlineDominoModel", "OnlineEagle3Model", + "OnlinePEagleModel", "QwenVLOnlineEagle3Model", ] diff --git a/specforge/core/peagle.py b/specforge/core/peagle.py new file mode 100644 index 000000000..f2ab62352 --- /dev/null +++ b/specforge/core/peagle.py @@ -0,0 +1,296 @@ +"""P-EAGLE (Parallel EAGLE) training wrapper with COD sampling.""" + +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.nn as nn +from torch.nn.attention.flex_attention import create_block_mask + +from specforge.core.loss import LogSoftmaxLoss +from specforge.modeling.draft.peagle import PEagleDraftModel + + +def generate_cod_sample_indices( + seq_length: int, + loss_mask: torch.Tensor, + num_depths: int = 8, + down_sample_ratio: float = 0.8, + down_sample_ratio_min: float = 0.2, + filter_position_zero: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + loss_mask = loss_mask.squeeze(0) + device = loss_mask.device + all_valid_indices = torch.where(loss_mask == 1)[0] + + sample_indices = [torch.arange(seq_length, device=device)] + n_per_depth = [seq_length] + prev_indices = all_valid_indices + + for d in range(1, num_depths): + valid_length = max(0, all_valid_indices.shape[0] - d) + ratio = max(down_sample_ratio**d, down_sample_ratio_min) + sample_size = int(valid_length * ratio) + + if sample_size <= 0: + break + + if prev_indices.shape[0] >= sample_size: + random_selection = torch.randperm(prev_indices.shape[0], device=device)[ + :sample_size + ] + sampled_idx = prev_indices[random_selection] + sampled_idx = torch.sort(sampled_idx)[0] + else: + sampled_idx = prev_indices + + next_candidates = (sampled_idx + 1) % seq_length + if filter_position_zero: + next_candidates = next_candidates[next_candidates != 0] + mask = torch.isin(next_candidates, all_valid_indices) + prev_indices = next_candidates[mask] + + sample_indices.append(sampled_idx - d) + n_per_depth.append(sampled_idx.shape[0]) + + anchor_pos = torch.cat(sample_indices) + depth = torch.cat( + [ + torch.full((n,), i, device=device, dtype=torch.long) + for i, n in enumerate(n_per_depth) + ] + ) + return anchor_pos, depth + + +def create_peagle_mask_mod(anchor_pos, depth, lengths, total_seq_len): + document_ids = torch.repeat_interleave( + torch.arange(lengths.shape[0], device=lengths.device, dtype=torch.long), + lengths, + ) + document_ids = torch.cat( + [ + document_ids, + -1 + * torch.ones( + total_seq_len - document_ids.shape[0], + device=lengths.device, + dtype=torch.long, + ), + ] + ).contiguous() + + def peagle_mask_mod(_b, _h, q_idx, kv_idx): + q_anchor_pos = anchor_pos[q_idx] + kv_anchor_pos = anchor_pos[kv_idx] + q_depth = depth[q_idx] + kv_depth = depth[kv_idx] + + same_document = document_ids[q_anchor_pos] == document_ids[kv_anchor_pos] + is_not_padding = document_ids[q_anchor_pos] != -1 + same_rollout = q_anchor_pos == kv_anchor_pos + kv_depth0 = kv_depth == 0 + in_depth_order = q_depth >= kv_depth + is_anchor_causal = q_anchor_pos >= kv_anchor_pos + + return ( + is_not_padding + & same_document + & ((kv_depth0 & is_anchor_causal) | (same_rollout & in_depth_order)) + ) + + return peagle_mask_mod + + +def compute_peagle_metrics( + logits: torch.Tensor, + targets: torch.Tensor, + loss_mask: torch.Tensor, + anchor_pos: torch.Tensor, + depth: torch.Tensor, + num_depths: int, + t2d: torch.Tensor, +) -> Tuple[torch.Tensor, Dict[str, Any]]: + device = logits.device + orig_positions = anchor_pos + depth + + # Ensure loss_mask is 2D [batch, seq_len] + if loss_mask.dim() == 3: + loss_mask = loss_mask.squeeze(-1) + + sampled_loss_mask = loss_mask[:, orig_positions].float() # [batch, total_sampled] + + # Map targets to draft vocabulary and skip positions whose target top-1 + # token is outside the draft vocabulary. + target_logits = targets[:, orig_positions, :] + if t2d is not None and t2d.dtype == torch.bool: + target_top1 = targets.argmax(dim=-1)[:, orig_positions] + target_in_draft_vocab = t2d[target_top1].to(sampled_loss_mask.dtype) + sampled_loss_mask = sampled_loss_mask * target_in_draft_vocab + target_logits = target_logits[:, :, t2d] + + target_p = torch.nn.functional.softmax(target_logits.float(), dim=-1) + position_mask = sampled_loss_mask.unsqueeze(-1) # [batch, total_sampled, 1] + total_positions = position_mask.shape[0] * position_mask.shape[1] + denominator = sampled_loss_mask.sum().clamp_min(1e-6) + loss = LogSoftmaxLoss.apply(logits, target_p, position_mask) * ( + total_positions / denominator + ) + + with torch.no_grad(): + pred_ids = torch.argmax(logits, dim=-1) # [batch, total_sampled] + target_ids = torch.argmax(target_p, dim=-1) # [batch, total_sampled] + + metrics: Dict[str, Any] = { + "loss_sum": loss.detach(), + "loss_total": torch.tensor(1.0, device=device), + } + + correct_total = torch.tensor(0.0, device=device) + count_total = torch.tensor(0.0, device=device) + for d in range(num_depths): + depth_mask = (depth == d).unsqueeze(0) & (sampled_loss_mask > 0.5) + d_correct = ((pred_ids == target_ids) & depth_mask).sum().float() + d_total = depth_mask.sum().float() + metrics[f"position_{d}_acc_sum"] = d_correct + metrics[f"position_{d}_acc_total"] = d_total + correct_total += d_correct + count_total += d_total + + metrics["full_acc_sum"] = correct_total + metrics["full_acc_total"] = count_total + + return loss, metrics + + +class OnlinePEagleModel(nn.Module): + + def __init__( + self, + draft_model: PEagleDraftModel, + mask_token_id: int, + num_depths: int = 8, + down_sample_ratio: float = 0.7, + down_sample_ratio_min: float = 0.2, + ): + super().__init__() + self.draft_model = draft_model + self.mask_token_id = mask_token_id + self.num_depths = num_depths + self.down_sample_ratio = down_sample_ratio + self.down_sample_ratio_min = down_sample_ratio_min + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + target: torch.Tensor, + loss_mask: torch.Tensor, + hidden_states: torch.Tensor, + lengths: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Dict[str, Any]]: + """P-EAGLE training forward pass. + + Args: + input_ids: [batch, seq_len] - input token IDs + attention_mask: [batch, seq_len] - padding mask + target: [batch, seq_len, vocab_size] - target logits from target model + loss_mask: [batch, seq_len] - which positions contribute to loss + hidden_states: [batch, seq_len, 3*hidden_size] - concatenated aux hidden states + lengths: [num_samples] - sequence lengths for multi-sample packing + """ + device = hidden_states.device + seq_length = input_ids.shape[1] + + # Ensure loss_mask is 2D [batch, seq_len] + if loss_mask.dim() == 3: + loss_mask = loss_mask.squeeze(-1) + + if lengths is None: + lengths = torch.tensor([seq_length], dtype=torch.long, device=device) + + # Step 1: COD sampling + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_length, + loss_mask=loss_mask, + num_depths=self.num_depths, + down_sample_ratio=self.down_sample_ratio, + down_sample_ratio_min=self.down_sample_ratio_min, + ) + total_sampled = anchor_pos.shape[0] + orig_positions = anchor_pos + depth + is_depth_0 = depth == 0 + + # Step 2: Build sampled input_ids + sampled_ids = torch.where( + is_depth_0, + input_ids[0, orig_positions], + torch.tensor(self.mask_token_id, dtype=input_ids.dtype, device=device), + ).unsqueeze(0) + + inputs_embeds = self.draft_model.embed_input_ids(sampled_ids).to( + hidden_states.dtype + ) + + # Step 3: Build sampled hidden states + mask_hidden = self.draft_model.mask_hidden.to( + device=device, dtype=hidden_states.dtype + ) + sampled_hidden = torch.where( + is_depth_0.unsqueeze(-1), + hidden_states[0, orig_positions], + mask_hidden.squeeze(0).expand(orig_positions.shape[0], -1), + ).unsqueeze(0) + + # Step 4: Project and concatenate + sampled_hidden = self.draft_model.project_hidden_states(sampled_hidden) + layer_input = torch.cat([inputs_embeds, sampled_hidden], dim=-1) + + # Step 5: Position IDs and rotary embeddings + position_ids = orig_positions.unsqueeze(0) + + # Step 6: Create flex attention mask + mask_mod = create_peagle_mask_mod( + anchor_pos=anchor_pos, + depth=depth, + lengths=lengths, + total_seq_len=seq_length, + ) + block_mask = create_block_mask( + mask_mod, + B=None, + H=None, + Q_LEN=total_sampled, + KV_LEN=total_sampled, + device=device, + ) + + # Step 7: Run through draft model layers + cos, sin = self.draft_model.rotary_emb( + layer_input, seq_len=position_ids.max().item() + 1 + ) + cos = cos.squeeze(0).squeeze(0) + sin = sin.squeeze(0).squeeze(0) + cos = cos[position_ids] + sin = sin[position_ids] + position_embeddings = (cos, sin) + + h = layer_input + for layer in self.draft_model.layers: + h = layer(h, block_mask, position_embeddings) + + # Step 8: Compute logits + logits = self.draft_model.compute_logits(h) + + # Step 9: Compute loss and metrics (target is already logits from target model) + loss, metrics = compute_peagle_metrics( + logits=logits, + targets=target, + loss_mask=loss_mask, + anchor_pos=anchor_pos, + depth=depth, + num_depths=self.num_depths, + t2d=self.draft_model.t2d, + ) + + return loss, metrics diff --git a/specforge/data/parse.py b/specforge/data/parse.py index b9a7cccdf..eca025509 100644 --- a/specforge/data/parse.py +++ b/specforge/data/parse.py @@ -102,6 +102,20 @@ def _sanitize_message(self, message: dict) -> dict: return cleaned + def _normalize_message(self, message: dict) -> dict: + role = message.get("role", message.get("from", "")) + content = message.get("content") or message.get("value") or "" + + if role in ("human", "user"): + role = "user" + elif role in ("gpt", "assistant"): + role = "assistant" + + normalized = {**message, "role": role, "content": content} + normalized.pop("from", None) + normalized.pop("value", None) + return normalized + _harmony_encoding = None @@ -157,6 +171,9 @@ def parse( **kwargs, ) -> Dict[str, List[torch.Tensor]]: if not preformatted: + conversation = [ + self._normalize_message(message) for message in conversation + ] messages = [] if conversation[0]["role"] == "system": @@ -171,6 +188,12 @@ def parse( if self.system_prompt: messages.append({"role": "system", "content": self.system_prompt}) + while conversation and conversation[0]["role"] != "user": + warnings.warn( + f"Dropping leading '{conversation[0]['role']}' message before the first user turn." + ) + conversation = conversation[1:] + for j, sentence in enumerate(conversation): role = sentence["role"] if j == 0: diff --git a/specforge/data/preprocessing.py b/specforge/data/preprocessing.py index ea34bb271..972fbd4e3 100644 --- a/specforge/data/preprocessing.py +++ b/specforge/data/preprocessing.py @@ -306,6 +306,7 @@ def build_eagle3_dataset( processor: Optional[ImageProcessingMixin] = None, is_preformatted: Optional[bool] = False, train_only_last_turn: Optional[bool] = False, + minimum_valid_tokens: Optional[int] = None, ) -> HFDataset: """ build eagle3 dataset @@ -333,10 +334,14 @@ def build_eagle3_dataset( If False, expects "conversations" column with ShareGPT format. train_only_last_turn: If True, only the last assistant turn contributes to the loss. Useful for thinking models where history may not contain thoughts. + minimum_valid_tokens: If set, drops samples with fewer trainable tokens. Returns: The processed HF dataset. """ + if minimum_valid_tokens is not None and minimum_valid_tokens < 0: + raise ValueError("minimum_valid_tokens must be >= 0") + if is_vlm: assert processor is not None, "processor must be provided when is_vlm is True" @@ -460,6 +465,30 @@ def preprocess_function(examples): cache_file_name=cache_file_name, ) + if minimum_valid_tokens is not None: + before_filter = len(dataset) + + def has_minimum_valid_tokens(example): + loss_mask = example["loss_mask"] + if isinstance(loss_mask, torch.Tensor): + valid_tokens = int(loss_mask.sum().item()) + else: + valid_tokens = sum( + int(token) + for row in loss_mask + for token in (row if isinstance(row, list) else [row]) + ) + return valid_tokens >= minimum_valid_tokens + + dataset = dataset.filter( + has_minimum_valid_tokens, + num_proc=num_proc, + desc=f"Filtering samples with >= {minimum_valid_tokens} trainable tokens", + ) + print( + f"Filtered dataset by trainable tokens: {before_filter} -> {len(dataset)}" + ) + dataset.set_format(type="torch") return dataset diff --git a/specforge/modeling/__init__.py b/specforge/modeling/__init__.py index 09999d60b..77e4b7be5 100644 --- a/specforge/modeling/__init__.py +++ b/specforge/modeling/__init__.py @@ -1,6 +1,7 @@ # from .auto import AutoDistributedTargetModel, AutoDraftModelConfig, AutoEagle3DraftModel from .auto import AutoDraftModelConfig, AutoEagle3DraftModel from .draft.llama3_eagle import LlamaForCausalLMEagle3 +from .draft.peagle import PEagleDraftModel from .target.eagle3_target_model import ( CustomEagle3TargetModel, HFEagle3TargetModel, @@ -10,6 +11,7 @@ __all__ = [ "LlamaForCausalLMEagle3", + "PEagleDraftModel", "SGLangEagle3TargetModel", "HFEagle3TargetModel", "CustomEagle3TargetModel", diff --git a/specforge/modeling/auto.py b/specforge/modeling/auto.py index 1e48a43e7..d52759b4e 100644 --- a/specforge/modeling/auto.py +++ b/specforge/modeling/auto.py @@ -133,6 +133,7 @@ class AutoDraftModelConfig: _config_mapping = { "LlamaForCausalLMEagle3": LlamaConfig, + "PEagleDraftModel": LlamaConfig, } @classmethod diff --git a/specforge/modeling/draft/__init__.py b/specforge/modeling/draft/__init__.py index 8bdc7e2f6..6130dcc63 100644 --- a/specforge/modeling/draft/__init__.py +++ b/specforge/modeling/draft/__init__.py @@ -6,11 +6,13 @@ sample, ) from .llama3_eagle import LlamaForCausalLMEagle3 +from .peagle import PEagleDraftModel __all__ = [ "Eagle3DraftModel", "DFlashDraftModel", "LlamaForCausalLMEagle3", + "PEagleDraftModel", "build_target_layer_ids", "extract_context_feature", "sample", diff --git a/specforge/modeling/draft/peagle.py b/specforge/modeling/draft/peagle.py new file mode 100644 index 000000000..e68f99849 --- /dev/null +++ b/specforge/modeling/draft/peagle.py @@ -0,0 +1,288 @@ +"""P-EAGLE (Parallel EAGLE) draft model with multi-layer architecture.""" + +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +from transformers.cache_utils import Cache +from transformers.models.llama.configuration_llama import LlamaConfig + +from specforge.modeling.draft.base import Eagle3DraftModel +from specforge.modeling.draft.flex_attention import compile_friendly_flex_attention +from specforge.modeling.draft.llama3_eagle import ( + LlamaMLP, + LlamaRMSNorm, + LlamaRotaryEmbedding, + rotate_half, +) + + +class PEagleAttention(nn.Module): + """Flex-attention layer for P-EAGLE. Accepts pre-computed BlockMask and position embeddings.""" + + def __init__(self, config: LlamaConfig, input_size: int): + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + + self.q_proj = nn.Linear(input_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear( + input_size, self.num_key_value_heads * self.head_dim, bias=False + ) + self.v_proj = nn.Linear( + input_size, self.num_key_value_heads * self.head_dim, bias=False + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, self.hidden_size, bias=False + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + + query_states = ( + self.q_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + + # position_embeddings: (cos, sin), each [batch, seq_len, head_dim], pre-indexed + cos, sin = position_embeddings + cos = cos.unsqueeze(1) # [batch, 1, seq_len, head_dim] + sin = sin.unsqueeze(1) + query_states = (query_states * cos) + (rotate_half(query_states) * sin) + key_states = (key_states * cos) + (rotate_half(key_states) * sin) + + attn_output = compile_friendly_flex_attention( + query=query_states, + key=key_states, + value=value_states, + block_mask=attention_mask, + enable_gqa=True, + kernel_options={ + "FORCE_USE_FLEX_ATTENTION": True, + "BLOCK_M": 64, + "BLOCK_N": 64, + "num_stages": 2, + "bwd_BLOCK_M1": 32, + "bwd_BLOCK_N1": 64, + "bwd_BLOCK_M2": 32, + "bwd_BLOCK_N2": 32, + }, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + return self.o_proj(attn_output) + + +class PEagleFirstLayer(nn.Module): + """Eagle3-style first decoder layer: splits 2*hidden_size input into embeds and hidden, + normalizes separately, then runs attention with 2*hidden_size Q/K/V projections.""" + + def __init__(self, config: LlamaConfig, norm_before_residual: bool = False): + super().__init__() + self.hidden_size = config.hidden_size + self.norm_before_residual = norm_before_residual + + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hidden_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = PEagleAttention(config, input_size=2 * config.hidden_size) + self.post_attention_layernorm = LlamaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp = LlamaMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + mid = hidden_states.shape[2] // 2 + embeds, hidden = hidden_states.split(mid, dim=-1) + residual = hidden + + embeds = self.input_layernorm(embeds) + hidden = self.hidden_norm(hidden) + if self.norm_before_residual: + residual = hidden + hidden_states = torch.cat([embeds, hidden], dim=-1) + + hidden_states = self.self_attn( + hidden_states, attention_mask, position_embeddings + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class PEagleStandardLayer(nn.Module): + """Standard decoder layer for subsequent P-EAGLE layers: hidden_size input.""" + + def __init__(self, config: LlamaConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = PEagleAttention(config, input_size=config.hidden_size) + self.post_attention_layernorm = LlamaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp = LlamaMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states, attention_mask, position_embeddings + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class PEagleDraftModel(Eagle3DraftModel): + """P-EAGLE draft model with multi-layer architecture. + + Architecture follows speculators Eagle3DraftModel: + - First layer: Eagle3-style with 2*hidden_size Q/K/V (splits embeds + hidden) + - Subsequent layers: Standard decoder layers with hidden_size Q/K/V + - External rotary embeddings shared across all layers + """ + + config_class = LlamaConfig + + def __init__( + self, + config: LlamaConfig, + norm_before_residual: bool = False, + ) -> None: + super().__init__(config) + self.config = config + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.draft_vocab_size = getattr(config, "draft_vocab_size", config.vocab_size) + self.norm_before_residual = norm_before_residual + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id + ) + + if hasattr(config, "target_hidden_size"): + fc_input_size = config.target_hidden_size * 3 + else: + fc_input_size = config.hidden_size * 3 + self.fc = nn.Linear(fc_input_size, config.hidden_size, bias=False) + self.mask_hidden = nn.Parameter(torch.randn(1, 1, fc_input_size)) + + num_layers = config.num_hidden_layers + layers: List[nn.Module] = [ + PEagleFirstLayer(config, norm_before_residual=norm_before_residual) + ] + for _ in range(1, num_layers): + layers.append(PEagleStandardLayer(config)) + self.layers = nn.ModuleList(layers) + + self.rotary_emb = LlamaRotaryEmbedding( + dim=getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ), + max_position_embeddings=config.max_position_embeddings, + base=getattr(config, "rope_theta", 10000), + ) + + self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.lm_head = nn.Linear(config.hidden_size, self.draft_vocab_size, bias=False) + + t2d = torch.ones(self.vocab_size, dtype=torch.bool) + d2t = torch.zeros(self.draft_vocab_size, dtype=torch.int64) + self.register_buffer("t2d", t2d) + self.register_buffer("d2t", d2t) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def project_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.fc(hidden_states) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(self.norm(hidden_states)) + + def backbone( + self, + input_embeds: torch.Tensor, + hidden_states: torch.Tensor, + cache_hidden: torch.Tensor = None, + attention_mask=None, + position_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: bool = True, + ) -> torch.Tensor: + """Run multi-layer forward pass. + + Args: + input_embeds: [batch, seq_len, hidden_size] - token embeddings + hidden_states: [batch, seq_len, hidden_size] - projected aux hidden states + attention_mask: BlockMask from flex_attention + position_ids: [batch, seq_len] - position indices + """ + layer_input = torch.cat([input_embeds, hidden_states], dim=-1) + + cos, sin = self.rotary_emb(layer_input, seq_len=position_ids.max().item() + 1) + cos = cos.squeeze(0).squeeze(0) + sin = sin.squeeze(0).squeeze(0) + cos = cos[position_ids] + sin = sin[position_ids] + position_embeddings = (cos, sin) + + h = layer_input + for layer in self.layers: + h = layer(h, attention_mask, position_embeddings) + return h + + def forward( + self, + hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor, + attention_mask=None, + position_ids: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """Convenience forward that projects hidden states and runs backbone.""" + projected = self.fc(hidden_states) + h = self.backbone( + input_embeds=inputs_embeds, + hidden_states=projected, + attention_mask=attention_mask, + position_ids=position_ids, + ) + return self.norm(h) diff --git a/tests/test_data/test_parser_normalization.py b/tests/test_data/test_parser_normalization.py new file mode 100644 index 000000000..9cef91ef1 --- /dev/null +++ b/tests/test_data/test_parser_normalization.py @@ -0,0 +1,88 @@ +import unittest + +from specforge.data.parse import GeneralParser +from specforge.data.template import ChatTemplate + + +class DummyTokenizer: + pad_token_id = 0 + unk_token_id = 0 + bos_token = None + + def __init__(self): + self.messages = None + + def apply_chat_template( + self, + messages, + tokenize=False, + add_generation_prompt=False, + tools=None, + **kwargs, + ): + self.messages = messages + return "".join( + f"<{message['role']}>{message['content']}" for message in messages + ) + + def __call__( + self, + text, + max_length, + truncation, + return_tensors, + add_special_tokens, + ): + class Encoding: + pass + + encoding = Encoding() + encoding.input_ids = self._ids(text, max_length)[None, :] + return encoding + + def encode(self, text, add_special_tokens=False, truncation=True, max_length=None): + return self._ids(text, max_length).tolist() + + def _ids(self, text, max_length=None): + import torch + + values = [ord(char) % 251 for char in text] + if max_length is not None: + values = values[:max_length] + return torch.tensor(values, dtype=torch.long) + + +class TestParserNormalization(unittest.TestCase): + def test_general_parser_normalizes_sharegpt_keys_and_drops_leading_non_user(self): + tokenizer = DummyTokenizer() + parser = GeneralParser( + tokenizer, + ChatTemplate( + assistant_header="", + user_header="", + system_prompt=None, + end_of_turn_token="", + ), + ) + + with self.assertWarnsRegex(Warning, "Dropping leading 'assistant'"): + parser.parse( + [ + {"from": "gpt", "value": "orphan assistant"}, + {"from": "human", "value": "question"}, + {"from": "gpt", "value": "answer"}, + ], + max_length=512, + ) + + self.assertEqual( + tokenizer.messages, + [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ], + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_data/test_peagle_data_preprocessing.py b/tests/test_data/test_peagle_data_preprocessing.py new file mode 100644 index 000000000..1cb110842 --- /dev/null +++ b/tests/test_data/test_peagle_data_preprocessing.py @@ -0,0 +1,110 @@ +import unittest + +import torch + +from datasets import Dataset +from specforge.data.preprocessing import build_eagle3_dataset +from specforge.data.template import TEMPLATE_REGISTRY, ChatTemplate + + +class DummyTokenizer: + pad_token_id = 0 + unk_token_id = 0 + bos_token = None + + def apply_chat_template( + self, + messages, + tokenize=False, + add_generation_prompt=False, + tools=None, + **kwargs, + ): + return "".join( + f"<{message['role']}>{message['content']}" for message in messages + ) + + def __call__( + self, + text, + max_length, + truncation, + return_tensors, + add_special_tokens, + ): + class Encoding: + pass + + encoding = Encoding() + encoding.input_ids = self._ids(text, max_length)[None, :] + return encoding + + def encode(self, text, add_special_tokens=False, truncation=True, max_length=None): + return self._ids(text, max_length).tolist() + + def _ids(self, text, max_length=None): + values = [ord(char) % 251 for char in text] + if max_length is not None: + values = values[:max_length] + return torch.tensor(values, dtype=torch.long) + + +class TestPEagleDataPreprocessing(unittest.TestCase): + template_name = "unit-test-peagle-data" + + @classmethod + def setUpClass(cls): + if cls.template_name not in TEMPLATE_REGISTRY.get_all_template_names(): + TEMPLATE_REGISTRY.register( + cls.template_name, + ChatTemplate( + assistant_header="", + user_header="", + system_prompt=None, + end_of_turn_token="", + ), + ) + + def test_minimum_valid_tokens_filters_empty_loss_samples(self): + dataset = Dataset.from_list( + [ + { + "conversations": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + }, + { + "conversations": [ + {"role": "user", "content": "question"}, + ] + }, + ] + ) + + result = build_eagle3_dataset( + dataset=dataset, + tokenizer=DummyTokenizer(), + chat_template=self.template_name, + max_length=128, + num_proc=1, + cache_dir=None, + cache_key=None, + minimum_valid_tokens=1, + ) + + self.assertEqual(len(result), 1) + self.assertGreater(result[0]["loss_mask"].sum().item(), 0) + + def test_minimum_valid_tokens_rejects_negative_values(self): + with self.assertRaisesRegex(ValueError, "minimum_valid_tokens"): + build_eagle3_dataset( + dataset=Dataset.from_list([]), + tokenizer=DummyTokenizer(), + chat_template=self.template_name, + minimum_valid_tokens=-1, + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_utils/test_peagle.py b/tests/test_utils/test_peagle.py new file mode 100644 index 000000000..cafa9a65e --- /dev/null +++ b/tests/test_utils/test_peagle.py @@ -0,0 +1,215 @@ +import unittest +from argparse import Namespace +from unittest.mock import MagicMock, patch + +import torch +from transformers import LlamaConfig + +from scripts.train_peagle import resolve_mask_token_id +from specforge.core.peagle import ( + OnlinePEagleModel, + compute_peagle_metrics, + create_peagle_mask_mod, + generate_cod_sample_indices, +) +from specforge.modeling.draft.peagle import PEagleDraftModel + + +class TestPEagleTrainingSemantics(unittest.TestCase): + def _tiny_config(self): + return LlamaConfig( + vocab_size=32, + draft_vocab_size=16, + hidden_size=16, + intermediate_size=32, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + max_position_embeddings=64, + pad_token_id=0, + rms_norm_eps=1e-5, + ) + + def test_mask_hidden_is_part_of_draft_checkpoint_state(self): + config = self._tiny_config() + model = PEagleDraftModel(config) + + with torch.no_grad(): + model.mask_hidden.fill_(3.0) + + reloaded = PEagleDraftModel(config) + reloaded.load_state_dict(model.state_dict()) + + torch.testing.assert_close(reloaded.mask_hidden, model.mask_hidden) + + def test_online_wrapper_uses_draft_model_mask_hidden(self): + config = self._tiny_config() + draft_model = PEagleDraftModel(config) + wrapper = OnlinePEagleModel(draft_model=draft_model, mask_token_id=0) + + self.assertIs(wrapper.draft_model.mask_hidden, draft_model.mask_hidden) + self.assertNotIn("mask_hidden", dict(wrapper.named_parameters(recurse=False))) + + def test_peagle_embeddings_are_trainable_by_default(self): + config = self._tiny_config() + model = PEagleDraftModel(config) + + self.assertTrue(model.embed_tokens.weight.requires_grad) + + def test_compute_metrics_masks_targets_outside_draft_vocab(self): + logits = torch.tensor( + [ + [ + [0.0, 4.0], + [4.0, 0.0], + [0.0, 4.0], + ] + ], + dtype=torch.float32, + ) + targets = torch.full((1, 3, 4), -10.0, dtype=torch.float32) + targets[0, 0, 1] = 10.0 + targets[0, 1, 2] = 10.0 + targets[0, 2, 0] = 10.0 + loss_mask = torch.ones(1, 3) + anchor_pos = torch.tensor([0, 1, 2]) + depth = torch.tensor([0, 0, 0]) + t2d = torch.tensor([True, True, False, False]) + + def fake_loss(logits, target_p, position_mask): + return torch.tensor(0.0, device=logits.device) + + with patch("specforge.core.peagle.LogSoftmaxLoss.apply", side_effect=fake_loss): + _loss, metrics = compute_peagle_metrics( + logits=logits, + targets=targets, + loss_mask=loss_mask, + anchor_pos=anchor_pos, + depth=depth, + num_depths=1, + t2d=t2d, + ) + + self.assertEqual(metrics["position_0_acc_total"].item(), 2.0) + self.assertEqual(metrics["position_0_acc_sum"].item(), 1.0) + + def test_cod_sampling_uses_valid_targets_for_parallel_depths(self): + torch.manual_seed(0) + loss_mask = torch.tensor([[0, 1, 1, 1, 0, 1]]) + + anchor_pos, depth = generate_cod_sample_indices( + seq_length=loss_mask.shape[1], + loss_mask=loss_mask, + num_depths=4, + down_sample_ratio=1.0, + down_sample_ratio_min=1.0, + ) + + self.assertEqual(anchor_pos[: loss_mask.shape[1]].tolist(), list(range(6))) + self.assertEqual(depth[: loss_mask.shape[1]].tolist(), [0] * 6) + + sampled_target_pos = anchor_pos + depth + parallel_depth_mask = depth > 0 + self.assertTrue(torch.all(sampled_target_pos[parallel_depth_mask] >= 0)) + self.assertTrue(torch.all(sampled_target_pos[parallel_depth_mask] < 6)) + self.assertTrue( + torch.all(loss_mask[0, sampled_target_pos[parallel_depth_mask]] == 1) + ) + + def test_peagle_mask_respects_documents_depth_order_and_padding(self): + anchor_pos = torch.tensor([0, 1, 1, 2, 4, 4, 5]) + depth = torch.tensor([0, 0, 1, 0, 0, 1, 0]) + lengths = torch.tensor([3, 2]) + mask_mod = create_peagle_mask_mod( + anchor_pos=anchor_pos, + depth=depth, + lengths=lengths, + total_seq_len=6, + ) + + def allowed(q_idx, kv_idx): + return bool( + mask_mod( + None, + None, + torch.tensor(q_idx), + torch.tensor(kv_idx), + ).item() + ) + + self.assertTrue(allowed(2, 1)) # same rollout, depth 1 attends depth 0 + self.assertTrue(allowed(2, 0)) # depth 1 also attends causal depth-0 context + self.assertFalse(allowed(1, 2)) # depth 0 cannot attend a future depth + self.assertFalse(allowed(4, 3)) # different packed documents + self.assertFalse(allowed(6, 6)) # padding anchor position + + +class TestPEagleMaskTokenResolution(unittest.TestCase): + def _args(self, mask_token_id=None): + return Namespace( + mask_token_id=mask_token_id, + target_model_path="target", + trust_remote_code=False, + ) + + def test_explicit_mask_token_is_validated(self): + with self.assertRaises(ValueError): + resolve_mask_token_id(self._args(mask_token_id=33), embedding_vocab_size=32) + with self.assertRaises(ValueError): + resolve_mask_token_id(self._args(mask_token_id=-1), embedding_vocab_size=32) + + self.assertEqual( + resolve_mask_token_id( + self._args(mask_token_id=31), embedding_vocab_size=32 + ), + 31, + ) + + @patch("scripts.train_peagle.AutoTokenizer") + def test_tokenizer_mask_token_takes_priority(self, mock_auto_tokenizer): + tokenizer = MagicMock() + tokenizer.mask_token_id = 7 + mock_auto_tokenizer.from_pretrained.return_value = tokenizer + + self.assertEqual(resolve_mask_token_id(self._args(), 32), 7) + + @patch("scripts.train_peagle.AutoTokenizer") + def test_unused_embedding_slot_takes_priority_over_pad(self, mock_auto_tokenizer): + tokenizer = MagicMock() + tokenizer.mask_token_id = None + tokenizer.pad_token_id = 3 + tokenizer.eos_token_id = 4 + tokenizer.unk_token_id = 5 + tokenizer.__len__.return_value = 30 + mock_auto_tokenizer.from_pretrained.return_value = tokenizer + + self.assertEqual(resolve_mask_token_id(self._args(), 32), 30) + + @patch("scripts.train_peagle.AutoTokenizer") + def test_pad_fallback_when_no_mask_or_unused_slot(self, mock_auto_tokenizer): + tokenizer = MagicMock() + tokenizer.mask_token_id = None + tokenizer.pad_token_id = 3 + tokenizer.eos_token_id = 4 + tokenizer.unk_token_id = 5 + tokenizer.__len__.return_value = 32 + mock_auto_tokenizer.from_pretrained.return_value = tokenizer + + self.assertEqual(resolve_mask_token_id(self._args(), 32), 3) + + @patch("scripts.train_peagle.AutoTokenizer") + def test_fallback_token_must_fit_embedding_vocab(self, mock_auto_tokenizer): + tokenizer = MagicMock() + tokenizer.mask_token_id = None + tokenizer.pad_token_id = 33 + tokenizer.eos_token_id = None + tokenizer.unk_token_id = None + tokenizer.__len__.return_value = 32 + mock_auto_tokenizer.from_pretrained.return_value = tokenizer + + with self.assertRaises(ValueError): + resolve_mask_token_id(self._args(), 32) + + +if __name__ == "__main__": + unittest.main(verbosity=2)