reference/code/SpecForge/scripts/train_peagle.py

"""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()