feat(speculative): add vLLM data synthesis pipeline and Nemotron dataset preparation scripts (#1176)

### What does this PR do?

Type of change: New feature, new example, bug fix

Adds a vLLM-based synthetic data generation pipeline for speculative
decoding draft model training, along with dataset preparation scripts
for NVIDIA's Nemotron Post-Training dataset collections.

**Data synthesis pipeline** (`tools/launcher/common/vllm/query.sh` +
`common/query.py`):
- Launch a vLLM server and run multi-turn inference to synthesize
training data from input conversation skeletons
- Fork-safe OpenAI client: reinitializes HTTP connection pool after
`datasets.map()` forks worker processes, preventing 400 errors from
corrupted connections
- Clear Docker `ENTRYPOINT` so vLLM containers (which default to `vllm
serve`) work correctly under NeMo Run's executor
- `--max-tokens` argument to bound generation length
- Local file loading support (`--data /path/to/file.jsonl`)
- Re-raise connection errors so `datasets.map()` halts the shard instead
of silently producing empty rows
- Map `developer` role to `system` (OpenAI format compatibility)

**Multi-turn reasoning trace handling** (`common/query.py`):
- Strip `<think>...</think>` blocks from intermediate assistant turns
before re-feeding to the model; preserve the full trace only on the
final turn

**Nemotron dataset preparation** (`examples/dataset/`):
- `make_nemotron_ptv2_dataset.py` — prepares
[nvidia/Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2)
(~3.3M rows generate, ~1.9M rows train)
- `make_nemotron_ptv3_dataset.py` — prepares the [Nemotron PTv3
collection](https://huggingface.co/collections/nvidia/nemotron-post-training-v3)
of 16 datasets (~3.4M rows generate, ~3.9M rows train)
- Both support `generate` mode (strips assistant turns for synthesis
input) and `train` mode (normalizes to clean OpenAI format for SFT)
- `conversation_utils.py` — shared utilities: `strip_assistant_turns`,
`normalize_messages`, `make_augment_fn`, `AugmentationSpec`
- `augmentations.yaml` — 12 language-redirect variants + style/format
hints, cycled across dataset rows
- Scripts live in `examples/dataset/` (not under
`speculative_decoding/`) to signal reusability beyond speculative
decoding

**Bug fixes**:
- `strip_assistant_turns()`: return `{"messages": []}` when no user
turns remain (system-only rows were previously passed through instead of
being filtered)
- `concatenate_datasets()`: guard against empty parts list
- SSH tunnel user precedence: explicit `user` arg now correctly
overrides `slurm_config.user`

### Usage

```bash
# Prepare PTv3 input conversations for synthesis (~3.4M rows):
python examples/dataset/make_nemotron_ptv3_dataset.py --output-dir /tmp/ptv3_gen

# Launch vLLM server + synthesize responses:
bash tools/launcher/common/vllm/query.sh \
    --model /path/to/model \
    --tensor-parallel-size 4 \
    -- \
    --data /tmp/ptv3_gen/default.jsonl \
    --save /tmp/ptv3_responses \
    --num-shards 10 --num-proc 4 --max-tokens 4096

# Prepare PTv2 for direct SFT training (~1.9M rows):
python examples/dataset/make_nemotron_ptv2_dataset.py --mode train --output-dir /tmp/ptv2_train
```

### Testing

Tested end-to-end on an NVIDIA GB10 node (119 GiB GPU memory) with
`vllm/vllm-openai:qwen3_5-cu130` container and `Qwen/Qwen3.5-4B`:
- vLLM server starts correctly with cleared Docker entrypoint
- `datasets.map(num_proc=4)` runs without connection errors (fork-safe
client)
- Multi-turn synthesis produces correct assistant responses with
thinking traces handled

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: N/A (data synthesis scripts;
tested manually)
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added dataset generation and augmentation capabilities for Nemotron
post-training datasets (v2 and v3)
* Enhanced query functionality with thinking-block filtering and
improved client management for robust parallel processing
* Added support for local dataset file paths alongside HuggingFace Hub
datasets

* **Bug Fixes**
* Fixed SLURM executor user resolution and Docker container entrypoint
configuration
* Improved error handling for connection failures during dataset
synthesis

* **Documentation**
* Updated dataset preparation guide with new generation modes and
augmentation configuration details

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>

Author: ChenhanYu

examples/dataset/README.md

@@ -0,0 +1,128 @@
+# Dataset Preparation Scripts
+
+Utilities for building conversation datasets from NVIDIA Nemotron Post-Training
+collections and other HuggingFace sources.  These scripts produce datasets in
+**standard OpenAI chat format** (`{"messages": [{"role": ..., "content": ...}]}`)
+and can be used for any downstream fine-tuning task — SFT, distillation,
+speculative decoding draft-model training, etc.
+
+## Files
+
+| File | Description |
+|---|---|
+| `make_nemotron_ptv3_dataset.py` | Build a dataset from the [Nemotron PT v3 collection](https://huggingface.co/collections/nvidia/nemotron-post-training-v3) using a configurable YAML mix |
+| `make_nemotron_ptv2_dataset.py` | Build a dataset from [Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) |
+| `make_dataset.py` | General-purpose mixer for arbitrary HuggingFace datasets (mtbench, sharegpt, ultrachat, magpie, etc.) |
+| `conversation_utils.py` | Shared utilities: augmentation, role normalization, assistant-turn stripping |
+| `add_nemotron_chat.py` | Add Nemotron v2 chat conversations to an existing dataset |
+| `augmentations.yaml` | Augmentation variants (language redirects, style hints) for `make_nemotron_pt*.py` |
+| `nemotron_ptv3_datasets.yaml` | Dataset mix config for `make_nemotron_ptv3_dataset.py` |
+| `example_data_config.yaml` | Example YAML config for `make_dataset.py` |
+
+## Quick Start
+
+### Install dependencies
+
+```bash
+pip install datasets huggingface_hub pyyaml
+huggingface-cli login   # required for gated datasets
+```text
+
+### Build a Nemotron PT v3 dataset
+
+```bash
+# Synthetic data generation inputs (strips last assistant turn so a model can regenerate it)
+python make_nemotron_ptv3_dataset.py --output-dir /tmp/ptv3_gen
+
+# Full conversations for direct SFT training
+python make_nemotron_ptv3_dataset.py --mode train --output-dir /tmp/ptv3_train
+
+# Use a custom dataset mix
+python make_nemotron_ptv3_dataset.py --config my_mix.yaml --output-dir /tmp/ptv3_custom
+```text
+
+### Build a Nemotron PT v2 dataset
+
+```bash
+python make_nemotron_ptv2_dataset.py --output-dir /tmp/ptv2_gen
+python make_nemotron_ptv2_dataset.py --mode train --output-dir /tmp/ptv2_train
+```text
+
+### Build a general-purpose mixed dataset
+
+```bash
+python make_dataset.py --config example_data_config.yaml --output-dir /tmp/mixed
+```text
+
+## Dataset Modes
+
+Both `make_nemotron_pt*.py` scripts support two modes:
+
+| Mode | Description | Use case |
+|---|---|---|
+| `generate` (default) | Strips assistant turns, optionally augments prompts | Input data for synthetic generation (query a target model to produce training responses) |
+| `train` | Keeps all turns, normalizes to clean OpenAI format | Direct SFT / distillation training |
+
+## Synthetic Generation Pipeline
+
+The `generate` mode produces conversation skeletons that are fed to a target model
+via `tools/launcher/common/query.py` (vLLM or TRT-LLM).  The output becomes training
+data for a draft model (e.g. EAGLE3 speculative decoding) or a distilled student:
+
+```text
+make_nemotron_ptv3_dataset.py --mode generate  →  skeleton.jsonl
+        ↓
+query.py  (target model generates responses turn-by-turn)
+        ↓
+training data for draft model / student
+```text
+
+## Augmentations
+
+`augmentations.yaml` defines language-redirect and style-hint variants that are
+applied cyclically across the dataset.  Each enabled entry produces one augmented
+copy of the source rows.
+
+To customize augmentations:
+- **Disable** a variant: add `enabled: false`
+- **Add** a language redirect: append a `user_suffix` entry
+- **Add** a system prompt: append a `system_prompt` entry
+
+```yaml
+augmentations:
+  - type: user_suffix
+    text: " Please reply in French instead of English."
+  - type: system_prompt
+    content: "You are a helpful assistant."
+    enabled: false   # disable without deleting
+```text
+
+## Dataset Mix Config (`nemotron_ptv3_datasets.yaml`)
+
+Edit this file to add, remove, or re-weight datasets without touching the script:
+
+```yaml
+datasets:
+  - repo_id: nvidia/Nemotron-Math-v2
+    splits: [high_part00, high_part01]
+    cap_per_split: 200000
+    augment: true
+
+  - repo_id: nvidia/OpenMathReasoning-mini
+    splits: [train]
+    augment: false   # multilingual — skip language-redirect augmentation
+```text
+
+## Output Format
+
+Every output row is a JSONL object with a single `messages` key:
+
+```json
+{"messages": [
+  {"role": "system",    "content": "You are a helpful assistant."},
+  {"role": "user",      "content": "What is 2+2?"},
+  {"role": "assistant", "content": "4"}
+]}
+```text
+
+In `generate` mode, assistant turns are stripped so the row ends with a user turn.

…input_conversations/add_nemotron_chat.py examples/dataset/add_nemotron_chat.pyexamples/speculative_decoding/prepare_input_conversations/add_nemotron_chat.py renamed to examples/dataset/add_nemotron_chat.py examples/speculative_decoding/prepare_input_conversations/add_nemotron_chat.py renamed to examples/dataset/add_nemotron_chat.py

@@ -0,0 +1,93 @@
+# Augmentation specs for make_nemotron_ptv2_dataset.py and make_nemotron_ptv3_dataset.py
+#
+# Each entry defines one augmentation variant applied cyclically across the dataset.
+# The augmented copy is the same size as the source — each row gets exactly one variant.
+#
+# Supported types:
+#
+#   user_suffix
+#     Appends `text` to the content of every user message in the conversation.
+#     Example use: language-redirect instructions, style/length hints.
+#
+#   system_prompt
+#     Prepends a {"role": "system", "content": <content>} message to the conversation.
+#     Use this for model-specific flags (e.g. /no_think) or persona instructions.
+#     Set `enabled: false` for variants that are not supported by your target model.
+#
+# To disable an entry without deleting it, add `enabled: false`.
+# To add a new variant, append a new entry following the same schema.
+
+augmentations:
+
+  # --- Language redirects (user_suffix) ------------------------------------
+
+  - type: user_suffix
+    text: " Please reply in French instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Italian instead of English."
+
+  - type: user_suffix
+    text: " Please reply in German instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Spanish instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Mandarin Chinese instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Japanese instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Korean instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Turkish instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Modern Standard Arabic instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Russian instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Brazilian Portuguese instead of English."
+
+  - type: user_suffix
+    text: " Please reply in Vietnamese instead of English."
+
+  # --- Style / format hints (user_suffix) ----------------------------------
+
+  - type: user_suffix
+    text: " Be concise and answer in as few words as possible."
+
+  - type: user_suffix
+    text: " Provide a detailed, step-by-step explanation."
+
+  - type: user_suffix
+    text: " Format your response using Markdown (headers, bullet points, code blocks where appropriate)."
+
+  - type: user_suffix
+    text: " Do not use Markdown formatting; reply in plain text only."
+
+  - type: user_suffix
+    text: " Explain your answer as if I am a complete beginner with no prior knowledge."
+
+  - type: user_suffix
+    text: " Assume I am an expert; skip basic explanations and go straight to the details."
+
+  - type: user_suffix
+    text: " Think step by step before giving your final answer."
+
+  # --- System-prompt variants (system_prompt) ------------------------------
+
+  # /no_think: suppresses chain-of-thought in models that support it (e.g. Qwen3).
+  # Set enabled: false if your target model does not support this flag.
+  - type: system_prompt
+    content: "You are a helpful assistant. /no_think"
+    enabled: false
+
+  # Generic helpful-assistant system prompt (no special flags).
+  - type: system_prompt
+    content: "You are a helpful, respectful, and honest assistant."