examples/speculative_decoding/main.py (1)

188-190: Good: honor EagleConfig.eagle_offline per prior feedback.
This aligns with the earlier request to use the existing config flag.

examples/speculative_decoding/prepare_input_conversations/utils.py (2)

26-36: Add a network timeout and avoid multi-context in one line for clarity.

Prevents indefinite hangs and improves readability; keep parent dir creation.

Apply:

 async def download_file(url: str, destination: Path) -> None:
     """Download a file from a URL to a specified destination."""
     destination.parent.mkdir(parents=True, exist_ok=True)
-    async with aiohttp.ClientSession() as session, session.get(url) as response:
-        if response.status != 200:
-            msg = f"Failed to download {url}: {response.status}"
-            raise RuntimeError(msg)
-        content = await response.read()
-        destination.write_bytes(content)
-        print(f"Downloaded {url} to {destination}")
+    timeout = aiohttp.ClientTimeout(total=600)
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+        async with session.get(url) as response:
+            if response.status != 200:
+                msg = f"Failed to download {url}: {response.status}"
+                raise RuntimeError(msg)
+            content = await response.read()
+            destination.write_bytes(content)
+            print(f"Downloaded {url} to {destination}")

---

45-86: Prevent duplicate IDs within the same update and validate required keys.

Currently, duplicates inside the provided conversations list can be appended multiple times because existing_ids isn’t updated as you add. Also, missing "conversations" will raise KeyError later; validate early.

Apply:

 def add_conversations_to_split(conversations: list, dataset_dir: Path, split: str) -> None:
@@
-    # Open the dataset file for the specified split, or create it if it doesn't exist
-    dataset_file = dataset_dir / f"{split}.jsonl"
+    # Ensure output directory exists and open/create split file
+    dataset_dir.mkdir(parents=True, exist_ok=True)
+    dataset_file = dataset_dir / f"{split}.jsonl"
@@
-    existing_ids = {entry["conversation_id"] for entry in all_conversations}
+    existing_ids = {entry["conversation_id"] for entry in all_conversations}
     num_new_entries = 0
     num_duplicates = 0
     for entry in conversations:
-        if entry.get("conversation_id") is None:
+        entry_id = entry.get("conversation_id")
+        if entry_id is None:
             raise ValueError("Each conversation must have a 'conversation_id' field.")
-        if entry["conversation_id"] not in existing_ids:
+        if "conversations" not in entry:
+            raise ValueError("Each conversation must have a 'conversations' field.")
+        if entry_id not in existing_ids:
             all_conversations.append(
                 {
-                    "conversation_id": entry["conversation_id"],
+                    "conversation_id": entry_id,
                     "conversations": entry["conversations"],
                 }
             )
             num_new_entries += 1
+            existing_ids.add(entry_id)
         else:
             num_duplicates += 1
@@
-    dataset_dir.mkdir(parents=True, exist_ok=True)
     with dataset_file.open("w", encoding="utf-8") as f:
         for entry in all_conversations:
             f.write(json.dumps(entry, ensure_ascii=False) + "\n")

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py (2)

27-29: Unify REMOVE_THINK_CHAT_TEMPLATE with training and guard None chat_template.

Avoid divergence from training preprocessing and prevent AttributeError when chat_template is None.

-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer
+try:
+    # Keep in sync with training; import if available when run from examples/speculative_decoding
+    from eagle_utils import REMOVE_THINK_CHAT_TEMPLATE
+except Exception:
+    # Fallback for standalone runs; ensure value matches eagle_utils.REMOVE_THINK_CHAT_TEMPLATE
+    REMOVE_THINK_CHAT_TEMPLATE = (
+        "{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}"
+    )
@@
-REMOVE_THINK_CHAT_TEMPLATE = (
-    "{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}"
-)
@@
-    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "")
+    if getattr(tokenizer, "chat_template", None):
+        tokenizer.chat_template = tokenizer.chat_template.replace(
+            REMOVE_THINK_CHAT_TEMPLATE, ""
+        )

Also applies to: 96-100

---

138-148: Clamp/deduplicate auxiliary layer indices; avoid out‑of‑range/duplicates on small models.

Current logic can pick invalid indices (e.g., 2 when num_hidden_layers < 3).

-            # Extract hidden states from layers with index (2, N/2, N-3), and the output hidden states
+            # Extract hidden states from early/mid/late layers; clamp within [0, N-1]
             hidden_states = outputs.hidden_states
-            selected_layer_indices = [
-                2,
-                max(0, num_hidden_layers // 2),
-                max(1, num_hidden_layers - 3),
-            ]
-            selected_layer_indices = sorted(set(selected_layer_indices))
+            candidates = [2, num_hidden_layers // 2, num_hidden_layers - 3]
+            selected_layer_indices = sorted(
+                {i for i in candidates if 0 <= i <= num_hidden_layers - 1}
+            )
+            if not selected_layer_indices:
+                selected_layer_indices = [0]
             aux_hidden_states = torch.cat(
                 [hidden_states[i].squeeze(0).cpu() for i in selected_layer_indices], dim=-1
             )

examples/speculative_decoding/main.py (3)

72-81: Annotate as Optional to match None default.

offline_data_path defaults to None but is typed as str. Make it str | None for consistency with eval_data_path and type checkers.

Apply this diff:

-    offline_data_path: str = field(
+    offline_data_path: str | None = field(

---

144-145: Make the offline switch robust to empty strings.

is not None treats "" as offline. Prefer truthiness.

Apply this diff:

-use_offline_training = data_args.offline_data_path is not None
+use_offline_training = bool(data_args.offline_data_path)

---

150-159: num_hidden_layers=0: add safe fallback and always set num_orig_hidden_layers.

Keep your space‑saving default but harden for models that can’t instantiate with 0 layers, and ensure num_orig_hidden_layers is set even when resuming or the override path changes.

Apply this diff:

-        model_kwargs = {"num_hidden_layers": 0} if use_offline_training else {}
-        model = transformers.AutoModelForCausalLM.from_pretrained(
-            model_args.model_name_or_path, torch_dtype="auto", **model_kwargs
-        )
-        if use_offline_training:
-            # When doing offline training, we need to set num_hidden_layers
-            # since we override it when loading the model for space savings
-            model_config = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path)
-            model.config.num_orig_hidden_layers = model_config.num_hidden_layers
+        model_kwargs = {"num_hidden_layers": 0} if use_offline_training else {}
+        try:
+            model = transformers.AutoModelForCausalLM.from_pretrained(
+                model_args.model_name_or_path, torch_dtype="auto", **model_kwargs
+            )
+        except Exception as e:
+            if use_offline_training:
+                print_rank_0("num_hidden_layers=0 failed; falling back to 1 for offline training.")
+                model = transformers.AutoModelForCausalLM.from_pretrained(
+                    model_args.model_name_or_path, torch_dtype="auto", num_hidden_layers=1
+                )
+            else:
+                raise
+        if use_offline_training and not hasattr(model.config, "num_orig_hidden_layers"):
+            # Record original depth for plugins that need it.
+            base_cfg = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path)
+            model.config.num_orig_hidden_layers = base_cfg.num_hidden_layers

Please run a quick smoke test on at least one model per family you target (e.g., Llama, Mistral, Qwen2, Phi‑3) to confirm no fallback is triggered unexpectedly and training proceeds.

examples/speculative_decoding/prepare_input_conversations/example_make_prompt_dataset.sh (2)

1-1: Add a shebang to fix ShellCheck SC2148 and enable script execution.

Without a shebang, shells may interpret this file incorrectly and ShellCheck flags SC2148.

Apply:

+#!/usr/bin/env bash
 # SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

---

16-25: Fix relative paths; make script location-agnostic; add strict mode and ensure output dir.

The script lives in prepare_input_conversations/, but calls into prepare_input_conversations/*.py again, which breaks when run from repo root or this directory. Also add set -euo pipefail and create a data dir.

Apply:

-# Example script to prepare a dataset of prompts for generation
-# Lines in this script can be uncommented to include specific datasets/splits in the prompt dataset.
+# Example script to prepare a dataset of prompts for generation
+# Lines in this script can be uncommented to include specific datasets/splits in the prompt dataset.
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# Optional: change DATASET_DIR or pass via env
+DATASET_DIR="${DATASET_DIR:-${SCRIPT_DIR}/data}"
+mkdir -p "${DATASET_DIR}"

-python3 prepare_input_conversations/add_daring_anteater.py --output-split-name train
-# python3 prepare_input_conversations/add_sharegpt.py --output-split-name train
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split train_sft --output-split-name train
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split train_gen --output-split-name train
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split test_sft --output-split-name mix_test
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split test_gen --output-split-name mix_test
-python3 prepare_input_conversations/add_mtbench.py --output-split-name mix_test
+python3 "${SCRIPT_DIR}/add_daring_anteater.py" --output-split-name train ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"}
+# python3 "${SCRIPT_DIR}/add_sharegpt.py" --output-split-name train ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"}
+# python3 "${SCRIPT_DIR}/add_ultrachat.py" --ultrachat-split train_sft --output-split-name train ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"}
+# python3 "${SCRIPT_DIR}/add_ultrachat.py" --ultrachat-split train_gen --output-split-name train ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"}
+# python3 "${SCRIPT_DIR}/add_ultrachat.py" --ultrachat-split test_sft --output-split-name mix_test ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"}
+# python3 "${SCRIPT_DIR}/add_ultrachat.py" --ultrachat-split test_gen --output-split-name mix_test ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"}
+python3 "${SCRIPT_DIR}/add_mtbench.py" --output-split-name mix_test ${DATASET_DIR:+--dataset-dir "${DATASET_DIR}"} 

Note: If the add_* CLIs don’t support --dataset-dir, the parameter will be omitted; the mkdir remains safe.

examples/speculative_decoding/prepare_input_conversations/utils.py (1)

93-125: Make ratio check tolerant to FP error and avoid mutating global RNG state.

isclose avoids false negatives; using a local Random(seed) preserves global RNG for callers.

Apply:

+import math
@@
-    if train_ratio + val_ratio + test_ratio != 1.0:
+    if not math.isclose(train_ratio + val_ratio + test_ratio, 1.0, rel_tol=0.0, abs_tol=1e-9):
         msg = "Ratios must sum to 1.0"
         raise ValueError(msg)
@@
-    if shuffle:
-        random.seed(seed)
-        random.shuffle(conversations)
+    if shuffle:
+        random.Random(seed).shuffle(conversations)

examples/speculative_decoding/collect_hidden_states/run_send_conversations.sh (1)

1-24: Add shebang/safe flags, fix debug flag name, and quote paths.

Without a shebang, running as an executable fails; ShellCheck SC2148 also flags this. The commented flag name doesn’t exist in the CLI (should be --debug-max-num-conversations). Also make invocation path-robust and quote args.

+#!/usr/bin/env bash
+set -euo pipefail
+
+# Resolve script directory to make path robust when invoked from anywhere.
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
-python3 collect_hidden_states/send_conversations_for_hiddens.py \
-  --model meta-llama/Llama-3.2-1B-Instruct \
-  --input-file synthetic_conversations/mtbench.jsonl \
-  --output-dir /mnt/md0/eagle-hidden-states/llama1b/mtbench/
-# --debug-max-num-conversations-per-split 1000
+python3 "$script_dir/send_conversations_for_hiddens.py" \
+  --model "meta-llama/Llama-3.2-1B-Instruct" \
+  --input-file "synthetic_conversations/mtbench.jsonl" \
+  --output-dir "/mnt/md0/eagle-hidden-states/llama1b/mtbench/"
+# Optional: limit processed conversations during debugging
+# --debug-max-num-conversations 1000

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py (3)

93-95: Set eval() before inference.

Minor but standard to disable dropout and ensure deterministic behavior.

-    model = AutoModel.from_pretrained(args.model, torch_dtype="auto", device_map="auto")
+    model = AutoModel.from_pretrained(args.model, torch_dtype="auto", device_map="auto")
+    model.eval()

---

103-127: Counter name is misleading; it tracks too‑short and too‑long.

Tiny naming nit; adjust for clarity.

-    num_skipped_too_long = 0
+    num_filtered_by_length = 0
@@
-        if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len:
-            num_skipped_too_long += 1
+        if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len:
+            num_filtered_by_length += 1
             continue
@@
-    if num_skipped_too_long > 0:
-        print(f"Skipped {num_skipped_too_long} conversations due to length constraints.")
+    if num_filtered_by_length > 0:
+        print(f"Skipped {num_filtered_by_length} conversations due to length constraints.")

Also applies to: 163-164

---

82-83: Async isn’t used; simplify to synchronous main.

Removes unnecessary asyncio plumbing.

-async def main(args: argparse.Namespace) -> None:
+def main(args: argparse.Namespace) -> None:
@@
-if __name__ == "__main__":
-    cli_args = parse_args()
-    asyncio.run(main(cli_args))
+if __name__ == "__main__":
+    cli_args = parse_args()
+    main(cli_args)

Also applies to: 176-178

examples/speculative_decoding/collect_hidden_states/run_hf_compute_hiddens.sh (1)

1-23: Add shebang/safe flags, quote args, and create output dir.

Prevents exec failures (SC2148) and path issues; ensures output dir exists.

+#!/usr/bin/env bash
+set -euo pipefail
+
+# Resolve script directory
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Output directory (edit as needed)
+OUTPUT_DIR="/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/"
+mkdir -p "$OUTPUT_DIR"
+
-python3 collect_hidden_states/compute_hidden_states_hf.py \
-  --model meta-llama/Llama-3.2-1B-Instruct \
-  --input-file synthetic_conversations/daring-anteater.jsonl \
-  --output-dir /mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+python3 "$script_dir/compute_hidden_states_hf.py" \
+  --model "meta-llama/Llama-3.2-1B-Instruct" \
+  --input-file "synthetic_conversations/daring-anteater.jsonl" \
+  --output-dir "$OUTPUT_DIR"
+# Optional:
+# --max-seq-len 3072

examples/speculative_decoding/collect_hidden_states/run_hf_compute_hiddens_dp.sh (1)

16-36: Harden DP runner: shebang/safe flags, temp workspace, quoting, and mkdir.

Avoids /tmp collisions, improves safety, and makes paths robust.

+#!/usr/bin/env bash
+set -euo pipefail
+
-INPUT_FILE=synthetic_conversations/daring-anteater.jsonl
-OUTPUT_DIR=/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+INPUT_FILE="synthetic_conversations/daring-anteater.jsonl"
+OUTPUT_DIR="/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/"
+mkdir -p "$OUTPUT_DIR"
+
-split -n l/8 --numeric-suffixes=0 -d --additional-suffix=.jsonl $INPUT_FILE /tmp/part-
+tmpdir="$(mktemp -d)"
+trap 'rm -rf "$tmpdir"' EXIT
+split -n l/8 --numeric-suffixes=0 -d --additional-suffix=.jsonl "$INPUT_FILE" "$tmpdir/part-"
 
 for i in $(seq 0 7)
 do
-CUDA_VISIBLE_DEVICES=$i python3 collect_hidden_states/compute_hidden_states_hf.py --model meta-llama/Llama-3.2-1B-Instruct --input-file /tmp/part-0${i}.jsonl --output-dir $OUTPUT_DIR &
+  CUDA_VISIBLE_DEVICES="$i" python3 collect_hidden_states/compute_hidden_states_hf.py \
+    --model "meta-llama/Llama-3.2-1B-Instruct" \
+    --input-file "$tmpdir/part-0${i}.jsonl" \
+    --output-dir "$OUTPUT_DIR" &
 done
 wait
-
-rm /tmp/part-*.jsonl
+# Temporary files cleaned via trap

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: benchislett
PR: NVIDIA/TensorRT-Model-Optimizer#300
File: examples/speculative_decoding/main.py:150-159
Timestamp: 2025-09-17T14:14:44.950Z
Learning: In examples/speculative_decoding/main.py, setting num_hidden_layers=0 for offline training has been successfully tested by benchislett with past models and works in their use case for EAGLE3 offline training.

Learnt from: benchislett
PR: NVIDIA/TensorRT-Model-Optimizer#300
File: examples/speculative_decoding/main.py:150-159
Timestamp: 2025-09-17T14:14:44.950Z
Learning: In examples/speculative_decoding/main.py, setting num_hidden_layers=0 for offline training has been successfully tested by benchislett with past models and works in their use case for EAGLE3 offline training.

examples/speculative_decoding/launch_train.sh (1)

127-136: Good: validates offline data path and builds args.

Consider also quoting the path to tolerate spaces.

Apply:

-    OFFLINE_TRAINING_ARGS="--offline-data-path $OFFLINE_DATA_PATH"
+    OFFLINE_TRAINING_ARGS="--offline-data-path \"$OFFLINE_DATA_PATH\""

modelopt/torch/speculative/plugins/transformers.py (1)

338-349: Clamp aux layer IDs when offline and num_layers <= 0.

If num_layers resolves to 0/None offline, current defaults may produce invalid indices (e.g., 1). Clamp to valid range and drop negatives.

Apply:

-        num_layers = self.config.num_hidden_layers
-        if self.eagle_offline and (num_layers is None or num_layers <= 0):
-            num_layers = getattr(self.config, "num_orig_hidden_layers", 0)
-
-        self.eagle_config.eagle_aux_hidden_state_layer_ids = [
-            1,
-            max(0, num_layers // 2 - 1),
-            max(0, num_layers - 4),
-        ]
-        self.eagle_config.eagle_aux_hidden_state_layer_ids = list(
-            set(self.eagle_config.eagle_aux_hidden_state_layer_ids)
-        )
+        num_layers = self.config.num_hidden_layers
+        if self.eagle_offline and (num_layers is None or num_layers <= 0):
+            num_layers = getattr(self.config, "num_orig_hidden_layers", 0)
+        # Propose defaults then clamp into [0, max(0, num_layers-1)] and dedupe/sort
+        candidate_ids = [1, max(0, num_layers // 2 - 1), max(0, num_layers - 4)]
+        hi = max(0, num_layers - 1)
+        clamped = sorted({min(max(i, 0), hi) for i in candidate_ids})
+        self.eagle_config.eagle_aux_hidden_state_layer_ids = clamped

examples/speculative_decoding/collect_hidden_states/run_send_conversations.sh (1)

16-23: Add shebang and safe shell flags; optional: parametrize endpoint via env.

Apply:

+#!/usr/bin/env bash
+set -euo pipefail
+
 python3 collect_hidden_states/send_conversations_for_hiddens.py \
   --model meta-llama/Llama-3.2-1B-Instruct \
   --input-file synthetic_conversations/mtbench.jsonl \
   --output-dir /mnt/md0/eagle-hidden-states/llama1b/mtbench/
 # --debug-max-num-conversations-per-split 1000

examples/speculative_decoding/prepare_input_conversations/example_make_prompt_dataset.sh (1)

16-26: Add shebang/safe flags and make calls cwd-agnostic.

Use script-relative paths so it works from any directory.

Apply:

+#!/usr/bin/env bash
+set -euo pipefail
+
+# Example script to prepare a dataset of prompts for generation
 # Lines in this script can be uncommented to include specific datasets/splits in the prompt dataset.
 
-python3 prepare_input_conversations/add_daring_anteater.py --output-split-name train
-# python3 prepare_input_conversations/add_sharegpt.py --output-split-name train
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split train_sft --output-split-name train
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split train_gen --output-split-name train
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split test_sft --output-split-name mix_test
-# python3 prepare_input_conversations/add_ultrachat.py --ultrachat-split test_gen --output-split-name mix_test
-python3 prepare_input_conversations/add_mtbench.py --output-split-name mix_test
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+python3 "$ROOT/prepare_input_conversations/add_daring_anteater.py" --output-split-name train
+# python3 "$ROOT/prepare_input_conversations/add_sharegpt.py" --output-split-name train
+# python3 "$ROOT/prepare_input_conversations/add_ultrachat.py" --ultrachat-split train_sft --output-split-name train
+# python3 "$ROOT/prepare_input_conversations/add_ultrachat.py" --ultrachat-split train_gen --output-split-name train
+# python3 "$ROOT/prepare_input_conversations/add_ultrachat.py" --ultrachat-split test_sft --output-split-name mix_test
+# python3 "$ROOT/prepare_input_conversations/add_ultrachat.py" --ultrachat-split test_gen --output-split-name mix_test
+python3 "$ROOT/prepare_input_conversations/add_mtbench.py" --output-split-name mix_test

examples/speculative_decoding/collect_hidden_states/run_hf_compute_hiddens_dp.sh (1)

16-36: Add shebang/safe flags; consider quoting vars and checking split results.

Apply:

+#!/usr/bin/env bash
+set -euo pipefail
+
 # Example usage of the script to compute the hidden states for a conversation dataset
 # This script computes hidden states using a Hugging Face model and saves them to
 # the specified output directory. It does so in a data-parallel manner across 8 GPUs, by splitting
 # the input file into 8 parts and running 8 processes in parallel, one on each GPU.
@@
-INPUT_FILE=synthetic_conversations/daring-anteater.jsonl
-OUTPUT_DIR=/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+INPUT_FILE="synthetic_conversations/daring-anteater.jsonl"
+OUTPUT_DIR="/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/"
 
-split -n l/8 --numeric-suffixes=0 -d --additional-suffix=.jsonl $INPUT_FILE /tmp/part-
+split -n l/8 --numeric-suffixes=0 -d --additional-suffix=.jsonl "$INPUT_FILE" /tmp/part-
 
 for i in $(seq 0 7)
 do
-CUDA_VISIBLE_DEVICES=$i python3 collect_hidden_states/compute_hidden_states_hf.py --model meta-llama/Llama-3.2-1B-Instruct --input-file /tmp/part-0${i}.jsonl --output-dir $OUTPUT_DIR &
+  CUDA_VISIBLE_DEVICES=$i python3 collect_hidden_states/compute_hidden_states_hf.py \
+    --model meta-llama/Llama-3.2-1B-Instruct \
+    --input-file "/tmp/part-0${i}.jsonl" \
+    --output-dir "$OUTPUT_DIR" &
 done
 wait
 
-rm /tmp/part-*.jsonl
+rm -f /tmp/part-*.jsonl

examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py (2)

63-63: Load tensors onto CPU to avoid CUDA dependency.

Apply:

-        data = torch.load(file)
+        data = torch.load(file, map_location="cpu")

---

64-73: Don’t require exact key set; accept supersets.

Future producers may add fields; only assert required keys present.

Apply:

-        expected_keys = [
+        expected_keys = [
             "input_ids",
             "hidden_states",
             "aux_hidden_states",
             "conversation_id",
         ]
-        if set(expected_keys) != set(data.keys()):
-            print(f"File {file} does not contain all expected keys: {expected_keys}")
-            print(f"  Found keys: {list(data.keys())}")
-            continue
+        missing = [k for k in expected_keys if k not in data]
+        if missing:
+            print(f"File {file} missing required keys: {missing}. Found: {list(data.keys())}")
+            continue

examples/speculative_decoding/prepare_input_conversations/add_ultrachat.py (2)

23-27: Make imports robust to script working directory/package execution.

from utils import ... will fail if invoked outside this folder or as a module. Prefer relative import with a fallback for script execution.

-from utils import (
+try:
+    from .utils import (
+        dataset_splits_explanation,
+        id_for_conversation,
+        update_dataset_file_with_conversations,
+    )
+except ImportError:
+    from utils import (
         dataset_splits_explanation,
         id_for_conversation,
         update_dataset_file_with_conversations,
-)
+    )

---

61-65: Prefer iterating rows directly to avoid repeated random access.

Hugging Face datasets support direct iteration; it’s clearer and avoids repeated indexing overhead.

-    for i in tqdm(
-        range(len(ds)),
+    for i in tqdm(
+        range(len(ds)),  # or: enumerate(ds)
         desc=f"Loading UltraChat split {args.ultrachat_split}",
         total=len(ds),
     ):

If acceptable, switch to for i, row in enumerate(tqdm(ds, ...)) and use row[...].

examples/speculative_decoding/prepare_input_conversations/add_daring_anteater.py (2)

23-27: Harden import path (same concern as other prep scripts).

Apply the same relative-import fallback pattern suggested for add_ultrachat.py to avoid ImportError when run outside the folder.

---

63-76: Defensive role parsing.

.lower() on a missing/None role will raise. The diff above normalizes via str(...).lower() and falls back cleanly.

examples/speculative_decoding/eagle_utils.py (2)

176-236: Offline dataset: load safely on CPU and validate keys/shapes.

• Use map_location="cpu" to avoid accidental CUDA deserialization errors.
• Validate presence of expected keys and provide a concise error.

-        offline_data = torch.load(offline_file_path)
+        offline_data = torch.load(offline_file_path, map_location="cpu")
+        for k in ("input_ids", "hidden_states", "aux_hidden_states"):
+            if k not in offline_data:
+                raise ValueError(f"Missing key '{k}' in offline data: {offline_file_path}")
         offline_data["input_ids"] = offline_data["input_ids"][:max_length]
         offline_data["hidden_states"] = offline_data["hidden_states"][:max_length, :]
         offline_data["aux_hidden_states"] = offline_data["aux_hidden_states"][:max_length, :]
 
         # Make sure the input_ids have the same shape
         if preprocessed_base["input_ids"].shape != offline_data["input_ids"].shape:
-            msg = f"""Input IDs from offline data do not match the preprocessed input IDs
-                                for offline data sample at {offline_file_path}."""
-            raise ValueError(msg)
+            raise ValueError(
+                f"Shape mismatch for input_ids at {offline_file_path}: "
+                f"offline={tuple(offline_data['input_ids'].shape)} "
+                f"preproc={tuple(preprocessed_base['input_ids'].shape)}"
+            )

---

251-256: Handle blank lines in JSONL and set explicit encoding.

Minor robustness: ignore empty lines and open with utf‑8.

-    with open(data_args.data_path) as f:
-        if data_args.data_path.endswith("jsonl"):
-            data_json = [json.loads(line) for line in f]
+    with open(data_args.data_path, "r", encoding="utf-8") as f:
+        if data_args.data_path.endswith("jsonl"):
+            data_json = [json.loads(line) for line in f if line.strip()]
         else:
             data_json = json.load(f)

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py (1)

95-95: Null‑safe chat template handling.

Some tokenizers have no chat_template. Guard before replace.

-    tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "")
+    if getattr(tokenizer, "chat_template", None):
+        tokenizer.chat_template = tokenizer.chat_template.replace(
+            REMOVE_THINK_CHAT_TEMPLATE, ""
+        )

examples/speculative_decoding/prepare_input_conversations/add_sharegpt.py (3)

23-28: Import robustness (same pattern as other prep scripts).

Apply the relative‑import fallback used in prior comments to avoid ImportError when running from different CWDs.

---

62-85: Be lenient on unknown roles; prefer skip over raise to avoid aborting long runs.

Unexpected roles in ShareGPT can abort the entire job. Treat them as “skip this conversation” rather than raising.

-        else:
-            err_msg = f"Unknown role in conversation: {turn.get('from')}"
-            raise ValueError(err_msg)
+        else:
+            # Unknown role; skip this conversation
+            return None

---

87-101: Download step: ensure parent directory exists.

~/.cache may not exist. download_file creates parents, but direct existence check may fail earlier if parent is missing. Create dirs before checking.

-        args.sharegpt_file = Path("~/.cache/sharegpt.json").expanduser().resolve()
+        args.sharegpt_file = Path("~/.cache/sharegpt.json").expanduser()
+        args.sharegpt_file.parent.mkdir(parents=True, exist_ok=True)
+        args.sharegpt_file = args.sharegpt_file.resolve()

examples/speculative_decoding/prepare_input_conversations/utils.py (2)

103-110: Avoid strict float equality for ratio sum.

Use a small tolerance to prevent false negatives due to FP rounding.

-    if train_ratio + val_ratio + test_ratio != 1.0:
+    total = train_ratio + val_ratio + test_ratio
+    if abs(total - 1.0) > 1e-9:
         msg = "Ratios must sum to 1.0"
         raise ValueError(msg)

---

115-124: Don’t mutate caller’s list or global RNG state when shuffling.

Shuffle a copy with a local RNG for determinism without side effects.

-    if shuffle:
-        random.seed(seed)
-        random.shuffle(conversations)
+    if shuffle:
+        rng = random.Random(seed)
+        pool = conversations.copy()
+        rng.shuffle(pool)
+    else:
+        pool = conversations
@@
-    train_conversations = conversations[:train_count]
-    val_conversations = conversations[train_count : train_count + val_count]
-    test_conversations = conversations[train_count + val_count :]
+    train_conversations = pool[:train_count]
+    val_conversations = pool[train_count : train_count + val_count]
+    test_conversations = pool[train_count + val_count :]

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: benchislett
PR: NVIDIA/TensorRT-Model-Optimizer#300
File: examples/speculative_decoding/main.py:150-159
Timestamp: 2025-09-17T14:14:44.961Z
Learning: In examples/speculative_decoding/main.py, setting num_hidden_layers=0 for offline training has been successfully tested by benchislett with past models and works in their use case for EAGLE3 offline training.

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py (3)

89-91: Ensure eval mode for deterministic inference.

Set model.eval() after loading.

-    model = AutoModel.from_pretrained(args.model, torch_dtype="auto", device_map="auto")
+    model = AutoModel.from_pretrained(args.model, torch_dtype="auto", device_map="auto")
+    model.eval()

---

18-20: Drop asyncio (not used) or add awaits.

main is async but contains no awaits; simplify to a regular function.

-import asyncio
@@
-async def main(args: argparse.Namespace) -> None:
+def main(args: argparse.Namespace) -> None:
@@
 if __name__ == "__main__":
     cli_args = parse_args()
-    asyncio.run(main(cli_args))
+    main(cli_args)

Also applies to: 82-82, 174-176

---

112-113: Sanitize conversation_id for filenames.

Untrusted conversation_id could contain path separators or illegal chars.

+import re
@@
-        conversation_id = entry.get("conversation_id", "{:08d}".format(idx))
+        raw_id = entry.get("conversation_id", "{:08d}".format(idx))
+        conversation_id = re.sub(r"[^A-Za-z0-9._-]", "_", str(raw_id))

Also applies to: 148-159, 18-26

modelopt/torch/speculative/plugins/transformers.py (1)

337-343: Safer fallback when num_orig_hidden_layers missing

If offline and num_hidden_layers<=0 but num_orig_hidden_layers is absent, you set num_layers=0, yet later include layer id 1. Use a non‑zero fallback.

-        if self.eagle_offline and (num_layers is None or num_layers <= 0):
-            num_layers = getattr(self.config, "num_orig_hidden_layers", 0)
+        if self.eagle_offline and (num_layers is None or num_layers <= 0):
+            num_layers = getattr(self.config, "num_orig_hidden_layers", self.config.num_hidden_layers)

examples/speculative_decoding/prepare_input_conversations/utils.py (3)

26-36: Add a network timeout to avoid hanging downloads

aiohttp.ClientSession() without timeout can hang indefinitely.

-async def download_file(url: str, destination: Path) -> None:
+async def download_file(url: str, destination: Path) -> None:
     """Download a file from a URL to a specified destination."""
     destination.parent.mkdir(parents=True, exist_ok=True)
-    async with aiohttp.ClientSession() as session, session.get(url) as response:
+    timeout = aiohttp.ClientTimeout(total=600)
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+        async with session.get(url) as response:
             if response.status != 200:
                 msg = f"Failed to download {url}: {response.status}"
                 raise RuntimeError(msg)
             content = await response.read()
             destination.write_bytes(content)
             print(f"Downloaded {url} to {destination}")

---

104-109: Float-sum equality is brittle

Use tolerance when validating ratios.

-    if train_ratio + val_ratio + test_ratio != 1.0:
+    if abs((train_ratio + val_ratio + test_ratio) - 1.0) > 1e-6:
         msg = "Ratios must sum to 1.0"
         raise ValueError(msg)

---

50-61: Validate 'conversations' presence on existing data

Guard against malformed prior entries to avoid KeyError later.

-    if any(not entry.get("conversation_id") for entry in all_conversations):
+    if any(not entry.get("conversation_id") for entry in all_conversations):
         msg = "All existing conversations must have a 'conversation_id' field."
         raise ValueError(msg)
+    if any("conversations" not in entry for entry in all_conversations):
+        msg = "All existing conversations must have a 'conversations' field."
+        raise ValueError(msg)

examples/speculative_decoding/prepare_input_conversations/example_make_prompt_dataset.sh (1)

1-26: Add shebang and strict mode; ensure output dir exists

Improves portability and early failure.

+#!/usr/bin/env bash
+set -euo pipefail
+
+# Ensure default output dir exists (scripts default to input_conversations/)
+mkdir -p input_conversations

examples/speculative_decoding/collect_hidden_states/run_send_conversations.sh (3)

1-1: Add a shebang (or ShellCheck directive) to satisfy SC2148 and make the script directly executable.

Without a shebang, shells and tooling can misinterpret the script.

Apply one of the following diffs (preferred: add a shebang):

+#!/usr/bin/env bash
 # SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0

Or, if you intentionally keep it sourced/run via bash file.sh, add a ShellCheck directive:

+# shellcheck shell=bash
 # SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0

---

16-23: Harden and parameterize the example so it runs from any cwd and fails fast on errors.

Resolve brittle relative paths, create output dir, and enforce required env for the OpenAI‑compatible endpoint.

Apply this diff:

-# Example usage of the script to send conversations for hidden state collection
-# This script sends conversations to a (local) OpenAI-compatible server for processing and collects hidden states.
-
-python3 collect_hidden_states/send_conversations_for_hiddens.py \
-  --model meta-llama/Llama-3.2-1B-Instruct \
-  --input-file synthetic_conversations/mtbench.jsonl \
-  --output-dir /mnt/md0/eagle-hidden-states/llama1b/mtbench/
-# --debug-max-num-conversations-per-split 1000
+# Example usage of the script to send conversations for hidden state collection.
+# This script sends conversations to a (local) OpenAI-compatible server for processing and collects hidden states.
+
+set -euo pipefail
+IFS=$'\n\t'
+
+# Resolve repo root relative to this script so it works from any cwd
+SCRIPT_DIR="$(cd -- "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+
+# Required env for OpenAI-compatible server (fail early if missing)
+: "${OPENAI_API_BASE:?Set OPENAI_API_BASE to your server URL, e.g. http://127.0.0.1:8000/v1}"
+: "${OPENAI_API_KEY:?Set OPENAI_API_KEY (dummy value is fine for many local servers)}"
+
+# Tunables (override via env: MODEL, INPUT_FILE, OUTPUT_DIR)
+MODEL="${MODEL:-meta-llama/Llama-3.2-1B-Instruct}"
+INPUT_FILE="${INPUT_FILE:-$REPO_ROOT/examples/speculative_decoding/synthetic_conversations/mtbench.jsonl}"
+OUTPUT_DIR="${OUTPUT_DIR:-/mnt/md0/eagle-hidden-states/llama1b/mtbench}"
+
+# Sanity checks
+[ -f "$INPUT_FILE" ] || { echo "Input not found: $INPUT_FILE" >&2; exit 1; }
+mkdir -p "$OUTPUT_DIR"
+
+python3 "$REPO_ROOT/collect_hidden_states/send_conversations_for_hiddens.py" \
+  --model "$MODEL" \
+  --input-file "$INPUT_FILE" \
+  --output-dir "$OUTPUT_DIR"
+# Optional:
+#   --debug-max-num-conversations-per-split 1000

---

19-23: Make the script executable in git.

If this is intended to be run directly, mark it executable.

#!/bin/bash
# From repo root
git update-index --chmod=+x examples/speculative_decoding/collect_hidden_states/run_send_conversations.sh

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: benchislett
PR: NVIDIA/TensorRT-Model-Optimizer#300
File: examples/speculative_decoding/main.py:150-159
Timestamp: 2025-09-17T14:14:44.961Z
Learning: In examples/speculative_decoding/main.py, setting num_hidden_layers=0 for offline training has been successfully tested by benchislett with past models and works in their use case for EAGLE3 offline training.

examples/speculative_decoding/collect_hidden_states/run_hf_compute_hiddens_dp.sh (2)

25-27: Quote vars and ensure output directory exists.

Prevent word-splitting; create the output directory.

-INPUT_FILE=synthetic_conversations/daring-anteater.jsonl
-OUTPUT_DIR=/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+INPUT_FILE="synthetic_conversations/daring-anteater.jsonl"
+OUTPUT_DIR="/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/"
+mkdir -p "$OUTPUT_DIR"

---

36-36: Rely on trap for cleanup; avoid risky rm patterns.

The trap above safely removes the entire temp directory.

-rm /tmp/part-*.jsonl
+# cleanup handled by trap

examples/speculative_decoding/collect_hidden_states/run_hf_compute_hiddens.sh (1)

16-19: Parameterize paths, quote, and ensure output dir exists.

Improves portability and robustness; keeps the example concise.

-# Example usage of the script to compute the hidden states for a conversation dataset
+# Example usage of the script to compute the hidden states for a conversation dataset
@@
-python3 collect_hidden_states/compute_hidden_states_hf.py \
-  --model meta-llama/Llama-3.2-1B-Instruct \
-  --input-file synthetic_conversations/daring-anteater.jsonl \
-  --output-dir /mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+MODEL="${MODEL:-meta-llama/Llama-3.2-1B-Instruct}"
+INPUT_FILE="${INPUT_FILE:-synthetic_conversations/daring-anteater.jsonl}"
+OUTPUT_DIR="${OUTPUT_DIR:-/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/}"
+mkdir -p "$OUTPUT_DIR"
+
+# Resolve Python entry relative to this script for robustness
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
+python3 "$SCRIPT_DIR/compute_hidden_states_hf.py" \
+  --model "$MODEL" \
+  --input-file "$INPUT_FILE" \
+  --output-dir "$OUTPUT_DIR"

examples/speculative_decoding/launch_train.sh (4)

33-36: CLI flag name is slightly misleading vs downstream arg; consider aliasing.

Downstream you pass --offline-data-path to main.py. Here the user-facing flag is --offline-data. Consider supporting both for clarity:

• Keep --offline-data (backward-compat)
• Add an alias --offline-data-path (sets the same var)

Apply this minimal alias:

-    --offline-data*)
+    --offline-data*|--offline-data-path*)
       if [[ "$1" != *=* ]]; then shift; fi
       OFFLINE_DATA_PATH="${1#*=}"
       ;;

---

81-84: Invalid-argument error prints the value, not the flag.

For inputs like --foo=bar, message shows “bar” instead of “--foo”. Print $1 directly.

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument %s\n" "$1"

---

146-171: Quoting/arg safety when paths contain spaces.

You already quote OFFLINE_DATA_PATH and DATA in the proposed patches. For consistency, consider quoting other path-like vars (MODEL, OUTPUT_DIR) or build CMD as an array to avoid word-splitting. Non-blocking.

---

107-109: Unused vars (REDRAFTER_*).

REDRAFTER_TOKENS and REDRAFTER_NUM_LAYERS aren’t used in this script. Remove or wire them, or add a comment if they’re placeholders for future flags.

examples/speculative_decoding/eagle_utils.py (2)

35-35: Guard chat_template mutation

Some tokenizers may lack chat_template or it may be None; also consider not mutating per-call.

-    tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "")
+    if getattr(tokenizer, "chat_template", None):
+        tokenizer.chat_template = tokenizer.chat_template.replace(
+            REMOVE_THINK_CHAT_TEMPLATE, ""
+        )

Optional follow-up: move this one-time adjustment to dataset init to avoid repeated mutation.

---

294-297: Shuffle before split to avoid ordering bias

Current 95/5 split preserves source order; shuffle with a seed.

-        num_train = int(len(valid_entries) * 0.95)
+        # Optional: shuffle to avoid order bias
+        import random
+        rnd = random.Random(getattr(data_args, "seed", 42))
+        rnd.shuffle(valid_entries)
+        num_train = int(len(valid_entries) * 0.95)

examples/speculative_decoding/collect_hidden_states/run_send_conversations.sh (2)

1-1: Add a shebang (and strict mode) so the script runs as a shell script and satisfies ShellCheck.

Without a shebang (SC2148), shells may guess incorrectly; strict mode avoids silent failures.

+#!/usr/bin/env bash
+set -euo pipefail

---

19-19: Make invocation robust regardless of the current working directory.

Resolve the Python script relative to this file’s location to avoid path errors.

-python3 collect_hidden_states/send_conversations_for_hiddens.py \
+python3 "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)/send_conversations_for_hiddens.py" \

examples/speculative_decoding/collect_hidden_states/send_conversations_for_hiddens.py (1)

189-191: Remove redundant continue.

Control flow already proceeds to the next iteration.

-        num_success += 1
-        continue
+        num_success += 1

examples/speculative_decoding/prepare_input_conversations/add_sharegpt.py (2)

87-96: Resolve output directory path for consistency.

Normalize --output-dir (supports ~, avoids relative path surprises).

Apply this diff:

 async def main(args: argparse.Namespace) -> None:
@@
     if not args.sharegpt_file.exists():
         err_msg = f"ShareGPT file {args.sharegpt_file} does not exist."
         raise FileNotFoundError(err_msg)
 
+    # Normalize output dir
+    args.output_dir = args.output_dir.expanduser().resolve()

---

62-83: PII/safety filter hook (optional).

ShareGPT "unfiltered" may contain PII/toxicity. Consider an optional sanitizer callback applied to msgs before hashing/emit.

I can add a lightweight regex-based scrubber and a --drop-pii flag if you want.