diff --git a/README.md b/README.md index 8b8c952a2..50f649ef8 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ Check out [**our documentation**](https://docs.sglang.ai/SpecForge/) to get star | Method | Description | Training | Example | Optimization | | --- | --- | --- | --- | --- | | **[EAGLE3](https://arxiv.org/abs/2503.01840)** | Feature-based autoregressive drafting | [`scripts/train_eagle3.py`](./scripts/train_eagle3.py) | [`examples/run_qwen3_8b_eagle3_online.sh`](./examples/run_qwen3_8b_eagle3_online.sh) | [LK loss](https://arxiv.org/pdf/2602.23881) +| **[EDSD](https://aclanthology.org/2026.acl-long.2145.pdf)** | EAGLE3 with entropy-driven design | [`scripts/train_edsd.py`](./scripts/train_edsd.py) | [`examples/run_qwen3_8b_edsd_online.sh`](./examples/run_qwen3_8b_edsd_online.sh) | | **[DFlash](https://arxiv.org/abs/2602.06036)** | Block-parallel drafting | [`scripts/train_dflash.py`](./scripts/train_dflash.py) | [`examples/run_qwen3_8b_dflash_online.sh`](./examples/run_qwen3_8b_dflash_online.sh) | [D-PACE](https://arxiv.org/abs/2605.18810) | **[Domino](https://arxiv.org/html/2605.29707v1)** | DFlash with GRU logit correction | [`scripts/train_domino.py`](./scripts/train_domino.py) | [`examples/run_qwen3_8b_domino_online.sh`](./examples/run_qwen3_8b_domino_online.sh) | @@ -47,6 +48,7 @@ SpecBundle is a collection of production-grade speculative decoding models that ## 🎉 News +- [2026-07] 🔥 Added EDSD online training for EAGLE3 draft models. - [2026-06] 🔥 Added D-PACE as an optional loss for DFlash training. - [2026-06] 🔥 Added Domino online training for DFlash draft models. - [2026-01] 🔥 Added DFlash block-parallel online training with SGLang serving support. diff --git a/configs/qwen3-8b-edsd.json b/configs/qwen3-8b-edsd.json new file mode 100644 index 000000000..63bd151de --- /dev/null +++ b/configs/qwen3-8b-edsd.json @@ -0,0 +1,32 @@ +{ + "architectures": [ + "EdsdDraftModel" + ], + "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": 1, + "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, + "target_layer_ids": [23, 34] +} diff --git a/examples/run_qwen3_8b_edsd_online.sh b/examples/run_qwen3_8b_edsd_online.sh new file mode 100644 index 000000000..cf6bf1dff --- /dev/null +++ b/examples/run_qwen3_8b_edsd_online.sh @@ -0,0 +1,31 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +ROOT_DIR=$(dirname $SCRIPT_DIR) +export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels + +# support tp8 train eagle3 for Qwen3-4B/8B/32B up to tp_size = 8 +NUM_GPUS=${1:-4} +TP_SIZE=${4:-1} +BUILD_DATASET_NUM_PROC=${BUILD_DATASET_NUM_PROC:-64} + +torchrun \ + --standalone \ + --nproc_per_node $NUM_GPUS \ + $ROOT_DIR/scripts/train_edsd.py \ + --target-model-path /root/spec-train/ckpt/Qwen3-8B \ + --draft-model-config $ROOT_DIR/configs/qwen3-8b-edsd.json \ + --train-data-path /root/spec-train/data/qwen3-8b-test.jsonl \ + --build-dataset-num-proc $BUILD_DATASET_NUM_PROC \ + --output-dir $ROOT_DIR/outputs/qwen3-8b-edsd-sharegpt \ + --num-epochs 10 \ + --batch-size 4 \ + --learning-rate 1e-4 \ + --max-length 4096 \ + --chat-template qwen \ + --cache-dir $ROOT_DIR/cache \ + --embedding-key model.embed_tokens.weight \ + --tp-size $TP_SIZE \ + --target-model-backend sglang \ + --drop-ratio-scale 0.03 \ + --report-to tensorboard diff --git a/scripts/train_edsd.py b/scripts/train_edsd.py new file mode 100644 index 000000000..1e69bb0b7 --- /dev/null +++ b/scripts/train_edsd.py @@ -0,0 +1,1309 @@ +# EDSD training script. +import argparse +import hashlib +import math +import os +import time +from argparse import ArgumentParser, Namespace +from typing import List, Optional, Tuple, Union + +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 Dataset +from specforge import ( + AutoDraftModelConfig, + OnlineEdsdModel, +) +from specforge.modeling.draft.edsd import EdsdDraftModel +from specforge.args import SGLangBackendArgs, TrackerArgs +from specforge.data import ( + build_eagle3_dataset, + build_offline_eagle3_dataset, + generate_vocab_mapping_file, + prepare_dp_dataloaders, +) +from specforge.distributed import ( + destroy_distributed, + get_dp_group, + get_draft_dp_group, + get_tp_group, + init_distributed, +) +from specforge.modeling.target import ( + Eagle3TargetModel, + TargetHead, + get_eagle3_target_model, +) +from specforge.optimizer import BF16Optimizer +from specforge.tracker import Tracker, create_tracker, get_tracker_class +from specforge.utils import ( + create_draft_config_from_target, + get_last_checkpoint, + print_args_with_dots, + print_on_rank0, + print_with_rank, + rank_0_priority, + safe_conversations_generator, +) + + +def print_cuda_memory_debug(label: str) -> None: + if os.getenv("SPECFORGE_CI_MEMORY_DEBUG") != "1" or not torch.cuda.is_available(): + return + + try: + torch.cuda.synchronize() + free_bytes, total_bytes = torch.cuda.mem_get_info() + allocated_bytes = torch.cuda.memory_allocated() + reserved_bytes = torch.cuda.memory_reserved() + except Exception as exc: + print(f"[memory-debug] {label}: failed to query CUDA memory: {exc}", flush=True) + return + + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else "NA" + local_rank = os.getenv("LOCAL_RANK", "NA") + print( + "[memory-debug] " + f"{label}: rank={rank} local_rank={local_rank} " + f"free={free_bytes / 1024**3:.2f}GiB " + f"used={(total_bytes - free_bytes) / 1024**3:.2f}GiB " + f"total={total_bytes / 1024**3:.2f}GiB " + f"torch_allocated={allocated_bytes / 1024**3:.2f}GiB " + f"torch_reserved={reserved_bytes / 1024**3:.2f}GiB", + flush=True, + ) + + +def build_parser() -> ArgumentParser: + """Build the training argument parser (import-safe seam for tests).""" + parser = argparse.ArgumentParser(description="Train Eagle3 with online data") + + # add model-related arguments + 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, + help="Draft model config path. If not provided, will auto-generate from target model.", + ) + model_group.add_argument( + "--embedding-key", + type=str, + default="model.embed_tokens.weight", + help="The key of the embedding weight to load from the target model", + ) + model_group.add_argument( + "--lm-head-key", + type=str, + default="lm_head.weight", + help="The key of the lm head weight to load from the target model, this is only required for offline training", + ) + model_group.add_argument( + "--target-model-backend", + type=str, + default="sglang", + choices=["sglang", "hf", "custom"], + help="The backend of the target model", + ) + + # dataset arguments + dataset_group = parser.add_argument_group("dataset") + dataset_group.add_argument("--train-data-path", type=str, required=True) + dataset_group.add_argument("--train-hidden-states-path", type=str, default=None) + dataset_group.add_argument("--eval-hidden-states-path", type=str, default=None) + 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", + help="Whether the input data is preformatted text with the chat template already applied to the conversation messages.", + ) + dataset_group.add_argument( + "--train-only-last-turn", + action="store_true", + help="If set, only the last assistant turn in each conversation contributes to the loss. " + "Useful for thinking models where conversation history may lack thought processes.", + ) + dataset_group.add_argument("--build-dataset-num-proc", type=int, default=8) + dataset_group.add_argument( + "--dataloader-num-workers", + type=int, + default=4, + help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", + ) + # training hyper params + 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, + help="The maximum number of steps to train. If not provided, will be calculated as num_epochs * steps_per_epoch", + ) + training_group.add_argument("--batch-size", type=int, default=1) + training_group.add_argument("--learning-rate", type=float, default=1e-4) + training_group.add_argument("--max-length", type=int, default=2048) + training_group.add_argument("--warmup-ratio", type=float, default=0.015) + training_group.add_argument( + "--total-steps", + type=int, + default=None, + help="Total training steps. If not provided, will be calculated as num_epochs * steps_per_epoch", + ) + training_group.add_argument("--max-grad-norm", type=float, default=0.5) + training_group.add_argument( + "--ttt-length", + type=int, + default=7, + help="The length for Test-Time Training (TTT).", + ) + # EDSD-specific args + training_group.add_argument( + "--drop-ratio-scale", + type=float, + default=0.02, + help=( + "EDSD mask_p: per-epoch curriculum drop ratio is " + "min(max((num_epochs - epoch_idx) * drop_ratio_scale, 0.0), 0.4). " + "Largest at epoch 0, decays to ~0 by the final epoch." + ), + ) + training_group.add_argument( + "--step-n-schedule", + type=int, + nargs="+", + default=None, + help=( + "EDSD step_n: explicit per-epoch TTT length list (monotonic " + "non-decreasing, each in [1, --ttt-length]). epoch >= len uses the " + "last value. If unset, the length is auto-scheduled by EDSD Eq. 6 " + "(n_t = ceil(1 + (S_max-1)*t/(T-1)))." + ), + ) + training_group.add_argument("--resume", action="store_true") + training_group.add_argument( + "--ckpt-dir", + type=str, + default=None, + help="directory includes the checkpoint to start training with", + ) + 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, + help="Log training metrics every N steps", + ) + training_group.add_argument("--seed", type=int, default=0) + training_group.add_argument("--draft-accumulation-steps", type=int, default=1) + + # LK / acceptance-rate loss arguments + lk_group = parser.add_argument_group("lk loss") + lk_group.add_argument( + "--lk-loss-type", + type=str, + default=None, + choices=["lambda", "alpha"], + help="Enable LK loss objective. Choices: lambda (hybrid KL+LK), alpha (pure acceptance-rate likelihood).", + ) + lk_group.add_argument( + "--kl-scale", + type=float, + default=1.0, + help="Scale for adaptive KL weight: kl_weight = kl_scale * exp(-kl_decay * acc). Used when --lk-loss-type=lambda.", + ) + lk_group.add_argument( + "--kl-decay", + type=float, + default=3.0, + help="Decay for adaptive KL weight. Used when --lk-loss-type=lambda.", + ) + + # data processing type + 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", + ) + # distributed training + optimization_group.add_argument("--sp-ulysses-size", type=int, default=1) + optimization_group.add_argument("--sp-ring-size", type=int, default=1) + optimization_group.add_argument( + "--attention-backend", + type=str, + default="flex_attention", + help="The attention backend for the draft model", + ) + + # other args + 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, + help="Timeout for collective communication in minutes", + ) + other_group.add_argument( + "--model-download-dir", + type=str, + default=None, + help="The directory to download the target model to", + ) + + # profiling related args + 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 target model backend related args + sglang_group = parser.add_argument_group("sglang target model backend") + SGLangBackendArgs.add_args(sglang_group) + + # tracker related args + tracker_group = parser.add_argument_group("tracker") + TrackerArgs.add_args(tracker_group) + + return parser + + +def parse_args() -> Tuple[ArgumentParser, Namespace]: + """Parse CLI arguments for the training script.""" + parser = build_parser() + args = parser.parse_args() + return parser, args + + +def build_tracker(args: Namespace, parser: ArgumentParser) -> Tracker: + """ + Build the experiment tracker according to the report_to argument. + + Args: + args: The arguments for the training script. + parser: The parser for the training script. + + Returns: + The experiment 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}") + tracker = create_tracker(args, args.output_dir) + return tracker + + +def build_target_model( + args: Namespace, draft_model_config: AutoDraftModelConfig, is_online: bool = True +) -> Tuple[Union[Eagle3TargetModel, TargetHead], None]: + """ + Build the target model according to the arguments. + + Args: + args: The arguments for the training script. + draft_model_config: The draft model config. + + Returns: + The target model. + """ + if is_online: + 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, + ) + + # set the aux hidden states layers + 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"] + ) + elif hasattr(draft_model_config, "target_layer_ids"): + target_model.set_aux_hidden_states_layers( + draft_model_config.target_layer_ids + ) + else: + target_model.set_aux_hidden_states_layers() + + return target_model, None + else: + target_head = TargetHead.from_pretrained( + model_path=args.target_model_path, + lm_head_key=args.lm_head_key, + cache_dir=args.model_download_dir, + trust_remote_code=args.trust_remote_code, + ) + return target_head, None + + +def sanity_check(args: Namespace) -> None: + """ + Perform sanity checks on the arguments. + + Args: + args: The arguments for the training script. + + Returns: + None + """ + if args.step_n_schedule is not None: + if len(args.step_n_schedule) == 0: + raise ValueError("--step-n-schedule must not be empty") + prev = 0 + for i, v in enumerate(args.step_n_schedule): + if not isinstance(v, int) or v < 1: + raise ValueError( + f"--step-n-schedule values must be positive ints, got {v} at index {i}" + ) + if v < prev: + raise ValueError( + f"--step-n-schedule must be monotonic non-decreasing, " + f"got {v} < {prev} at index {i}" + ) + prev = v + if args.step_n_schedule[-1] > args.ttt_length: + raise ValueError( + f"--step-n-schedule values must be <= --ttt-length " + f"({args.ttt_length}), got max {args.step_n_schedule[-1]}" + ) + args.dp_size = dist.get_world_size() // args.tp_size + args.target_batch_size = args.tp_size * args.batch_size + if args.kl_scale < 0: + raise ValueError(f"--kl-scale must be non-negative, got {args.kl_scale}") + if args.kl_decay < 0: + raise ValueError(f"--kl-decay must be non-negative, got {args.kl_decay}") + if args.attention_backend == "usp": + sp_sanity_check(args) + + +def sp_sanity_check(args: Namespace) -> None: + args.draft_accumulation_steps = ( + args.draft_accumulation_steps * args.sp_ulysses_size * args.sp_ring_size + ) + assert ( + args.batch_size == 1 + ), f"USP only supports batch_size=1, got batch_size={args.batch_size}" + + assert args.sp_ring_size * args.sp_ulysses_size > 1, ( + f"USP requires sp_ring_size * sp_ulysses_size > 1. " + f"Got sp_ring_size={args.sp_ring_size}, sp_ulysses_size={args.sp_ulysses_size}." + ) + + assert args.train_hidden_states_path is not None, f"USP only support offline mode" + + if args.eval_data_path is not None and args.eval_hidden_states_path is not None: + raise ValueError( + "Cannot set both eval_data_path and eval_hidden_states_path. " + "For online mode, set only eval_data_path. " + "For offline mode, set only eval_hidden_states_path." + ) + + +def build_draft_model(args: Namespace) -> Tuple[AutoDraftModelConfig, nn.Module]: + # ckpt info(epoch, step) + ckpt_info = (0, 0) + + # Handle draft model config + if args.draft_model_config is None: + # Auto-generate and save config file + 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) + else: + # Use provided config file + draft_model_config = AutoDraftModelConfig.from_file(args.draft_model_config) + + # Handle base ckpt, config file + 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_last_checkpoint = args.ckpt_dir + 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." + ) + + # detecting last ckpt for draft model + if args.resume and os.path.isdir(args.output_dir): + print_on_rank0(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 + + # EDSD uses its own draft model class (EdsdDraftModel). We construct it + # directly rather than via AutoEagle3DraftModel, whose single config->class + # mapping is already occupied by LlamaForCausalLMEagle3. + # Tag the config so the saved config.json advertises EdsdDraftModel (useful + # for serving / re-loading), regardless of how it was generated. + draft_model_config.architectures = ["EdsdDraftModel"] + if draft_model_last_checkpoint: + draft_model = EdsdDraftModel.from_pretrained( + draft_model_last_checkpoint, + attention_backend=args.attention_backend, + torch_dtype=torch.bfloat16, + ).cuda() + else: + draft_model = EdsdDraftModel( + draft_model_config, + attention_backend=args.attention_backend, + ) + draft_model = draft_model.to(dtype=torch.bfloat16).cuda() + + # Load training state (optimizer, scheduler, epoch, step) for true resume + resume_state = None + 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']}" + ) + + draft_model.load_embedding(args.target_model_path, embedding_key=args.embedding_key) + draft_model.freeze_embedding() + return draft_model_config, draft_model, ckpt_info, resume_state + + +def build_dataloaders( + args: Namespace, + draft_model_config: AutoDraftModelConfig, +) -> Tuple[DataLoader, str, Optional[DataLoader]]: + # build dataloaders + tokenizer = AutoTokenizer.from_pretrained( + args.target_model_path, trust_remote_code=args.trust_remote_code + ) + + # convert to dataloader + dataset_cache_params_string = ( + f"{args.train_data_path}-" + f"{args.max_length}-" + f"{args.chat_template}-" + f"{args.target_model_path}" # Tokenizer may also different + ) + vocab_cache_params_string = ( + f"{dataset_cache_params_string}-" + f"{draft_model_config.draft_vocab_size}-" + f"{draft_model_config.vocab_size}" + ) + cache_key = hashlib.md5(dataset_cache_params_string.encode()).hexdigest() + vocab_cache_key = hashlib.md5(vocab_cache_params_string.encode()).hexdigest() + train_dataset = Dataset.from_generator( + generator=safe_conversations_generator, + gen_kwargs={"file_path": args.train_data_path}, + ) + is_online = ( + args.train_data_path is not None and args.train_hidden_states_path is None + ) + 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_preformatted=args.is_preformatted, + num_proc=args.build_dataset_num_proc, + train_only_last_turn=args.train_only_last_turn, + ) + vocab_mapping_path = generate_vocab_mapping_file( + dataset=train_eagle3_dataset, + target_vocab_size=draft_model_config.vocab_size, + draft_vocab_size=draft_model_config.draft_vocab_size, + cache_dir=os.path.join(args.cache_dir, "vocab_mapping"), + cache_key=vocab_cache_key, + ) + + if not is_online: + train_eagle3_dataset = build_offline_eagle3_dataset( + args.train_hidden_states_path, + args.max_length, + ttt_length=args.ttt_length, + use_usp_preprocess=(args.attention_backend == "usp"), + ) + + train_dataloader = prepare_dp_dataloaders( + train_eagle3_dataset, + args.target_batch_size, + num_workers=args.dataloader_num_workers, + shuffle=True, + process_group=( + get_draft_dp_group() + if args.attention_backend == "usp" and not is_online + else get_dp_group() + ), + ) + if args.eval_data_path is not None or args.eval_hidden_states_path is not None: + if args.eval_data_path is not None: + eval_dataset = Dataset.from_generator( + generator=safe_conversations_generator, + gen_kwargs={"file_path": args.eval_data_path}, + ) + eval_eagle3_dataset = build_eagle3_dataset( + eval_dataset, + tokenizer, + args.chat_template, + args.max_length, + num_proc=args.build_dataset_num_proc, + is_preformatted=args.is_preformatted, + train_only_last_turn=args.train_only_last_turn, + ) + elif args.eval_hidden_states_path is not None: + eval_eagle3_dataset = build_offline_eagle3_dataset( + args.eval_hidden_states_path, + args.max_length, + ttt_length=args.ttt_length, + use_usp_preprocess=(args.attention_backend == "usp"), + ) + eval_dataloader = prepare_dp_dataloaders( + eval_eagle3_dataset, + args.target_batch_size, + num_workers=args.dataloader_num_workers, + shuffle=False, + process_group=( + get_draft_dp_group() + if args.attention_backend == "usp" and not is_online + else get_dp_group() + ), + ) + print_with_rank("Initialized eval dataloader") + else: + eval_dataloader = None + return ( + train_dataloader, + vocab_mapping_path, + eval_dataloader, + ) + + +def filter_draft_state_dict(model_state_dict: dict) -> dict: + """Keep only draft-model weights for the serving checkpoint (drop embeddings). + + Embeddings are intentionally excluded (SGLang loads them from the target); the + target ``TargetHead`` is never part of the wrapped draft model, so no teacher state + can leak. + """ + return { + k.replace("draft_model.", ""): v + for k, v in model_state_dict.items() + if "draft_model." in k and "embed" not in k.lower() + } + + +def save_checkpoints( + args: Namespace, + epoch: int, + step: int, + eagle3_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(eagle3_model, StateDictType.FULL_STATE_DICT): + model_state_dict = eagle3_model.state_dict() + state_to_save = { + "epoch": epoch, + "global_step": step, + "args": args, + } + state_to_save.update(optimizer.state_dict()) + draft_model_state_dict = filter_draft_state_dict(model_state_dict) + + if dist.get_rank() == 0: + torch.save( + state_to_save, + os.path.join(epoch_output_dir, "training_state.pt"), + ) + print_on_rank0( + f"Saved full training state to {epoch_output_dir}/training_state.pt" + ) + eagle3_model.draft_model.save_pretrained( + epoch_output_dir, + state_dict=draft_model_state_dict, + ) + print_on_rank0(f"Saved model configuration to {epoch_output_dir}") + dist.barrier() + + +def run_forward( + args: Namespace, + eagle3_model: nn.Module, + data: dict, + target_model: Optional[Eagle3TargetModel] = None, + is_online: bool = True, + epoch: int = 0, + total_epochs: int = 1, +) -> Tuple[ + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], +]: + if is_online: + # generate eagle3 target data from the target model (text-only) + 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 = eagle3_data.input_ids + attention_mask = eagle3_data.attention_mask + loss_mask = eagle3_data.loss_mask + target = eagle3_data.target + hidden_states = eagle3_data.hidden_states + else: + # we generate the logits using the hidden states loaded from disk + attention_mask = data["attention_mask"].cuda() + hidden_states = data["hidden_state"].cuda() + input_ids, target, loss_mask = target_model.preprocess( + data["input_ids"], data["target"], data["loss_mask"] + ) + input_ids = input_ids.cuda() + loss_mask = loss_mask.cuda() + target = target_model(target.cuda()) + ( + plosses, + acceptance_rates, + acces, + acc_corrects, + acc_denoms, + metric_losses, + metric_loss_denoms, + ) = eagle3_model( + input_ids=input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + target=target, + hidden_states=hidden_states, + position_ids=( + data["position_ids"].cuda() if "position_ids" in data else None + ), + is_vlm=False, + epoch_idx=epoch, + total_epochs=total_epochs, + ) + return ( + plosses, + acceptance_rates, + acces, + acc_corrects, + acc_denoms, + metric_losses, + metric_loss_denoms, + ) + + +def run_backward_and_update( + args: Namespace, plosses: List[torch.Tensor], optimizer: Optimizer, global_step: int +) -> Optional[torch.Tensor]: + ploss_weight = [0.8**i for i in range(len(plosses))] + ploss = ( + sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))]) + / args.draft_accumulation_steps + ) + ploss.backward() + + if global_step % args.draft_accumulation_steps == 0: + grad_norm = optimizer.step() + if dist.is_initialized(): + grad_norm = grad_norm.detach().float() + if torch.cuda.is_available(): + grad_norm = grad_norm.to(torch.cuda.current_device()) + grad_norm = grad_norm.pow(2) + dist.all_reduce(grad_norm, op=dist.ReduceOp.SUM) + grad_norm = grad_norm.sqrt() + return grad_norm + return None + + +def record_metrics( + args: Namespace, + accuracies: List[torch.Tensor], + acceptance_rates: List[torch.Tensor], + plosses: List[torch.Tensor], + global_step: int, + tracker: Tracker, + optimizer: Optional[Optimizer] = None, + mode: str = "train", + acc_corrects: Optional[List[torch.Tensor]] = None, + acc_denoms: Optional[List[torch.Tensor]] = None, + ploss_denoms: Optional[List[torch.Tensor]] = None, + grad_norms: Optional[List[torch.Tensor]] = None, +) -> None: + logdict = {} + + if mode == "train" and optimizer is not None: + logdict["train/lr"] = optimizer.get_learning_rate() + + plosses = torch.stack(plosses) + + if acc_corrects is not None and acc_denoms is not None: + corrects = torch.stack(acc_corrects) + denoms = torch.stack(acc_denoms) + assert corrects.shape[0] >= 1, "expected at least one TTT position" + dist.all_reduce(corrects, op=dist.ReduceOp.SUM) + dist.all_reduce(denoms, op=dist.ReduceOp.SUM) + accuracies = corrects / denoms.clamp_min(1e-6) + else: + accuracies = torch.stack(accuracies) + assert accuracies.shape[0] >= 1, "expected at least one TTT position" + dist.all_reduce(accuracies, op=dist.ReduceOp.AVG) + + accuracies = accuracies.cpu().tolist() + for i in range(len(accuracies)): + logdict[f"{mode}/acc_{i}"] = accuracies[i] + print_on_rank0( + f"Eval - Step {global_step} [{global_step + 1}/{args.num_epochs}], position {i}, Acc: {accuracies[i]:.2f}" + ) + + acceptance_rates = torch.stack(acceptance_rates) + assert acceptance_rates.shape[0] >= 1, "expected at least one TTT position" + dist.all_reduce(acceptance_rates, op=dist.ReduceOp.AVG) + acceptance_rates = acceptance_rates.cpu().tolist() + for i in range(len(acceptance_rates)): + logdict[f"{mode}/acceptance_rate_{i}"] = acceptance_rates[i] + print_on_rank0( + f"Eval - Step {global_step} [{global_step + 1}/{args.num_epochs}], position {i}, Acceptance Rate: {acceptance_rates[i]:.4f}" + ) + + if ploss_denoms is not None: + ploss_denoms = torch.stack(ploss_denoms) + dist.all_reduce(plosses, op=dist.ReduceOp.SUM) + dist.all_reduce(ploss_denoms, op=dist.ReduceOp.SUM) + plosses = plosses / ploss_denoms.clamp_min(1e-6) + else: + dist.all_reduce(plosses, op=dist.ReduceOp.AVG) + + plosses = plosses.cpu().tolist() + for i in range(len(plosses)): + logdict[f"{mode}/ploss_{i}"] = plosses[i] + print_on_rank0( + f"Eval - Step {global_step} [{global_step + 1}/{args.num_epochs}], position {i}, pLoss: {plosses[i]}" + ) + + if grad_norms: + grad_norm = torch.stack([norm.detach().float() for norm in grad_norms]).mean() + logdict[f"{mode}/grad_norm"] = grad_norm.item() + + tracker.log(logdict, step=global_step) + + +def get_progress_metrics( + acc_corrects: List[torch.Tensor], + acc_denoms: List[torch.Tensor], + metric_losses: List[torch.Tensor], + metric_loss_denoms: List[torch.Tensor], + grad_norm: Optional[torch.Tensor] = None, + last_grad_norm: Optional[float] = None, +) -> Tuple[float, float, Optional[float]]: + loss_num = sum( + loss.detach() * denom.detach() + for loss, denom in zip(metric_losses, metric_loss_denoms) + ) + loss_den = sum(denom.detach() for denom in metric_loss_denoms) + acc_num = sum(correct.detach() for correct in acc_corrects) + acc_den = sum(denom.detach() for denom in acc_denoms) + + metric_values = [ + loss_num.float(), + loss_den.float(), + acc_num.float(), + acc_den.float(), + ] + metrics = torch.stack(metric_values) + dist.all_reduce(metrics, op=dist.ReduceOp.SUM) + loss = metrics[0] / metrics[1].clamp_min(1e-6) + acc = metrics[2] / metrics[3].clamp_min(1e-6) + if grad_norm is not None: + last_grad_norm = grad_norm.detach().float().item() + return loss.item(), acc.item(), last_grad_norm + + +def main(): + # ================================================ + # 1. Initialize + # ================================================ + parser, args = parse_args() + set_seed(args.seed) + init_distributed( + timeout=args.dist_timeout, + tp_size=args.tp_size, + sp_ring_size=args.sp_ring_size, + sp_ulysses_size=args.sp_ulysses_size, + ) + is_online = ( + args.train_data_path is not None and args.train_hidden_states_path is None + ) + + sanity_check(args) + print_args_with_dots(args) + print_with_rank("Initialized distributed environment") + print_cuda_memory_debug("after init_distributed") + + # ================================================ + # 2. Build models + # ================================================ + print_cuda_memory_debug("before build_draft_model") + draft_model_config, draft_model, ckpt_info, resume_state = build_draft_model(args) + print_cuda_memory_debug("after build_draft_model") + print_cuda_memory_debug("before build_target_model") + target_model, _ = build_target_model(args, draft_model_config, is_online) + print_cuda_memory_debug("after build_target_model") + + # ================================================ + # 3. Build dataloader + # ================================================ + print_cuda_memory_debug("before build_dataloaders") + train_dataloader, vocab_mapping_path, eval_dataloader = build_dataloaders( + args, draft_model_config + ) + print_cuda_memory_debug("after build_dataloaders") + + # we load the vocab mapping then + draft_model.load_vocab_mapping(vocab_mapping_path) + print_with_rank("Loaded vocab mapping") + print_cuda_memory_debug("after load_vocab_mapping") + + # Calculate total steps if not provided + 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} (num_epochs={args.num_epochs} * steps_per_epoch={steps_per_epoch})" + ) + else: + print_with_rank(f"Using provided total_steps: {args.total_steps}") + + # ================================================ + # 4. Build Eagle3 model + # ================================================ + if is_online: + eagle3_model = OnlineEdsdModel( + target_model=target_model, + draft_model=draft_model, + length=args.ttt_length, + attention_backend=args.attention_backend, + lk_loss_type=args.lk_loss_type, + kl_scale=args.kl_scale, + kl_decay=args.kl_decay, + drop_ratio_scale=args.drop_ratio_scale, + step_n_schedule=args.step_n_schedule, + ) + else: + # offline: the target_model is TargetHead, not a full model + eagle3_model = OnlineEdsdModel( + draft_model=draft_model, + length=args.ttt_length, + attention_backend=args.attention_backend, + lk_loss_type=args.lk_loss_type, + kl_scale=args.kl_scale, + kl_decay=args.kl_decay, + drop_ratio_scale=args.drop_ratio_scale, + step_n_schedule=args.step_n_schedule, + ) + eagle3_model = FSDP( + eagle3_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, # the draft model should run dp for all processes + ) + print_with_rank("Initialized Eagle3 FSDP model") + + # ================================================ + # 5. Build optimizer and scheduler + # ================================================ + optimizer = BF16Optimizer( + draft_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") + + # Restore optimizer/scheduler state for true resume + 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] + + # Calculate how many steps to skip in the current epoch (for dataloader fast-forward) + skip_steps = global_step - start_epoch * len(train_dataloader) + + # ================================================ + # 6. Build tracker + # ================================================ + tracker = build_tracker(args, parser) + dist.barrier() + + last_time = time.time() + metric_correct_sums = None + metric_denom_sums = None + metric_loss_weighted_sums = None + metric_loss_denom_sums = None + grad_norms = [] + last_grad_norm = None + + # ================================================ + # 7. Start training + # ================================================ + print_on_rank0( + f"Starting training from epoch:{start_epoch} step:{global_step}" + ) + + # ================================================ + # Worst-case memory probe before training begins. + # Forces actual_length = --ttt-length (the final-epoch curriculum value), + # which is the maximum any real training step will ever reach. If this + # single forward+backward fits in memory, no subsequent epoch can OOM on + # the actual_length axis. If it OOMs here, the user can shrink batch/seq + # before wasting hours of training. Probe grads are discarded (no + # optimizer.step), so the real first step starts clean. + # ================================================ + probe_iter = iter(train_dataloader) + try: + probe_data = next(probe_iter) + except StopIteration: + probe_iter = None + probe_data = None + + if probe_data is not None: + draft_model.train() + print_on_rank0("Full memory probe at max TTT length...") + try: + ( + probe_plosses, + *_, + ) = run_forward( + args, + eagle3_model, + probe_data, + target_model, + is_online, + epoch=args.num_epochs - 1, # force actual_length = S_max + total_epochs=args.num_epochs, + ) + probe_loss = sum(probe_plosses) / len(probe_plosses) + probe_loss.backward() + except torch.OutOfMemoryError as e: + torch.cuda.empty_cache() + raise torch.OutOfMemoryError( + "Worst-case memory probe failed at max TTT length. " + "Reduce batch size or sequence length before launching training." + ) from e + finally: + # Drop the probe's gradients so the real first step starts clean. + # BF16Optimizer only clears grads inside .step(); the probe must not + # rely on it, and must never call .step(). + with torch.no_grad(): + for p in optimizer.model_params: + p.grad = None + del probe_plosses, probe_loss + torch.cuda.empty_cache() + print_on_rank0("Memory probe passed. Starting training.") + + for epoch in range(start_epoch, args.num_epochs): + # Run training + 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): + # Skip steps already processed in the current epoch when resuming + if epoch == start_epoch and step_in_epoch < skip_steps: + continue + + global_step += 1 + + # ================================================ + # 7.0 Profiling + # ================================================ + if args.profile: + # we add the step by 1 to align with global step + 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{torch.distributed.get_rank()}_{time.time()}.trace.json.gz", + ) + print(f"End profile {output_path=}") + torch_profiler.stop() + torch_profiler.export_chrome_trace(output_path) + + # ================================================ + # 7.1 Training Step + # ================================================ + ( + plosses, + acceptance_rates, + acces, + acc_corrects, + acc_denoms, + metric_losses, + metric_loss_denoms, + ) = run_forward( + args, + eagle3_model, + data, + target_model, + is_online, + epoch=epoch, + total_epochs=args.num_epochs, + ) + grad_norm = run_backward_and_update(args, plosses, optimizer, global_step) + if grad_norm is not None: + grad_norms.append(grad_norm) + + with torch.no_grad(): + # EDSD's TTT length can change across an epoch boundary within a + # single logging interval; reset interval accumulators if the + # current step's length no longer matches the accumulated one. + if ( + metric_correct_sums is not None + and len(metric_correct_sums) != len(acc_corrects) + ): + metric_correct_sums = None + metric_denom_sums = None + metric_loss_weighted_sums = None + metric_loss_denom_sums = None + if metric_correct_sums is None: + metric_correct_sums = [ + correct.detach().clone() for correct in acc_corrects + ] + metric_denom_sums = [denom.detach().clone() for denom in acc_denoms] + metric_loss_weighted_sums = [ + loss.detach() * denom.detach() + for loss, denom in zip(metric_losses, metric_loss_denoms) + ] + metric_loss_denom_sums = [ + denom.detach().clone() for denom in metric_loss_denoms + ] + else: + for i in range(len(acc_corrects)): + metric_correct_sums[i] += acc_corrects[i].detach() + metric_denom_sums[i] += acc_denoms[i].detach() + metric_loss_weighted_sums[i] += ( + metric_losses[i].detach() * metric_loss_denoms[i].detach() + ) + metric_loss_denom_sums[i] += metric_loss_denoms[i].detach() + + # log training metrics + if global_step % (args.log_interval * args.draft_accumulation_steps) == 0: + record_metrics( + args, + acces, + acceptance_rates, + metric_loss_weighted_sums, + global_step // args.draft_accumulation_steps, + tracker, + optimizer, + mode="train", + acc_corrects=metric_correct_sums, + acc_denoms=metric_denom_sums, + ploss_denoms=metric_loss_denom_sums, + grad_norms=grad_norms, + ) + metric_correct_sums = None + metric_denom_sums = None + metric_loss_weighted_sums = None + metric_loss_denom_sums = None + grad_norms = [] + + avg_loss, avg_acc, last_grad_norm = get_progress_metrics( + acc_corrects, + acc_denoms, + metric_losses, + metric_loss_denoms, + grad_norm, + last_grad_norm, + ) + if dist.get_rank() == 0: + time_per_step = time.time() - last_time + last_time = time.time() + avg_acceptance_rate = sum(ar for ar in acceptance_rates) / len( + acceptance_rates + ) + postfix = { + "loss": f"{avg_loss:.2f}", + "acc": f"{avg_acc:.2f}", + "acceptance_rate": f"{avg_acceptance_rate:.2f}", + "time": f"{time_per_step:.2f}s", + } + if last_grad_norm is not None: + postfix["grad_norm"] = f"{last_grad_norm:.2f}" + progress_bar.set_postfix(postfix) + + # ================================================ + # 7.2 Evaluation Step + # ================================================ + should_evaluate = ( + args.eval_data_path is not None + or args.eval_hidden_states_path is not None + ) + if ( + should_evaluate + and global_step % (args.eval_interval * args.draft_accumulation_steps) + == 0 + ): + # Run evaluation + draft_model.eval() + eval_acces = [[] for _ in range(eagle3_model.length)] + eval_acceptance_rates = [[] for _ in range(eagle3_model.length)] + eval_plosses = [[] for _ in range(eagle3_model.length)] + + for data in tqdm(eval_dataloader, desc=f"Evaluating Epoch {epoch}"): + with torch.no_grad(): + ( + plosses, + acceptance_rates, + acces, + _, + _, + _, + _, + ) = run_forward( + args, + eagle3_model, + data, + target_model, + is_online, + epoch=epoch, + total_epochs=args.num_epochs, + ) + eval_acces = [ + eval_acces[i] + [acces[i]] for i in range(len(acces)) + ] + eval_acceptance_rates = [ + eval_acceptance_rates[i] + [acceptance_rates[i]] + for i in range(len(acceptance_rates)) + ] + eval_plosses = [ + eval_plosses[i] + [plosses[i]] for i in range(len(plosses)) + ] + + # compute average over all minibatches (only positions that + # actually produced outputs, in case the eval TTT length varies) + eval_acces = [torch.stack(acc).mean() for acc in eval_acces if len(acc) > 0] + eval_acceptance_rates = [ + torch.stack(ar).mean() for ar in eval_acceptance_rates if len(ar) > 0 + ] + eval_plosses = [torch.stack(pl).mean() for pl in eval_plosses if len(pl) > 0] + + record_metrics( + args, + eval_acces, + eval_acceptance_rates, + eval_plosses, + global_step // args.draft_accumulation_steps, + tracker, + mode="eval", + ) + draft_model.train() + # ================================================ + # 7.3 Save Checkpoints + # ================================================ + if global_step % (args.save_interval * args.draft_accumulation_steps) == 0: + # Save the model + save_checkpoints(args, epoch, global_step, eagle3_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 + # Save final checkpoint if training ended without saving + 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, eagle3_model, optimizer) + + # Close the tracker + tracker.close() + destroy_distributed() + + +if __name__ == "__main__": + main() diff --git a/specforge/core/__init__.py b/specforge/core/__init__.py index 4d5dcc644..179560718 100644 --- a/specforge/core/__init__.py +++ b/specforge/core/__init__.py @@ -1,10 +1,12 @@ from .dflash import OnlineDFlashModel from .domino import OnlineDominoModel +from .edsd import OnlineEdsdModel from .eagle3 import OnlineEagle3Model, QwenVLOnlineEagle3Model from .peagle import OnlinePEagleModel __all__ = [ "OnlineDFlashModel", + "OnlineEdsdModel", "OnlineDominoModel", "OnlineEagle3Model", "OnlinePEagleModel", diff --git a/specforge/core/edsd.py b/specforge/core/edsd.py new file mode 100644 index 000000000..4365b1c8b --- /dev/null +++ b/specforge/core/edsd.py @@ -0,0 +1,503 @@ +from typing import Callable, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.cache_utils import DynamicCache + +from specforge.core.eagle3 import ( + OnlineEagle3Model, + _compute_loss_and_acceptance_rate, + _compute_metric_counts, +) +from specforge.core.eagle3_adapters import BackendAdapter, SdpaLikeAdapter, UspAdapter +from specforge.core.loss import LogSoftmaxLoss +from specforge.modeling.draft import Eagle3DraftModel +from specforge.utils import padding + +# --------------------------------------------------------------------------- +# Target distribution computation with curriculum learning mask +# --------------------------------------------------------------------------- + + +@torch.compile(dynamic=None) +def _edsd_compute_target_p_core(target, t2d, loss_mask, compute_on_draft=False): + """Core target distribution computation (kept inside torch.compile). + + Same as EAGLE3's ``_compute_target_p`` but with an optional + ``compute_on_draft`` flag to skip the expensive logsumexp when LK loss + is not used. + """ + target_head = target.float() + target_max_token = target_head.argmax(-1) + target_mask = t2d[target_max_token] + target_mask = target_mask[..., None].int() + position_mask = target_mask * loss_mask + draft_target_head = target_head[..., t2d] + target_p = nn.Softmax(dim=2)(draft_target_head) + target_p = target_p.detach() + target_token_ids = target_max_token.detach() + + # Target probabilities on the full vocabulary restricted to draft tokens. + # Expensive: full-vocab logsumexp. Only computed when LK loss is enabled. + target_p_on_draft = None + if compute_on_draft: + target_logsumexp = torch.logsumexp(target_head, dim=-1, keepdim=True) + target_p_on_draft = torch.exp(draft_target_head - target_logsumexp) + target_p_on_draft = target_p_on_draft.detach() + + return target_p, target_p_on_draft, target_token_ids, position_mask + + +def _edsd_apply_curriculum_mask( + target_p, position_mask, epoch_idx, drop_ratio_scale, total_epochs +): + """Apply entropy-based curriculum learning mask (outside torch.compile). + + Drops the top-``drop_ratio`` fraction of valid positions ranked by entropy + (confusing tokens). ``drop_ratio = min(max((total_epochs - epoch_idx) * + drop_ratio_scale, 0.0), 0.4)``: it is largest at epoch 0 and decays to ~0 by + the final epoch, so higher epochs retain more positions. Placed outside + ``torch.compile`` because ``Tensor.item()`` and dynamic-shape ``topk`` + cause graph breaks. + """ + + eps = 1e-12 + entropy = -(target_p * (target_p + eps).log()).sum(dim=-1) + valid_mask_flat = position_mask.squeeze(-1).view(-1).bool() + entropy_flat = entropy.view(-1) + + drop_ratio = min( + max((total_epochs - epoch_idx - 1) * drop_ratio_scale, 0.0), 0.4 + ) + num_valid = valid_mask_flat.sum() + k = int((drop_ratio * num_valid).item()) + confusion_mask_flat = torch.ones_like( + entropy_flat, dtype=position_mask.dtype + ) + + if k > 0 and num_valid > 0: + valid_indices = valid_mask_flat.nonzero(as_tuple=False).squeeze(-1) + valid_entropy = entropy_flat[valid_indices] + _, topk_rel_indices = torch.topk(valid_entropy, k) + drop_indices = valid_indices[topk_rel_indices] + confusion_mask_flat[drop_indices] = 0 + + confusion_mask = confusion_mask_flat.view_as(entropy)[..., None] + return position_mask * confusion_mask + + +def _edsd_compute_target_p_padded( + target, t2d, loss_mask, length, epoch_idx=0, compute_on_draft=False, + drop_ratio_scale=0.02, total_epochs=1, +): + """Pad target distributions for TTT unrolling (EDSD variant). + + Same as EAGLE3's ``_compute_target_p_padded`` but applies EDSD's + curriculum learning mask after the compiled core computation. + """ + with torch.no_grad(): + ( + target_p, + target_p_on_draft, + target_token_ids, + position_mask, + ) = _edsd_compute_target_p_core( + target=target, + t2d=t2d, + loss_mask=loss_mask, + compute_on_draft=compute_on_draft, + ) + + # Curriculum mask is applied outside torch.compile to avoid + # graph breaks from Tensor.item() and dynamic-shape topk. + position_mask = _edsd_apply_curriculum_mask( + target_p, position_mask, epoch_idx, drop_ratio_scale, total_epochs + ) + + assert len(target_p.shape) == 3 + target_p_padded = F.pad( + target_p, + pad=(0, 0, 0, length), + mode="constant", + value=1 / target_p.shape[-1], + ) + + if target_p_on_draft is not None: + target_p_on_draft_padded = F.pad( + target_p_on_draft, + pad=(0, 0, 0, length), + mode="constant", + value=0.0, + ) + else: + target_p_on_draft_padded = None + + target_token_ids_padded = F.pad( + target_token_ids, + pad=(0, length), + mode="constant", + value=0, + ) + + return ( + target_p_padded, + target_p_on_draft_padded, + target_token_ids_padded, + position_mask, + ) + + +# --------------------------------------------------------------------------- +# Online EDSD Model +# --------------------------------------------------------------------------- + + +class OnlineEdsdModel(OnlineEagle3Model): + """EDSD online training wrapper. + + Inherits from ``OnlineEagle3Model`` and overrides: + + 1. ``forward`` — TTT length grows with ``epoch_idx``; uses EDSD's + curriculum-learning-aware ``_compute_target_p_padded``. + 2. ``_acc_and_loss`` — always computes ``acceptance_rate`` as a + monitoring metric (matching EAGLE3 behaviour). + + All other logic (adapter, position_ids, attention mask, padding, etc.) + is inherited from EAGLE3. + """ + + def __init__( + self, + draft_model: Eagle3DraftModel, + length: int = 7, + attention_backend: str = "sdpa", + target_model: Optional[nn.Module] = None, + lk_loss_type: Optional[str] = None, + kl_scale: float = 1.0, + kl_decay: float = 1.0, + drop_ratio_scale: float = 0.02, + step_n_schedule: Optional[List[int]] = None, + ): + super().__init__( + draft_model=draft_model, + length=length, + attention_backend=attention_backend, + target_model=target_model, + lk_loss_type=lk_loss_type, + kl_scale=kl_scale, + kl_decay=kl_decay, + ) + # Always compute target_p_on_draft so acceptance_rate can be + # reported as a monitoring metric even when LK loss is off. + self._compute_on_draft = True + self.drop_ratio_scale = drop_ratio_scale + self.step_n_schedule = step_n_schedule + + @staticmethod + def compute_step_n( + current_epoch: int, + total_epochs: int, + s_max: int = 7, + ) -> int: + """Compute the Step-n TTT length for the current epoch. + + From the EDSD paper (Eq. 6): + n_t = ceil(1 + (S_max - 1) * t / (T - 1)) + + At epoch 0: n_0 = 1 (minimum TTT length) + At epoch T-1: n_{T-1} = S_max (maximum TTT length) + + Args: + current_epoch: Current training epoch (0-indexed). + total_epochs: Total number of training epochs T. + s_max: Maximum simulation steps (``self.length``). + + Returns: + TTT length n_t for the current epoch. + """ + import math + if total_epochs <= 1: + return s_max + return math.ceil(1 + (s_max - 1) * current_epoch / (total_epochs - 1)) + + def _compute_actual_length(self, epoch_idx: int, total_epochs: int) -> int: + """Compute the actual TTT unroll length for the given epoch. + + If ``step_n_schedule`` is provided, it takes priority: the value at + ``epoch_idx`` is used directly (clamped to the last value if out of range). + Otherwise, delegates to ``compute_step_n`` (EDSD paper Eq. 6): + n_t = ceil(1 + (S_max - 1) * t / (T - 1)) + + At epoch 0: n_0 = 1 (minimum TTT length) + At epoch T-1: n_{T-1} = S_max (maximum TTT length) + """ + if self.step_n_schedule is not None: + if epoch_idx < len(self.step_n_schedule): + return self.step_n_schedule[epoch_idx] + return self.step_n_schedule[-1] + return self.compute_step_n( + current_epoch=epoch_idx, + total_epochs=total_epochs, + s_max=self.length, + ) + + def _acc_and_loss( + self, + *, + logits: torch.Tensor, + target_p: torch.Tensor, + target_p_on_draft: Optional[torch.Tensor], + target_token_ids: torch.Tensor, + position_mask: torch.Tensor, + loss_mask: torch.Tensor, + adapter: BackendAdapter, + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + """Compute accuracy metric, acceptance_rate, and loss.""" + with torch.no_grad(): + pred_draft_token_ids = logits.argmax(-1) + pred_target_token_ids = ( + pred_draft_token_ids + self.draft_model.d2t[pred_draft_token_ids] + ) + local_correct = ( + (pred_target_token_ids == target_token_ids) * loss_mask.squeeze(-1) + ).sum() + local_denom = loss_mask.sum().clamp_min(1e-6) + local_correct, local_denom = adapter.reduce_metrics( + local_correct=local_correct, local_denom=local_denom + ) + acc = local_correct / local_denom + + # Always compute acceptance_rate (as a monitoring metric) and loss. + # When lk_loss_type is None, acceptance_rate is computed with + # gradients disabled inside _compute_loss_and_acceptance_rate. + acceptance_rate, loss = _compute_loss_and_acceptance_rate( + logits=logits, + target_p=target_p, + target_p_on_draft=target_p_on_draft, + position_mask=position_mask, + lk_loss_type=self.lk_loss_type, + kl_scale=self.kl_scale, + kl_decay=self.kl_decay, + reduce_metrics_fn=adapter.reduce_metrics, + reduce_loss_fn=adapter.reduce_loss, + ) + + loss_denom = torch.tensor( + logits.shape[0] * logits.shape[1], + device=logits.device, + dtype=torch.float32, + ) + return ( + acc, + acceptance_rate, + loss, + local_correct, + local_denom, + loss.detach(), + loss_denom, + ) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + target: torch.Tensor, + loss_mask: torch.Tensor, + hidden_states: torch.Tensor, + past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_ids: Optional[torch.Tensor] = None, + image_grid_thw: Optional[torch.Tensor] = None, + is_vlm: bool = False, + epoch_idx: int = 0, + total_epochs: int = 1, + **kwargs, + ) -> Tuple[ + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + List[torch.Tensor], + ]: + """EDSD forward with variable TTT length and curriculum learning. + + Args: + ... (same as OnlineEagle3Model.forward) + epoch_idx: Training epoch index (0-indexed). Controls TTT unroll + length via Step-n (Eq. 6) and curriculum masking. At epoch 0, + n=1 with no masking; at epoch T-1, n=S_max. + total_epochs: Total number of training epochs. Used by Step-n + scheduling when ``step_n_schedule`` is not provided. + """ + # Step 1: precompute the EDSD variable TTT length (Step-n). This must + # happen before Step 2 because target_p padding depends on it. + actual_length = self._compute_actual_length(epoch_idx, total_epochs) + + # Step 2: handle vocab size (with EDSD curriculum-aware target_p) + ( + target_p_padded, + target_p_on_draft_padded, + target_token_ids_padded, + position_mask, + ) = _edsd_compute_target_p_padded( + target=target, + t2d=self.draft_model.t2d, + loss_mask=loss_mask, + length=actual_length, + epoch_idx=epoch_idx, + compute_on_draft=self._compute_on_draft, + drop_ratio_scale=self.drop_ratio_scale, + total_epochs=total_epochs, + ) + del target + torch.cuda.empty_cache() + + # basic info + batch_size, seq_length, _ = hidden_states.shape + seq_length_with_past = seq_length + past_key_values_length = 0 + + # Step 3: project the concatenated hidden states to the target hidden size + hidden_states = self.draft_model.project_hidden_states(hidden_states) + + # Step 4: process kv cache, position ids + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + position_ids = self._prepare_position_ids( + position_ids=position_ids, + seq_length=seq_length, + past_key_values_length=past_key_values_length, + device=hidden_states.device, + is_vlm=is_vlm, + input_ids=input_ids, + image_grid_thw=image_grid_thw, + ) + + # Step 5: handle attention mask + if attention_mask is None: + attention_mask = torch.ones( + (batch_size, seq_length_with_past), + dtype=torch.bool, + device=hidden_states.device, + ) + if self.attention_backend == "sdpa": + attention_mask = self.draft_model.prepare_decoder_attention_mask( + attention_mask=attention_mask, + hidden_states=hidden_states, + batch_size=batch_size, + seq_length=seq_length, + past_key_values_length=past_key_values_length, + ) + + # Step 6: run TTT with variable length + # (actual_length computed above in Step 1) + + plosses = [] + acceptance_rates = [] + acces = [] + metric_corrects = [] + metric_denoms = [] + metric_losses = [] + metric_loss_denoms = [] + adapter = self._make_adapter() + global_input_ids = input_ids + if self.attention_backend in ["sdpa", "fa", "usp"]: + cache_hidden = [[], []] + past_key_values = None + elif self.attention_backend == "flex_attention": + cache_hidden = None + past_key_values = DynamicCache() + else: + raise ValueError(f"Unknown attention backend: {self.attention_backend}") + + for idx in range(actual_length): + state = adapter.step_view( + idx=idx, + ttt_length=actual_length, + global_input_ids=global_input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + position_ids=position_ids, + hidden_states=hidden_states, + target_p_padded=target_p_padded, + target_p_on_draft_padded=target_p_on_draft_padded, + target_token_ids_padded=target_token_ids_padded, + position_mask=position_mask, + seq_length=seq_length, + ) + is_last = idx == actual_length - 1 + + # Step 6.1: embed the input ids + inputs_embeds = self.draft_model.embed_input_ids(state.input_ids) + inputs_embeds = inputs_embeds.to(hidden_states.dtype) + + # Step 6.2: run the draft model backbone + hidden_states_out = self.draft_model.backbone( + input_embeds=inputs_embeds, + hidden_states=state.hidden_states, + cache_hidden=cache_hidden, + attention_mask=state.attention_mask, + position_ids=state.position_ids, + past_key_values=past_key_values, + use_cache=True, + ) + + # update hidden states for next step + hidden_states = hidden_states_out + + # Step 6.4: get logits + logits = self.draft_model.compute_logits(hidden_states) + + # Step 6.5 + 6.6: metric and loss + ( + acc, + acceptance_rate, + loss, + correct, + denom, + metric_loss, + loss_denom, + ) = self._acc_and_loss( + logits=logits, + target_p=state.target_p, + target_p_on_draft=state.target_p_on_draft, + target_token_ids=state.target_token_ids, + position_mask=state.position_mask, + loss_mask=state.loss_mask, + adapter=adapter, + ) + acces.append(acc) + acceptance_rates.append(acceptance_rate) + plosses.append(loss) + metric_corrects.append(correct) + metric_denoms.append(denom) + metric_losses.append(metric_loss) + metric_loss_denoms.append(loss_denom) + + if not is_last: + # Step 6.7: we need to update the loss mask + global_input_ids = padding(global_input_ids, left=False) + position_mask = padding(position_mask, left=False) + loss_mask = padding(loss_mask, left=False) + # Flex attention mask shrinking is handled inside attention module + + return ( + plosses, + acceptance_rates, + acces, + metric_corrects, + metric_denoms, + metric_losses, + metric_loss_denoms, + ) diff --git a/specforge/modeling/auto.py b/specforge/modeling/auto.py index d52759b4e..4684c41f1 100644 --- a/specforge/modeling/auto.py +++ b/specforge/modeling/auto.py @@ -133,6 +133,7 @@ class AutoDraftModelConfig: _config_mapping = { "LlamaForCausalLMEagle3": LlamaConfig, + "EdsdDraftModel": LlamaConfig, "PEagleDraftModel": LlamaConfig, } diff --git a/specforge/modeling/draft/__init__.py b/specforge/modeling/draft/__init__.py index 6130dcc63..baf6d686d 100644 --- a/specforge/modeling/draft/__init__.py +++ b/specforge/modeling/draft/__init__.py @@ -5,12 +5,14 @@ extract_context_feature, sample, ) +from .edsd import EdsdDraftModel from .llama3_eagle import LlamaForCausalLMEagle3 from .peagle import PEagleDraftModel __all__ = [ "Eagle3DraftModel", "DFlashDraftModel", + "EdsdDraftModel", "LlamaForCausalLMEagle3", "PEagleDraftModel", "build_target_layer_ids", diff --git a/specforge/modeling/draft/edsd.py b/specforge/modeling/draft/edsd.py new file mode 100644 index 000000000..0db11d53e --- /dev/null +++ b/specforge/modeling/draft/edsd.py @@ -0,0 +1,133 @@ +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 .llama3_eagle import LlamaDecoderLayer, LlamaForCausalLMEagle3, LlamaMLP, LlamaRMSNorm + + +class EDFuse(LlamaMLP): + + def forward(self, x: torch.Tensor, embed_tokens: torch.Tensor) -> torch.Tensor: + if self.config.pretraining_tp > 1: + slice_ = self.intermediate_size // self.config.pretraining_tp + gate_proj_slices = self.gate_proj.weight.split(slice_, dim=0) + up_proj_slices = self.up_proj.weight.split(slice_, dim=0) + down_proj_slices = self.down_proj.weight.split(slice_, dim=1) + + gate_proj = torch.cat( + [torch.nn.functional.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], + dim=-1, + ) + up_proj = torch.cat( + [torch.nn.functional.linear(embed_tokens, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], + dim=-1, + ) + intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice_, dim=2) + down_proj = sum( + torch.nn.functional.linear(intermediate_states[i], down_proj_slices[i]) + for i in range(self.config.pretraining_tp) + ) + else: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(embed_tokens)) + + return down_proj + + +class EdsdDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: LlamaConfig, attention_backend: str = "sdpa"): + super().__init__(config, attention_backend=attention_backend) + self.input_layernorm = nn.Identity() + for name in ("q_proj", "k_proj", "v_proj"): + old = getattr(self.self_attn, name) + setattr(self.self_attn, name, nn.Linear(config.hidden_size, old.out_features, bias=False)) + + def forward( + self, + input_emb, + hidden_states: torch.Tensor, + cache_hidden: Optional[List[List[torch.Tensor]]] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + + hidden_states = self.hidden_norm(hidden_states) + + hidden_states = self.self_attn( + cache_hidden=cache_hidden, + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + 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 EdsdDraftModel(LlamaForCausalLMEagle3): + + config_class = LlamaConfig + + def _build_midlayer(self, config, attention_backend): + return EdsdDecoderLayer(config, attention_backend=attention_backend) + + def _build_fc(self, config): + if hasattr(config, "target_hidden_size"): + return torch.nn.Linear( + config.target_hidden_size * 2, config.hidden_size, bias=False + ) + return torch.nn.Linear( + config.hidden_size * 2, config.hidden_size, bias=False + ) + + def __init__(self, config, quant_config=None, attention_backend="sdpa") -> None: + assert hasattr(config, "target_layer_ids") and len(config.target_layer_ids) == 2, \ + f"EdsdDraftModel requires exactly 2 target layers, got {getattr(config, 'target_layer_ids', None)}" + + super().__init__(config, quant_config=quant_config, attention_backend=attention_backend) + + self.edfuse = EDFuse(config) + self.embnorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + + def project_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.size(-1) == self.config.hidden_size * 2 + return self.fc(hidden_states) + + def backbone( + self, + input_embeds: torch.Tensor, + hidden_states: torch.Tensor, + cache_hidden: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + past_key_values: Optional[Cache] = None, + use_cache: bool = True, + ) -> torch.Tensor: + input_embeds = self.embnorm(input_embeds) + hidden_states = self.edfuse(hidden_states, input_embeds) + return self.midlayer( + input_emb=None, + hidden_states=hidden_states, + cache_hidden=cache_hidden, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=False, + use_cache=False, + ) + diff --git a/specforge/modeling/draft/llama3_eagle.py b/specforge/modeling/draft/llama3_eagle.py index e96286bdf..b487d6aa7 100644 --- a/specforge/modeling/draft/llama3_eagle.py +++ b/specforge/modeling/draft/llama3_eagle.py @@ -1628,6 +1628,18 @@ class LlamaForCausalLMEagle3(Eagle3DraftModel): config_class = LlamaConfig + def _build_midlayer(self, config, attention_backend): + return LlamaDecoderLayer(config, attention_backend=attention_backend) + + def _build_fc(self, config): + if hasattr(config, "target_hidden_size"): + return torch.nn.Linear( + config.target_hidden_size * 3, config.hidden_size, bias=False + ) + return torch.nn.Linear( + config.hidden_size * 3, config.hidden_size, bias=False + ) + def __init__(self, config, quant_config=None, attention_backend="sdpa") -> None: super().__init__(config) self.config = config @@ -1638,16 +1650,8 @@ def __init__(self, config, quant_config=None, attention_backend="sdpa") -> None: self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, config.pad_token_id ) - self.midlayer = LlamaDecoderLayer(config, attention_backend=attention_backend) - - if hasattr(config, "target_hidden_size"): - self.fc = torch.nn.Linear( - config.target_hidden_size * 3, config.hidden_size, bias=False - ) - else: - self.fc = torch.nn.Linear( - config.hidden_size * 3, config.hidden_size, bias=False - ) + self.midlayer = self._build_midlayer(config, attention_backend) + self.fc = self._build_fc(config) self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.lm_head = nn.Linear( diff --git a/specforge/modeling/target/eagle3_target_model.py b/specforge/modeling/target/eagle3_target_model.py index b89d1010b..78b5d96d4 100644 --- a/specforge/modeling/target/eagle3_target_model.py +++ b/specforge/modeling/target/eagle3_target_model.py @@ -132,9 +132,11 @@ def set_aux_hidden_states_layers( num_layers - 4, ] self.aux_hidden_states_layers = aux_hidden_states_layers - assert ( - len(self.aux_hidden_states_layers) == 3 - ), "aux_hidden_states_layers is expected to be 3 layers for EAGLE3" + if len(aux_hidden_states_layers) not in (2, 3): + raise ValueError( + "Expected 2 target layers for EDSD or " + "3 target layers for EAGLE3." + ) class HFEagle3TargetModel(Eagle3TargetModel): @@ -257,18 +259,22 @@ def hook(module, input, output): handle.remove() # Verify we captured everything - if len(captured_states) != 3: + n_layers = len(self.aux_hidden_states_layers) + if len(captured_states) != n_layers: raise RuntimeError( - f"Expected to capture 3 layers, but captured {len(captured_states)}" + f"Expected to capture {n_layers} layers, but captured {len(captured_states)}" ) # Extract in the correct order - hidden_states0 = captured_states[target_indices[0]] - hidden_states1 = captured_states[target_indices[1]] - hidden_states2 = captured_states[target_indices[2]] + #hidden_states0 = captured_states[target_indices[0]] + #hidden_states1 = captured_states[target_indices[1]] + #hidden_states2 = captured_states[target_indices[2]] + #hidden_states = torch.cat( + # (hidden_states0, hidden_states1, hidden_states2), dim=-1 + #) hidden_states = torch.cat( - (hidden_states0, hidden_states1, hidden_states2), dim=-1 + [captured_states[target_indices[i]] for i in range(n_layers)], dim=-1 ) # apply pading diff --git a/tests/test_core/test_edsd_core.py b/tests/test_core/test_edsd_core.py new file mode 100644 index 000000000..07cba9b9c --- /dev/null +++ b/tests/test_core/test_edsd_core.py @@ -0,0 +1,215 @@ +import unittest + +import torch + +from specforge.core.edsd import ( + OnlineEdsdModel, + _edsd_apply_curriculum_mask, + _edsd_compute_target_p_padded, +) + + +class TestComputeStepN(unittest.TestCase): + """Tests for OnlineEdsdModel.compute_step_n (EDSD paper Eq. 6).""" + + def test_epoch0_returns_1(self): + """At epoch 0 the TTT length should be 1 (minimum).""" + self.assertEqual(OnlineEdsdModel.compute_step_n(0, 10, s_max=7), 1) + + def test_last_epoch_returns_smax(self): + """At the last epoch the TTT length should equal s_max.""" + self.assertEqual(OnlineEdsdModel.compute_step_n(9, 10, s_max=7), 7) + + def test_single_epoch_returns_smax(self): + """When total_epochs <= 1, always return s_max.""" + self.assertEqual(OnlineEdsdModel.compute_step_n(0, 1, s_max=7), 7) + self.assertEqual(OnlineEdsdModel.compute_step_n(0, 0, s_max=5), 5) + + def test_monotonic_increase(self): + """TTT length should be non-decreasing across epochs.""" + s_max = 7 + total = 10 + lengths = [OnlineEdsdModel.compute_step_n(e, total, s_max) for e in range(total)] + for i in range(len(lengths) - 1): + self.assertGreaterEqual(lengths[i + 1], lengths[i]) + + def test_two_epochs(self): + """With 2 epochs: epoch 0 -> 1, epoch 1 -> s_max.""" + self.assertEqual(OnlineEdsdModel.compute_step_n(0, 2, s_max=7), 1) + self.assertEqual(OnlineEdsdModel.compute_step_n(1, 2, s_max=7), 7) + + def test_smax_1(self): + """When s_max=1, all epochs should return 1.""" + for epoch in range(5): + self.assertEqual(OnlineEdsdModel.compute_step_n(epoch, 5, s_max=1), 1) + + +class TestComputeActualLength(unittest.TestCase): + """Tests for OnlineEdsdModel._compute_actual_length.""" + + def _make_model(self, length=7, step_n_schedule=None): + from specforge.modeling.draft.edsd import EdsdDraftModel + from transformers import LlamaConfig + + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=256, + draft_vocab_size=64, + target_layer_ids=[1, 2], + ) + draft = EdsdDraftModel(config) + return OnlineEdsdModel( + draft_model=draft, + length=length, + step_n_schedule=step_n_schedule, + ) + + def test_schedule_overrides_eq6(self): + """step_n_schedule takes priority over Eq. 6.""" + model = self._make_model(length=7, step_n_schedule=[1, 3, 5, 7]) + self.assertEqual(model._compute_actual_length(0, 10), 1) + self.assertEqual(model._compute_actual_length(1, 10), 3) + self.assertEqual(model._compute_actual_length(2, 10), 5) + self.assertEqual(model._compute_actual_length(3, 10), 7) + + def test_schedule_clamps_to_last(self): + """Out-of-range epochs use the last schedule value.""" + model = self._make_model(length=7, step_n_schedule=[1, 3, 5]) + self.assertEqual(model._compute_actual_length(5, 10), 5) + self.assertEqual(model._compute_actual_length(100, 10), 5) + + def test_no_schedule_uses_eq6(self): + """Without step_n_schedule, delegates to compute_step_n.""" + model = self._make_model(length=7, step_n_schedule=None) + self.assertEqual(model._compute_actual_length(0, 10), 1) + self.assertEqual(model._compute_actual_length(9, 10), 7) + + +class TestCurriculumMask(unittest.TestCase): + """Tests for _edsd_apply_curriculum_mask.""" + + def _make_inputs(self, batch=2, seq=8, draft_vocab=16, num_valid=6): + """Create simple target_p and position_mask tensors.""" + target_p = torch.zeros(batch, seq, draft_vocab) + # Uniform distribution on first num_valid positions + target_p[:, :, :num_valid] = 1.0 / num_valid + position_mask = torch.ones(batch, seq, 1) + return target_p, position_mask + + def test_no_drop_at_high_epoch(self): + """When epoch_idx >= total_epochs, drop_ratio is 0, mask unchanged.""" + target_p, position_mask = self._make_inputs() + result = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=10, drop_ratio_scale=0.02, total_epochs=10) + self.assertTrue(torch.equal(result, position_mask)) + + def test_no_drop_with_zero_scale(self): + """When drop_ratio_scale=0, no positions are dropped.""" + target_p, position_mask = self._make_inputs() + result = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=0, drop_ratio_scale=0.0, total_epochs=10) + self.assertTrue(torch.equal(result, position_mask)) + + def test_drops_some_at_early_epoch(self): + """At epoch 0 with non-zero scale, some valid positions should be dropped.""" + target_p, position_mask = self._make_inputs(num_valid=10) + result = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=0, drop_ratio_scale=0.04, total_epochs=10) + # drop_ratio = (10 - 0) * 0.04 = 0.4, capped at 0.4 + # Some positions should have been zeroed out + num_original = position_mask.sum().item() + num_after = result.sum().item() + self.assertLess(num_after, num_original) + + def test_mask_capped_at_0p4(self): + """Maximum drop ratio is 0.4.""" + target_p, position_mask = self._make_inputs(num_valid=100) + # epoch_idx=0, drop_ratio_scale=1.0 -> (10-0)*1.0 = 10, capped at 0.4 + result = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=0, drop_ratio_scale=1.0, total_epochs=10) + num_original = position_mask.sum().item() + num_after = result.sum().item() + # At most 40% dropped, so at least 60% remain + self.assertGreaterEqual(num_after / num_original, 0.59) + + def test_preserves_zero_positions(self): + """Positions that were already 0 in position_mask stay 0.""" + target_p, position_mask = self._make_inputs() + position_mask[:, 5:, :] = 0 # zero out last 3 positions + result = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=0, drop_ratio_scale=0.04, total_epochs=10) + self.assertTrue(torch.all(result[:, 5:, :] == 0)) + + def test_progressive_unmasking(self): + """Higher epochs should retain more positions than lower epochs.""" + target_p, position_mask = self._make_inputs(batch=1, seq=20, num_valid=20) + # Make entropy vary so curriculum has something to drop + target_p[0, :, :10] = 0.08 + target_p[0, :, 10:] = 0.02 + result_0 = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=0, drop_ratio_scale=0.04, total_epochs=10) + result_5 = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=5, drop_ratio_scale=0.04, total_epochs=10) + result_10 = _edsd_apply_curriculum_mask(target_p, position_mask, epoch_idx=10, drop_ratio_scale=0.04, total_epochs=10) + self.assertLessEqual(result_0.sum().item(), result_5.sum().item()) + self.assertLessEqual(result_5.sum().item(), result_10.sum().item()) + + +class TestEdsdComputeTargetPPadded(unittest.TestCase): + """Tests for _edsd_compute_target_p_padded.""" + + def _make_inputs(self, batch=2, seq=4, vocab=32, draft_vocab=8, length=3): + target = torch.randn(batch, seq, vocab) + t2d = torch.zeros(vocab, dtype=torch.bool) + # Select draft_vocab tokens evenly spaced + indices = torch.linspace(0, vocab - 1, draft_vocab).long() + t2d[indices] = True + loss_mask = torch.ones(batch, seq, 1) + return target, t2d, loss_mask, length + + def test_output_shapes(self): + """Check output tensor shapes match expectations.""" + target, t2d, loss_mask, length = self._make_inputs(draft_vocab=8, length=3) + target_p, target_p_on_draft, target_ids, position_mask = _edsd_compute_target_p_padded( + target, t2d, loss_mask, length, epoch_idx=0, compute_on_draft=True + ) + batch, seq = 2, 4 + self.assertEqual(target_p.shape, (batch, seq + length, 8)) + self.assertEqual(target_p_on_draft.shape, (batch, seq + length, 8)) + self.assertEqual(target_ids.shape, (batch, seq + length)) + self.assertEqual(position_mask.shape, (batch, seq, 1)) + + def test_padding_values(self): + """Padded positions should have uniform target_p and zero target_p_on_draft.""" + target, t2d, loss_mask, length = self._make_inputs(draft_vocab=8, length=3) + target_p, target_p_on_draft, target_ids, _ = _edsd_compute_target_p_padded( + target, t2d, loss_mask, length, epoch_idx=0, compute_on_draft=True + ) + # Padded region: last `length` positions + padded_target_p = target_p[:, -length:, :] + self.assertTrue(torch.allclose(padded_target_p, torch.tensor(1.0 / 8))) + padded_on_draft = target_p_on_draft[:, -length:, :] + self.assertTrue(torch.allclose(padded_on_draft, torch.zeros_like(padded_on_draft))) + + def test_compute_on_draft_false(self): + """When compute_on_draft=False, target_p_on_draft should be None.""" + target, t2d, loss_mask, length = self._make_inputs() + _, target_p_on_draft, _, _ = _edsd_compute_target_p_padded( + target, t2d, loss_mask, length, epoch_idx=0, compute_on_draft=False + ) + self.assertIsNone(target_p_on_draft) + + def test_curriculum_mask_applied(self): + """Curriculum masking should reduce position_mask sum at early epochs.""" + target, t2d, loss_mask, length = self._make_inputs(batch=1, seq=20, vocab=64, draft_vocab=16, length=3) + # epoch_idx=0 with non-zero scale -> some positions dropped + _, _, _, pm_early = _edsd_compute_target_p_padded( + target, t2d, loss_mask, length, epoch_idx=0, drop_ratio_scale=0.04, total_epochs=10 + ) + # epoch_idx=10 -> no dropping + _, _, _, pm_late = _edsd_compute_target_p_padded( + target, t2d, loss_mask, length, epoch_idx=10, drop_ratio_scale=0.04, total_epochs=10 + ) + # Original (non-padded) region only + orig_len = target.shape[1] + self.assertLessEqual(pm_early[:, :orig_len].sum().item(), pm_late[:, :orig_len].sum().item()) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/test_modeling/test_draft/test_edsd.py b/tests/test_modeling/test_draft/test_edsd.py new file mode 100644 index 000000000..99004a198 --- /dev/null +++ b/tests/test_modeling/test_draft/test_edsd.py @@ -0,0 +1,224 @@ +import os +import shutil +import tempfile +import unittest +from unittest.mock import patch + +import torch +from transformers import LlamaConfig + +from specforge.modeling.draft.edsd import EDFuse, EdsdDecoderLayer, EdsdDraftModel +from specforge.modeling.draft.llama3_eagle import LlamaForCausalLMEagle3, LlamaRMSNorm + + +class TestEdsdDraftModel(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + + # Small config mirroring qwen3-8b-edsd.json structure + # target_layer_ids=[1,2] means 2 layers concatenated -> hidden_size * 2 + self.config_dict = { + "architectures": ["EdsdDraftModel"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 64, + "initializer_range": 0.02, + "intermediate_size": 128, + "max_position_embeddings": 512, + "model_type": "llama", + "num_attention_heads": 4, + "num_key_value_heads": 2, + "num_hidden_layers": 1, + "pad_token_id": 0, + "rms_norm_eps": 1e-5, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "use_cache": True, + "vocab_size": 256, + "draft_vocab_size": 64, + "target_layer_ids": [1, 2], + } + self.config = LlamaConfig(**self.config_dict) + + def tearDown(self): + shutil.rmtree(self.temp_dir) + + # ------------------------------------------------------------------ + # Initialization + # ------------------------------------------------------------------ + + def test_model_initialization(self): + model = EdsdDraftModel(self.config) + + # EDSD-specific modules + self.assertIsInstance(model.edfuse, EDFuse) + self.assertIsInstance(model.embnorm, LlamaRMSNorm) + self.assertIsInstance(model.midlayer, EdsdDecoderLayer) + + # Inherited modules + self.assertIsInstance(model.embed_tokens, torch.nn.Embedding) + self.assertIsInstance(model.lm_head, torch.nn.Linear) + self.assertIsInstance(model.norm, LlamaRMSNorm) + + def test_edsd_decoder_layer_qkv_input_dim(self): + """EdsdDecoderLayer replaces q/k/v input dim from hidden_size*2 to hidden_size. + + The output dims are preserved from the original GQA layout: + - q_proj: hidden_size -> num_heads * head_dim + - k_proj: hidden_size -> num_key_value_heads * head_dim + - v_proj: hidden_size -> num_key_value_heads * head_dim + """ + model = EdsdDraftModel(self.config) + attn = model.midlayer.self_attn + + head_dim = self.config.hidden_size // self.config.num_attention_heads + q_out = self.config.num_attention_heads * head_dim + kv_out = self.config.num_key_value_heads * head_dim + + for name, expected_out in [("q_proj", q_out), ("k_proj", kv_out), ("v_proj", kv_out)]: + proj = getattr(attn, name) + self.assertIsInstance(proj, torch.nn.Linear) + self.assertEqual(proj.in_features, self.config.hidden_size) + self.assertEqual(proj.out_features, expected_out) + self.assertIsNone(proj.bias) + + def test_fc_input_dim_matches_target_layer_count(self): + """fc input dim must equal hidden_size * num_target_layers (2 for EDSD).""" + model = EdsdDraftModel(self.config) + # EdsdDraftModel.project_hidden_states asserts hidden_size * 2, + # so fc.in_features must also be hidden_size * 2. + self.assertEqual(model.fc.in_features, self.config.hidden_size * 2) + self.assertEqual(model.fc.out_features, self.config.hidden_size) + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + def test_edfuse_forward(self): + model = EdsdDraftModel(self.config) + model.eval() + batch_size, seq_len = 2, 10 + x = torch.randn(batch_size, seq_len, self.config.hidden_size) + embed_tokens = torch.randn(batch_size, seq_len, self.config.hidden_size) + + with torch.no_grad(): + out = model.edfuse(x, embed_tokens) + + self.assertEqual(out.shape, (batch_size, seq_len, self.config.hidden_size)) + + def test_project_hidden_states(self): + model = EdsdDraftModel(self.config) + model.eval() + batch_size, seq_len = 2, 10 + # EDSD uses 2 target layers -> hidden_size * 2 + hidden_states = torch.randn(batch_size, seq_len, self.config.hidden_size * 2) + + with torch.no_grad(): + projected = model.project_hidden_states(hidden_states) + + self.assertEqual(projected.shape, (batch_size, seq_len, self.config.hidden_size)) + + def test_backbone_forward(self): + model = EdsdDraftModel(self.config) + model.eval() + batch_size, seq_len = 2, 10 + input_embeds = torch.randn(batch_size, seq_len, self.config.hidden_size) + hidden_states = torch.randn(batch_size, seq_len, self.config.hidden_size) + # backbone() passes attention_mask directly to SDPA, which requires + # 4D format (batch, num_heads, q_len, kv_len) or None for causal. + # Passing None triggers is_causal=True in the attention layer. + position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + + with torch.no_grad(): + output = model.backbone( + input_embeds=input_embeds, + hidden_states=hidden_states, + cache_hidden=None, + attention_mask=None, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + ) + + self.assertEqual(output.shape, (batch_size, seq_len, self.config.hidden_size)) + + def test_full_forward_pass(self): + """End-to-end: embed -> project -> backbone -> logits.""" + model = EdsdDraftModel(self.config) + model.eval() + batch_size, seq_len = 2, 10 + input_ids = torch.randint(0, self.config.vocab_size, (batch_size, seq_len)) + hidden_states = torch.randn(batch_size, seq_len, self.config.hidden_size * 2) + position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + + with torch.no_grad(): + input_embeds = model.embed_input_ids(input_ids) + projected = model.project_hidden_states(hidden_states) + output = model.backbone( + input_embeds=input_embeds, + hidden_states=projected, + cache_hidden=None, + attention_mask=None, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + ) + logits = model.compute_logits(output) + + self.assertEqual( + logits.shape, (batch_size, seq_len, self.config.draft_vocab_size) + ) + + # ------------------------------------------------------------------ + # Save / load + # ------------------------------------------------------------------ + + def test_save_pretrained(self): + model = EdsdDraftModel(self.config) + self.config.save_pretrained(self.temp_dir) + model_path = os.path.join(self.temp_dir, "pytorch_model.bin") + torch.save(model.state_dict(), model_path) + self.assertTrue(os.path.exists(os.path.join(self.temp_dir, "config.json"))) + self.assertTrue(os.path.exists(model_path)) + + def test_state_dict_compatibility(self): + model1 = EdsdDraftModel(self.config) + model2 = EdsdDraftModel(self.config) + state_dict = model1.state_dict() + model2.load_state_dict(state_dict) + + for name, param1 in model1.named_parameters(): + param2 = dict(model2.named_parameters())[name] + self.assertTrue(torch.equal(param1, param2)) + + @patch("transformers.modeling_utils.PreTrainedModel.from_pretrained") + def test_from_pretrained_mock(self, mock_from_pretrained): + mock_model = EdsdDraftModel(self.config) + mock_from_pretrained.return_value = mock_model + + loaded_model = EdsdDraftModel.from_pretrained(self.temp_dir) + mock_from_pretrained.assert_called_once_with(self.temp_dir) + self.assertIsInstance(loaded_model, EdsdDraftModel) + + # ------------------------------------------------------------------ + # Config validation + # ------------------------------------------------------------------ + + def test_config_validation(self): + invalid_config = LlamaConfig( + vocab_size=1000, + hidden_size=127, + num_attention_heads=4, + num_key_value_heads=2, + ) + with self.assertRaises(AssertionError): + EdsdDraftModel(invalid_config) + + +if __name__ == "__main__": + suite = unittest.TestSuite() + suite.addTest(unittest.makeSuite(TestEdsdDraftModel)) + runner = unittest.TextTestRunner(verbosity=2) + runner.run(suite) \ No newline at end of file