Update EAGLE offline flags to use EagleConfig

Signed-off-by: Benjamin Chislett <[email protected]>

Author: benchislett

examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py

@@ -166,12 +166,6 @@ async def main(args: argparse.Namespace) -> None:
 
         try:
             # Send the message to the OpenAI-compatible endpoint
-            # await client.chat.completions.create(
-            #     model=args.model,
-            #     messages=conversations,
-            #     temperature=0.0,
-            #     max_tokens=1,
-            # )
             await client.completions.create(
                 model=args.model,
                 prompt=input_string,

examples/speculative_decoding/eagle_utils.py

@@ -225,14 +225,17 @@ def __getitem__(self, i) -> dict[str, torch.Tensor]:
             raise ValueError(msg)
 
         ret = {**preprocessed_base}  # Shallow copy so we don't accidentally modify the cache
-        ret["hidden_states"] = offline_data["hidden_states"]
-        ret["aux_hidden_states"] = offline_data["aux_hidden_states"]
-
+        ret["kwargs"] = {
+            "base_model_outputs": {
+                "base_model_hidden_states": offline_data["hidden_states"],
+                "aux_hidden_states": offline_data["aux_hidden_states"],
+            }
+        }
         return ret
 
 
 def make_eagle_supervised_data_module(
-    tokenizer: transformers.PreTrainedTokenizer, data_args
+    tokenizer: transformers.PreTrainedTokenizer, data_args, use_offline_training: bool
 ) -> dict:
     """Make dataset and collator for supervised fine-tuning.
 
@@ -250,11 +253,14 @@ def make_eagle_supervised_data_module(
         else:
             data_json = json.load(f)
 
-    if data_args.offline_training:
+    if use_offline_training:
         print_rank_0("Loading pre-processed data for offline training...")
         dataset_cls = OfflineSupervisedDataset
 
         # Glob for all .pt files in the data_path directory
+        assert data_args.offline_data_path is not None, (
+            "offline_data_path must be provided for offline training."
+        )
         offline_data_path = Path(data_args.offline_data_path)
         all_files = {str(p) for p in offline_data_path.glob("*.pt")}
         if not all_files:
@@ -346,24 +352,30 @@ def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]:
 class DataCollatorForOffline(DataCollatorWithPadding):
     def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]:
         base_batch = super().__call__(features)
-        if "hidden_states" not in features[0]:
-            print(features[0].keys())
-            print(features[0])
-            print(features)
-            raise ValueError("Features do not contain 'hidden_states' key.")
-        max_hs_length = max(item["hidden_states"].shape[0] for item in features)
+        if "kwargs" not in features[0]:
+            raise ValueError("No kwargs found in batch features. Offline data required.")
+
+        features = [item["kwargs"]["base_model_outputs"] for item in features]
+        max_hs_length = max(item["base_model_hidden_states"].shape[0] for item in features)
 
         batch_hidden_states = torch.stack(
-            [self.paddingtensor2d(item["hidden_states"], max_hs_length) for item in features]
+            [
+                self.paddingtensor2d(item["base_model_hidden_states"], max_hs_length)
+                for item in features
+            ]
         )
         batch_aux_hidden_states = torch.stack(
             [self.paddingtensor2d(item["aux_hidden_states"], max_hs_length) for item in features]
         )
 
         batch = {
             **base_batch,
-            "hidden_states": batch_hidden_states,
-            "aux_hidden_states": batch_aux_hidden_states,
+            "kwargs": {
+                "base_model_outputs": {
+                    "base_model_hidden_states": batch_hidden_states,
+                    "aux_hidden_states": batch_aux_hidden_states,
+                }
+            },
         }
 
         return batch

examples/speculative_decoding/gen_synthetic_conversations/run_vllm_server.sh

@@ -14,25 +14,4 @@
 # limitations under the License.
 
 # Example launch configuration for a vLLM server
-# On 8xB200, Llama 3.3 70B runs comfortably with TP=2 at high batch sizes.
-
-# Achieve data parallelism by running multiple vLLM servers on different GPUs.
-# CUDA_VISIBLE_DEVICES=0,1 vllm serve meta-llama/Llama-3.3-70B-Instruct --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8000 &
-# CUDA_VISIBLE_DEVICES=2,3 vllm serve meta-llama/Llama-3.3-70B-Instruct --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8001 &
-# CUDA_VISIBLE_DEVICES=4,5 vllm serve meta-llama/Llama-3.3-70B-Instruct --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8002 &
-# CUDA_VISIBLE_DEVICES=6,7 vllm serve meta-llama/Llama-3.3-70B-Instruct --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8003 &
-
-# Alternatively, use vLLM's built-in data parallelism.
-# vllm serve meta-llama/Llama-3.3-70B-Instruct --tensor-parallel-size 2 --data-parallel-size 4 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8000
-
-# Default to a small model for testing.
-# vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --data-parallel-size 8 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8000
-
-CUDA_VISIBLE_DEVICES=0 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --max-num-seqs 1024 --port 8000
-# CUDA_VISIBLE_DEVICES=1 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8001 &
-# CUDA_VISIBLE_DEVICES=2 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8002 &
-# CUDA_VISIBLE_DEVICES=3 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8003 &
-# CUDA_VISIBLE_DEVICES=4 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8004 &
-# CUDA_VISIBLE_DEVICES=5 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8005 &
-# CUDA_VISIBLE_DEVICES=6 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8006 &
-# CUDA_VISIBLE_DEVICES=7 vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8007 &
+vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 1 --data-parallel-size 8 --max-num-batched-tokens 32768 --max-seq-len 8192 --disable-log-requests --max-num-seqs 1024 --port 8000

examples/speculative_decoding/launch.sh

@@ -30,6 +30,10 @@ while [ $# -gt 0 ]; do
       if [[ "$1" != *=* ]]; then shift; fi
       DATA="${1#*=}"
       ;;
+    --offline-data*)
+      if [[ "$1" != *=* ]]; then shift; fi
+      OFFLINE_DATA_PATH="${1#*=}"
+      ;;
     --mode*)
       if [[ "$1" != *=* ]]; then shift; fi
       MODE="${1#*=}"
@@ -87,7 +91,7 @@ set -x
 # Get the default value for save_steps based on the available number of GPUs
 GPU_COUNT=$(python -c "import torch; print(torch.cuda.device_count())")
 # Calculate save_steps
-DEFAULT_SAVE_STEPS=$((8192 / GPU_COUNT))
+DEFAULT_SAVE_STEPS=$((192 / GPU_COUNT))
 
 MODEL=${MODEL:-"TinyLlama/TinyLlama-1.1B-Chat-v1.0"}
 MODE=${MODE:-"eagle3"}
@@ -104,7 +108,8 @@ REDRAFTER_TOKENS=${REDRAFTER_TOKENS:-1}
 REDRAFTER_NUM_LAYERS=${REDRAFTER_NUM_LAYERS:-1}
 FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP=${FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP:-"LlamaDecoderLayer"}
 NUM_GPU=${NUM_GPU:-1}
-TRAINING_SEQ_LEN=${TRAINING_SEQ_LEN:-512}
+TRAINING_SEQ_LEN=${TRAINING_SEQ_LEN:-2048}
+OFFLINE_DATA_PATH=${OFFLINE_DATA_PATH:-""}
 
 if [[ "$MODE" == "medusa" ]]; then
   SPECULATIVE_ARGS="--medusa_num_heads $MEDUSA_NUM_HEADS --medusa_num_layers $MEDUSA_NUM_LAYERS"
@@ -119,6 +124,17 @@ else
   exit 1
 fi
 
+if [[ "$OFFLINE_DATA_PATH" != "" ]]; then
+  if [[ ! -d "$OFFLINE_DATA_PATH" ]]; then
+    echo "Offline data path $OFFLINE_DATA_PATH does not exist or is not a directory."
+    exit 1
+  else
+    OFFLINE_TRAINING_ARGS="--offline-data-path $OFFLINE_DATA_PATH"
+  fi
+else
+  OFFLINE_TRAINING_ARGS=""
+fi
+
 if [[ "$NUM_GPU" == 1 ]]; then
   MULTI_GPU=""
 else
@@ -149,6 +165,7 @@ CMD="accelerate launch $MULTI_GPU --mixed_precision bf16 main.py \
     --logging_steps 100 \
     --tf32 True \
     --data_path $DATA \
+    $OFFLINE_TRAINING_ARGS \
     $SPECULATIVE_ARGS
 "