merge online and offline into one script (#308)

* merge online and offline into one script

* polish

---------

Co-authored-by: root <FrankLeeeee>

Author: FrankLeeeee

docs/advanced_features/customization.md

@@ -6,7 +6,7 @@
 torchrun \
     --standalone \
     --nproc_per_node 8 \
-    ./scripts/train_eagle3_online.py \
+    ./scripts/train_eagle3.py \
     --target-model-path meta-llama/Llama-3.1-8B-Instruct \
     --draft-model-config ./configs/llama3-8B-eagle3.json \
     --train-data-path ./cache/dataset/sharegpt.jsonl \
@@ -19,7 +19,7 @@ torchrun \
     --cache-dir ./cache
 ```
 
-If you wish to understand what each argument does, you can run `python scripts/train_eagle3_online.py --help` to see the full list of arguments. Particularly, we will discuss some important arguments below.
+If you wish to understand what each argument does, you can run `python scripts/train_eagle3.py --help` to see the full list of arguments. Particularly, we will discuss some important arguments below.
 - `--chat-template`: This should be the chat template to use for the model, so please make sure you set it to the correct value.
 - `--cache-dir`: This directory contains the dataset cache including the `input_ids`, `loss_mask`, `attention_mask` and `vocab_mapping`. These caches can make your data loading much faster once a cache is generated. The cache file has a name which is obtained by hashing the dataset path to avoid cache collision.
 

docs/basic_usage/data_preparation.md

@@ -100,7 +100,7 @@ To use pre-formatted datasets, add the `--is-preformatted` flag to your training
 
 ```bash
 torchrun --standalone --nproc_per_node 8 \
-    scripts/train_eagle3_online.py \
+    scripts/train_eagle3.py \
     --is-preformatted \
     --chat-template qwen \
     --train-data-path ./your_preformatted_dataset.jsonl \

examples/README.md

@@ -10,7 +10,7 @@ TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels \
 torchrun \
     --standalone \
     --nproc_per_node $NUM_GPUS \
-    $ROOT_DIR/scripts/train_eagle3_online.py \
+    $ROOT_DIR/scripts/train_eagle3.py \
     --target-model-path meta-llama/Llama-3.1-8B-Instruct \
     --draft-model-config $ROOT_DIR/configs/llama3-8B-eagle3.json \
     --train-data-path $ROOT_DIR/cache/dataset/sharegpt.jsonl \

examples/prepare_hidden_states.sh

@@ -7,15 +7,16 @@ NUM_GPUS=${1:-8}
 torchrun \
     --standalone \
     --nproc_per_node $NUM_GPUS \
-    $ROOT_DIR/scripts/train_eagle3_online.py \
-    --target-model-path DeepSeek-V2-Lite \
+    $ROOT_DIR/scripts/train_eagle3.py \
+    --target-model-path deepseek-ai/DeepSeek-V2-Lite \
     --draft-model-config $ROOT_DIR/configs/deepseek-v2-lite-eagle3.json \
-    --train-data-path $ROOT_DIR/cache/dataset/sharegpt.jsonl \
-    --output-dir $ROOT_DIR/outputs/deepseek-v2-lite-eagle3 \
+    --train-data-path $ROOT_DIR/cache/dataset/sharegpt_train.jsonl \
+    --output-dir $ROOT_DIR/outputs/deepseek-v2-lite-eagle3-sharegpt \
     --num-epochs 10 \
     --batch-size 1 \
     --tp-size 1 \
     --learning-rate 1e-4 \
-    --max-length 2048 \
+    --max-length 4096 \
     --chat-template deepseek \
-    --cache-dir $ROOT_DIR/cache \
+    --target-model-backend hf \
+    --cache-dir $ROOT_DIR/cache

examples/run_deepseek_v2_lite_eagle3_online.sh

@@ -7,7 +7,7 @@ NUM_GPUS=${1:-8}
 torchrun \
     --standalone \
     --nproc_per_node $NUM_GPUS \
-    $ROOT_DIR/scripts/train_eagle3_online.py \
+    $ROOT_DIR/scripts/train_eagle3.py \
     --target-model-path openai/gpt-oss-120b \
     --draft-model-config $ROOT_DIR/configs/gpt-oss-20B-eagle3.json \
     --train-data-path $ROOT_DIR/cache/dataset/perfect-blend-gptoss-20B.jsonl \
@@ -16,7 +16,7 @@ torchrun \
     --num-epochs 10 \
     --batch-size 1 \
     --learning-rate 1e-4 \
-    --max-length 2048 \
+    --max-length 4096 \
     --chat-template gpt-oss \
     --cache-dir $ROOT_DIR/cache \
     --dist-timeout 60

examples/run_gpt_oss_120b_eagle3_online.sh

@@ -7,18 +7,15 @@ NUM_GPUS=${1:-8}
 torchrun \
     --standalone \
     --nproc_per_node $NUM_GPUS \
-    $ROOT_DIR/scripts/train_eagle3_online.py \
+    $ROOT_DIR/scripts/train_eagle3.py \
     --target-model-path openai/gpt-oss-20b \
     --draft-model-config $ROOT_DIR/configs/gpt-oss-20B-eagle3.json \
     --train-data-path $ROOT_DIR/cache/dataset/perfect-blend-gptoss-20B.jsonl \
     --output-dir $ROOT_DIR/outputs/perfect-blend-gptoss-20b-eagle3 \
     --num-epochs 10 \
     --batch-size 1 \
     --learning-rate 1e-4 \
-    --max-length 2048 \
+    --max-length 4096 \
     --chat-template gpt-oss \
     --cache-dir $ROOT_DIR/cache \
     --dist-timeout 60
-
-
-# --train-data-path $ROOT_DIR/cache/dataset/perfect-blend-gptoss-20B.jsonl \

examples/run_gpt_oss_120b_eagle3_sgl_online.sh

@@ -7,7 +7,7 @@ NUM_GPUS=${1:-8}
 torchrun \
     --standalone \
     --nproc_per_node $NUM_GPUS \
-    $ROOT_DIR/scripts/train_eagle3_online.py \
+    $ROOT_DIR/scripts/train_eagle3.py \
     --target-model-path meta-llama/Llama-3.3-70B-Instruct \
     --draft-model-config $ROOT_DIR/configs/llama3-8B-eagle3.json \
     --train-data-path $ROOT_DIR/cache/dataset/sharegpt.jsonl \
@@ -16,6 +16,6 @@ torchrun \
     --batch-size 1 \
     --tp-size 8 \
     --learning-rate 1e-4 \
-    --max-length 2048 \
+    --max-length 4096 \
     --chat-template llama3 \
     --cache-dir $ROOT_DIR/cache

examples/run_gpt_oss_20b_eagle3_online.sh

@@ -1,75 +1,35 @@
-#!/bin/bash
-export PERSIST_DIR=/tmp # Please Change this to your own directory
-export MODEL_PATH="meta-llama/Llama-3.1-8B-Instruct"
-export DATASET_PATH=$PERSIST_DIR/dataset/
-export CACHE_DIR=$PERSIST_DIR/cache/
-export OUTPUT_DIR=$PERSIST_DIR/outputs/
-export HIDDEN_STATES_DIR=$PERSIST_DIR/hidden_states/
-export MAX_LENGTH=2048
-export CHAT_TEMPLATE=llama3
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+ROOT_DIR=$(dirname $SCRIPT_DIR)
+NUM_GPUS=${1:-8}
 
-hf download $MODEL_PATH
-hf download Aeala/ShareGPT_Vicuna_unfiltered --repo-type dataset
-
-python scripts/prepare_data.py --dataset sharegpt --output-path $DATASET_PATH --split-eval
-python scripts/build_eagle3_dataset.py \
-    --model-path $MODEL_PATH \
-    --data-path $DATASET_PATH \
-    --cache-dir $CACHE_DIR \
-    --chat-template $CHAT_TEMPLATE \
-    --max-length $MAX_LENGTH \
-
-CUDA_VISIBLE_DEVICES=1,2,3,4 torchrun --nproc_per_node=4 \
-    scripts/prepare_hidden_states.py \
-    --data-path $DATASET_PATH/sharegpt_test.jsonl \
-    --model-path $MODEL_PATH \
-    --cache-dir $CACHE_DIR \
-    --output-path $HIDDEN_STATES_DIR/sharegpt_test \
-    --chat-template $CHAT_TEMPLATE \
-    --max-length $MAX_LENGTH \
-    --enable-aux-hidden-states \
-    --tp-size 4 \
-    --batch-size 4 \
-    --mem-frac=0.75
-
-CUDA_VISIBLE_DEVICES=1,2,3,4 torchrun --nproc_per_node=4 \
+# generate hidden states
+torchrun \
+    --standalone \
+    --nproc_per_node $NUM_GPUS \
     scripts/prepare_hidden_states.py \
-    --data-path $DATASET_PATH/sharegpt_train.jsonl \
-    --model-path $MODEL_PATH \
-    --cache-dir $CACHE_DIR \
-    --output-path $HIDDEN_STATES_DIR/sharegpt_train \
-    --chat-template $CHAT_TEMPLATE \
-    --max-length $MAX_LENGTH \
+    --target-model-path meta-llama/Llama-3.1-8B-Instruct \
     --enable-aux-hidden-states \
-    --tp-size 4 \
-    --batch-size 4 \
-    --mem-frac=0.75
-
-# python scripts/view_data.py --data-path $HIDDEN_STATES_DIR/all_test/rows_0-5000/data_100.ckpt --tokenizer $MODEL_PATH
-# python scripts/view_data.py --data-path $HIDDEN_STATES_DIR/all_train/rows_0-5000/data_100.ckpt --tokenizer $MODEL_PATH
+    --data-path $ROOT_DIR/cache/dataset/sharegpt_train.jsonl \
+    --output-path $ROOT_DIR/cache/hidden_states/sharegpt_train_Llama-3.1-8B-Instruct \
+    --chat-template llama3 \
+    --max-length 4096 \
+    --tp-size 1 \
+    --batch-size 32
 
-export NUM_GPUS=4
-CUDA_VISIBLE_DEVICES=1,2,3,4 torchrun \
+# train eagle3 offline
+torchrun \
     --standalone \
     --nproc_per_node $NUM_GPUS \
-    scripts/train_eagle3_offline.py \
-    --target-model-path $MODEL_PATH \
-    --draft-model-config ./configs/llama3-8B-eagle3.json \
-    --train-data-path $DATASET_PATH/sharegpt_train.jsonl \
-    --train-hidden-states-path $HIDDEN_STATES_DIR/sharegpt_train/ \
-    --eval-data-path $DATASET_PATH/sharegpt_test.jsonl \
-    --eval-hidden-states-path $HIDDEN_STATES_DIR/sharegpt_test/ \
-    --output-dir $OUTPUT_DIR \
+    $ROOT_DIR/scripts/train_eagle3.py \
+    --target-model-path meta-llama/Llama-3.1-8B-Instruct \
+    --draft-model-config $ROOT_DIR/configs/llama3-8B-eagle3.json \
+    --train-data-path $ROOT_DIR/cache/dataset/sharegpt_train.jsonl \
+    --train-hidden-states-path $ROOT_DIR/cache/hidden_states/sharegpt_train_Llama-3.1-8B-Instruct \
+    --output-dir $ROOT_DIR/outputs/llama3-8b-eagle3-sharegpt-offline \
     --num-epochs 10 \
-    --draft-global-batch-size 16 \
-    --draft-micro-batch-size 1 \
-    --learning-rate 5e-5 \
-    --draft-attention-backend flex_attention \
-    --max-length $MAX_LENGTH \
-    --chat-template $CHAT_TEMPLATE \
-    --cache-dir $CACHE_DIR \
-    --dist-timeout=10 \
-    --log-steps 1 \
-    --report-to wandb \
-    --wandb-project llama3-8b-eagle3 \
-    --wandb-name offline-100k-4gpus
+    --batch-size 1 \
+    --tp-size 1 \
+    --learning-rate 1e-4 \
+    --max-length 4096 \
+    --chat-template llama3 \
+    --cache-dir $ROOT_DIR/cache