diff --git a/examples/speculative_decoding/recipes/README_dflash_cosmos3_nano.md b/examples/speculative_decoding/recipes/README_dflash_cosmos3_nano.md new file mode 100644 index 00000000000..3f829b23668 --- /dev/null +++ b/examples/speculative_decoding/recipes/README_dflash_cosmos3_nano.md @@ -0,0 +1,122 @@ + + +# Training a DFlash Draft Model for Cosmos3 Nano + +An end-to-end multimodal DFlash recipe: synthesize training data from image and +video prompts, train the draft model, export it, and smoke-test it under vLLM. + +The executable pipeline is a ModelOpt Launcher example: + +```bash +cd tools/launcher +uv run launch.py --yaml examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml --yes +``` + +This document covers the reasoning behind that pipeline. For per-step +configuration, read the YAML — every task is commented inline. + +## Why synthesize the training data + +DFlash trains a draft model to predict what the *target* model will say. It +benefits far more from the target model's own completions than from a large +collection of human-written answers, because the draft's job is to match the +target's distribution, not to be independently correct. + +So the pipeline does not consume an off-the-shelf SFT dataset. It takes prompts +from several sources, replays every prompt through the target model, and trains +on the target's responses. Each prompt is generated at several temperatures, +which widens coverage of the target's output distribution; the merge step then +removes near-duplicate completions. + +## The data sources + +Four complementary sources, three of which the pipeline builds itself: + +1. **PAI-Understanding** — representative video usage. +2. **VQA v2** — image visual reasoning. +3. **Multilingual prompts** — high-quality text prompts + (`nvidia/Speculative-Decoding-Multilingual-Prompt-v2`). +4. **Curated text** — optional, supplied by you. + +The PAI and VQA sources keep their media; the other two are text-only. Mixing +text into a multimodal draft matters: a draft trained only on media prompts +degrades on the plain-text turns that dominate real conversations. + +### Curated text (optional fourth source) + +To build the Nemotron Chat component, follow +`recipes/train_eagle_head_cosmos_reason2.ipynb`: accept the Hugging Face licence +for `nvidia/Nemotron-Post-Training-Dataset-v2`, then run + +```bash +python ../prepare_input_conversations/add_nemotron_chat.py --mapping-file nemotron_mapping.bin +``` + +which writes `input_conversations/nemotron-chat.jsonl` (89,511 conversations). +Point a `--source curated_text=` at that file, or at any privacy-reviewed +JSONL of assistant-completed conversations in the same `messages`/`conversations` +format. To use both, concatenate them into one valid JSONL first. + +## The merge step + +Merging is not just concatenation. It does two things the training job depends on: + +- **Resolves every image/video reference to an absolute path.** This is why + training runs with `data.vlm_img_dir=/`. +- **Deduplicates conservatively.** Records are grouped by identical prompt and + media, and only near-identical completions within a group are dropped. The + temperature sweep intentionally produces varied answers to the same prompt; + the goal is to remove redundancy, not diversity. + +## Training constraints worth knowing + +- `training_seq_len` must be divisible by `dflash_block_size`. The example uses + `16384` and `8`. +- The `VLM_*` environment variables cap text and visual token growth *before* + tokenization. They matter because visual token count depends on the actual + resolution and frame count of each sample: without caps, one high-resolution + video can expand past `training_seq_len`, which the collator rejects rather + than silently truncating (a silent truncation would corrupt the DFlash labels). +- Setting `data.vlm_processor` is what selects the multimodal collator. Leave it + unset and training uses the text-only path, even for a VLM target. + +## Deployment + +The pipeline's last task runs a vLLM smoke test against the exported draft. + +Serve the original Cosmos3 Nano target with the exported draft. A DFlash block +size of eight yields seven speculative tokens — the remaining position is the +context/bonus token. The command below is a single-GPU smoke test: it binds to +localhost and caps the context at 4096 tokens to keep startup and memory +bounded. Raise `--max-model-len` only after sizing the context and concurrency +you need. Add `--trust-remote-code` only for a checkpoint you trust. + +```bash +SERVED_MODEL_NAME=cosmos3-nano-dflash +vllm serve "$MODEL_PATH" \ + --host 127.0.0.1 --port 8000 \ + --max-model-len 4096 \ + --served-model-name "$SERVED_MODEL_NAME" \ + --speculative-config "{\"method\":\"dflash\",\"model\":\"$EXPORT_PATH\",\"num_speculative_tokens\":7}" + +# In another terminal: +curl -fsS http://127.0.0.1:8000/health +curl -fsS http://127.0.0.1:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d "{\"model\":\"$SERVED_MODEL_NAME\",\"messages\":[{\"role\":\"user\",\"content\":\"Reply with exactly: DFlash deployment smoke test passed.\"}],\"temperature\":0,\"max_tokens\":16}" +``` diff --git a/examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb b/examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb deleted file mode 100644 index 6c020bcdf01..00000000000 --- a/examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb +++ /dev/null @@ -1,732 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "7fb27b941602401d91542211134fc71a", - "metadata": {}, - "source": [ - "# Training a DFlash Draft Model for Cosmos3 Nano\n", - "\n", - "This production runbook trains an online DFlash draft model with Cosmos3 Nano as the frozen vision-language target. DFlash learns a compact block-diffusion draft model.\n", - "\n", - "| Step | Description |\n", - "| :---: | :--- |\n", - "| 1 | Configure the training-data build |\n", - "| 2 | Build and merge the training data |\n", - "| 3 | Submit the eight-GPU training job |\n", - "| 4 | Export a checkpoint for deployment |\n", - "\n", - "> **Hardware** – the example requests eight GPUs by default.\n" - ] - }, - { - "cell_type": "markdown", - "id": "acae54e37e7d407bbb7b55eff062a284", - "metadata": {}, - "source": [ - "## Step 1 – Configure the Training-Data Build\n", - "\n", - "Build the merged Cosmos3 Nano DFlash JSONL from the sources below, then use this cell to set the model, data, and output locations shared by the remaining steps. Treat the provided source mix as a starting point and extend it with privacy-reviewed, task-relevant user data when available.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9a63283cbaf04dbcab1f6479b197f3a8", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from pathlib import Path\n", - "\n", - "REPO_ROOT = Path.cwd().resolve().parents[2]\n", - "SPEC_ROOT = REPO_ROOT / \"examples\" / \"speculative_decoding\"\n", - "if not (REPO_ROOT / \"pyproject.toml\").is_file():\n", - " raise RuntimeError(\n", - " f\"Run this notebook from examples/speculative_decoding/recipes; got {Path.cwd()}\"\n", - " )\n", - "\n", - "TRAINING_DATA = Path(\n", - " os.environ.get(\"TRAINING_DATA\", str(SPEC_ROOT / \"CR3_data\" / \"cosmos3_nano_dflash_train.jsonl\"))\n", - ").resolve()\n", - "MODEL_PATH = os.environ.get(\n", - " \"MODEL_PATH\",\n", - " \"/lustre/fsw/portfolios/coreai/users/skierat/ngc/ngc-cli/cosmos3-nano-reasoner_vbf16-final\",\n", - ")\n", - "OUTPUT_DIR = Path(\n", - " os.environ.get(\"OUTPUT_DIR\", str(REPO_ROOT / \"ckpts\" / \"cosmos3-nano-dflash-first-run\"))\n", - ").resolve()\n", - "os.environ.update(\n", - " REPO_ROOT=str(REPO_ROOT),\n", - " TRAINING_DATA=str(TRAINING_DATA),\n", - " MODEL_PATH=MODEL_PATH,\n", - " OUTPUT_DIR=str(OUTPUT_DIR),\n", - ")\n", - "print(f\"Training data will be written to: {TRAINING_DATA}\")\n", - "if TRAINING_DATA.is_file() and TRAINING_DATA.stat().st_size:\n", - " print(f\"Existing file size: {TRAINING_DATA.stat().st_size / 2**30:.1f} GiB\")\n", - "print(f\"Training output: {OUTPUT_DIR}\")" - ] - }, - { - "cell_type": "markdown", - "id": "8edb47106e1a46a883d545849b8ab81b", - "metadata": {}, - "source": [ - "## Step 2 – Build the Training Data\n", - "\n", - "Prepare prompt shards, generate Cosmos3 Nano completions in a Slurm allocation, then merge the results into the `TRAINING_DATA` path configured above. The paths below default to this checkout rather than a prior run.\n", - "\n", - "These cells require a Jupyter environment with the example requirements installed. Run the configuration cell below before any data-preparation cell, and set `PLAIN_TEXT_INPUT` to curated Nemotron Chat or compatible user-data JSONL before merging.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "10185d26023b46108eb7d9f57d49d2b3", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "SPEC_ROOT = Path.cwd().resolve().parent\n", - "REPO_ROOT = SPEC_ROOT.parents[1]\n", - "DATA_ROOT = Path(os.environ.get(\"DATA_ROOT\", str(SPEC_ROOT / \"CR3_data\"))).resolve()\n", - "MODEL_PATH = os.environ.get(\n", - " \"MODEL_PATH\",\n", - " \"/lustre/fsw/portfolios/coreai/users/skierat/ngc/ngc-cli/cosmos3-nano-reasoner_vbf16-final\",\n", - ")\n", - "TRAINING_DATA = Path(\n", - " os.environ.get(\"TRAINING_DATA\", str(DATA_ROOT / \"cosmos3_nano_dflash_train.jsonl\"))\n", - ").resolve()\n", - "OUTPUT_DIR = Path(\n", - " os.environ.get(\"OUTPUT_DIR\", str(REPO_ROOT / \"ckpts\" / \"cosmos3-nano-dflash-first-run\"))\n", - ").resolve()\n", - "\n", - "PAI_SHARDS = Path(\n", - " os.environ.get(\"PAI_SHARDS\", str(DATA_ROOT / \"pai_understanding_native_shards\"))\n", - ").resolve()\n", - "VQA_SHARD_PATH = Path(\n", - " os.environ.get(\"VQA_SHARD_PATH\", str(DATA_ROOT / \"vqa_v2_train_shards\"))\n", - ").resolve()\n", - "TEXT_SHARDS = Path(\n", - " os.environ.get(\"TEXT_SHARDS\", str(DATA_ROOT / \"specdec_multilingual_prompt_full_shards\"))\n", - ").resolve()\n", - "PAI_OUTPUT = Path(\n", - " os.environ.get(\"PAI_OUTPUT\", str(SPEC_ROOT / \"pai_understanding_synthetic_outputs\"))\n", - ").resolve()\n", - "VQA_OUTPUT = Path(\n", - " os.environ.get(\"VQA_OUTPUT\", str(SPEC_ROOT / \"vqa_v2_synthetic_outputs\"))\n", - ").resolve()\n", - "MULTILINGUAL_OUTPUT = Path(\n", - " os.environ.get(\n", - " \"MULTILINGUAL_OUTPUT\", str(SPEC_ROOT / \"specdec_multilingual_prompt_synthetic_outputs\")\n", - " )\n", - ").resolve()\n", - "PLAIN_TEXT_INPUT = os.environ.get(\"PLAIN_TEXT_INPUT\", \"\")\n", - "\n", - "PAI_REVISION = os.environ.get(\"PAI_REVISION\", \"\")\n", - "TEXT_PROMPT_REVISION = os.environ.get(\"TEXT_PROMPT_REVISION\", \"\")\n", - "PAI_NUM_GENERATION_SHARDS = int(os.environ.get(\"PAI_NUM_GENERATION_SHARDS\", \"5\"))\n", - "PAI_LINES_PER_SHARD = int(os.environ.get(\"PAI_LINES_PER_SHARD\", \"128\"))\n", - "PAI_SHUFFLE_SEED = int(os.environ.get(\"PAI_SHUFFLE_SEED\", \"42\"))\n", - "VQA_NUM_SAMPLES = int(os.environ.get(\"VQA_NUM_SAMPLES\", \"20000\"))\n", - "VQA_SHUFFLE_SEED = int(os.environ.get(\"VQA_SHUFFLE_SEED\", \"42\"))\n", - "DEDUP_WORD_OVERLAP = float(os.environ.get(\"DEDUP_WORD_OVERLAP\", \"0.90\"))\n", - "DEDUP_CACHE_CONTEXTS = int(os.environ.get(\"DEDUP_CACHE_CONTEXTS\", \"25000\"))\n", - "MERGE_WORKERS = int(os.environ.get(\"MERGE_WORKERS\", \"8\"))\n", - "\n", - "if PAI_NUM_GENERATION_SHARDS <= 0 or PAI_LINES_PER_SHARD <= 0:\n", - " raise ValueError(\"PAI_NUM_GENERATION_SHARDS and PAI_LINES_PER_SHARD must be positive.\")\n", - "if not 0 < DEDUP_WORD_OVERLAP <= 1:\n", - " raise ValueError(\"DEDUP_WORD_OVERLAP must be in (0, 1].\")\n", - "if DEDUP_CACHE_CONTEXTS <= 0 or MERGE_WORKERS <= 0:\n", - " raise ValueError(\"DEDUP_CACHE_CONTEXTS and MERGE_WORKERS must be positive.\")\n", - "\n", - "for name, value in {\n", - " \"REPO_ROOT\": REPO_ROOT,\n", - " \"DATA_ROOT\": DATA_ROOT,\n", - " \"MODEL_PATH\": MODEL_PATH,\n", - " \"TRAINING_DATA\": TRAINING_DATA,\n", - " \"OUTPUT_DIR\": OUTPUT_DIR,\n", - " \"PAI_SHARDS\": PAI_SHARDS,\n", - " \"VQA_SHARD_PATH\": VQA_SHARD_PATH,\n", - " \"TEXT_SHARDS\": TEXT_SHARDS,\n", - " \"PAI_OUTPUT\": PAI_OUTPUT,\n", - " \"VQA_OUTPUT\": VQA_OUTPUT,\n", - " \"MULTILINGUAL_OUTPUT\": MULTILINGUAL_OUTPUT,\n", - " \"PLAIN_TEXT_INPUT\": PLAIN_TEXT_INPUT,\n", - " \"PAI_REVISION\": PAI_REVISION,\n", - " \"TEXT_PROMPT_REVISION\": TEXT_PROMPT_REVISION,\n", - " \"PYTHON_BIN\": sys.executable,\n", - " \"PAI_NUM_GENERATION_SHARDS\": PAI_NUM_GENERATION_SHARDS,\n", - " \"PAI_LINES_PER_SHARD\": PAI_LINES_PER_SHARD,\n", - " \"PAI_SHUFFLE_SEED\": PAI_SHUFFLE_SEED,\n", - " \"VQA_NUM_SAMPLES\": VQA_NUM_SAMPLES,\n", - " \"VQA_SHUFFLE_SEED\": VQA_SHUFFLE_SEED,\n", - " \"DEDUP_WORD_OVERLAP\": DEDUP_WORD_OVERLAP,\n", - " \"DEDUP_CACHE_CONTEXTS\": DEDUP_CACHE_CONTEXTS,\n", - " \"MERGE_WORKERS\": MERGE_WORKERS,\n", - "}.items():\n", - " os.environ[name] = str(value)\n", - "\n", - "print(f\"Training data: {TRAINING_DATA}\")\n", - "print(f\"Training output: {OUTPUT_DIR}\")" - ] - }, - { - "cell_type": "markdown", - "id": "8763a12b2bbd4a93a75aff182afb95dc", - "metadata": {}, - "source": [ - "### Sources and Workflow\n", - "\n", - "DFlash benefits most from target-model completions rather than a large collection of human-written answers. This build combines four complementary sources:\n", - "\n", - "1. **Representative usage samples** – PAI-Understanding video prompts.\n", - "2. **Visual reasoning data** – VQA v2 image-question prompts.\n", - "3. **Curated text** – either selected Nemotron Chat conversations or approved user data.\n", - "4. **High-quality text prompts** – multilingual prompt completions.\n", - "\n", - "The PAI and VQA sources retain their media; curated and multilingual text are text-only. The merge resolves media to absolute paths and conservatively deduplicates within matching prompt/media groups. This is why training uses `data.vlm_img_dir=/`.\n", - "\n", - "To create the Curated Nemotron Chat component, follow `recipes/train_eagle_head_cosmos_reason2.ipynb`: accept the Hugging Face licences for `nvidia/Nemotron-Post-Training-Dataset-v2`, then run `python ../prepare_input_conversations/add_nemotron_chat.py --mapping-file nemotron_mapping.bin` from that notebook. It writes `input_conversations/nemotron-chat.jsonl` with exactly 89,511 conversations. Set `PLAIN_TEXT_INPUT` to that file or to a privacy-reviewed JSONL of assistant-completed user conversations in the same `messages`/`conversations` format. To include both, combine them into one valid JSONL first.\n", - "\n", - "For DFlash, `training_seq_len` must be divisible by `dflash_block_size`; the production job uses `16384` and `8`.\n" - ] - }, - { - "cell_type": "markdown", - "id": "7623eae2785240b9bd12b16a66d81610", - "metadata": {}, - "source": [ - "### Representative Usage: PAI-Understanding\n", - "\n", - "Use a modest representative sample. In production, replace this benchmark with a privacy-reviewed sample of real requests when available. The target model produces the assistant response used for training. The CPU-only setup cell creates a deterministic shuffled sample. Its default produces five shards; set `PAI_NUM_GENERATION_SHARDS` to a multiple of the number of allocated nodes.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7cdc8c89c7104fffa095e18ddfef8986", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "# CPU-only node: download PAI-Understanding and make native-video prompt shards.\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${PAI_NUM_GENERATION_SHARDS:?Run the Step 2 configuration cell before this cell.}\"\n", - "PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding}\n", - "PAI_SHARDS=${PAI_SHARDS:-$SPEC/CR3_data/pai_understanding_native_shards}\n", - "PAI_LINES_PER_SHARD=${PAI_LINES_PER_SHARD:-128}\n", - "PAI_NUM_SAMPLES=${PAI_NUM_SAMPLES:-$((PAI_NUM_GENERATION_SHARDS * PAI_LINES_PER_SHARD))}\n", - "PAI_REVISION_ARGS=()\n", - "[ -n \"${PAI_REVISION:-}\" ] && PAI_REVISION_ARGS+=(--revision \"$PAI_REVISION\")\n", - "HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache}\n", - "\n", - "# Invoke the CLI through the active kernel interpreter to avoid a stale `hf` executable on PATH.\n", - "\"$PYTHON_BIN\" -m huggingface_hub.cli.hf --version\n", - "# If networking on the CPU-only node stalls, retry this command with HF_HUB_DISABLE_XET=1.\n", - "\"$PYTHON_BIN\" -m huggingface_hub.cli.hf download shi-labs/physical-ai-bench-understanding \\\n", - " --repo-type dataset \\\n", - " --local-dir \"$PAI_ROOT\" \\\n", - " \"${PAI_REVISION_ARGS[@]}\" \\\n", - " --max-workers 8\n", - "\n", - "HF_DATASETS_CACHE=\"$HF_DATASETS_CACHE\" \"$PYTHON_BIN\" \"$SPEC/recipes/prepare_multimodal_synthetic_shards.py\" \\\n", - " --dataset pai_understanding \\\n", - " --dataset_dir \"$PAI_ROOT\" \\\n", - " --media_root \"$PAI_ROOT\" \\\n", - " --output_dir \"$PAI_SHARDS\" \\\n", - " --max_lines_per_shard \"$PAI_LINES_PER_SHARD\" \\\n", - " --num_samples \"$PAI_NUM_SAMPLES\" \\\n", - " --shuffle_seed \"$PAI_SHUFFLE_SEED\" \\\n", - " --overwrite\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b118ea5561624da68c537baed56e602f", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "# Compute node: run a representative PAI shard slice in an existing Slurm GPU allocation.\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${MODEL_PATH:?Set MODEL_PATH and run Step 2 before this compute-node cell.}\"\n", - "PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding}\n", - "export MODEL_PATH\n", - "export DATASET_DIR=\"$PAI_ROOT\"\n", - "export MEDIA_ROOT=\"$PAI_ROOT\"\n", - "export SHARD_PATH=${PAI_SHARDS:-$SPEC/CR3_data/pai_understanding_native_shards}\n", - "export PREPARE_SHARDS=0\nexport OUTPUT_PATH=\"${PAI_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n", - "NODE_NAMES=\"$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | paste -sd, -)\"\n", - ": \"${PAI_NUM_GENERATION_SHARDS:?Run the Step 2 configuration cell before this cell.}\"\n", - "PAI_START_SHARD=${PAI_START_SHARD:-0}\n", - "PAI_NUM_AVAILABLE_SHARDS=$(find \"$SHARD_PATH\" -maxdepth 1 -type f -name 'train-*.jsonl' | wc -l)\n", - "PAI_NUM_NODES=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | wc -l)\n", - "[ \"$PAI_NUM_AVAILABLE_SHARDS\" -gt 0 ] || { echo \"No PAI shards found in $SHARD_PATH\" >&2; exit 1; }\n", - "[ \"$PAI_NUM_NODES\" -gt 0 ] || { echo \"No allocated nodes found\" >&2; exit 1; }\n", - "(( PAI_START_SHARD + PAI_NUM_GENERATION_SHARDS <= PAI_NUM_AVAILABLE_SHARDS )) || { echo \"Requested PAI shard range exceeds $PAI_NUM_AVAILABLE_SHARDS available shards\" >&2; exit 1; }\n", - "(( PAI_NUM_GENERATION_SHARDS % PAI_NUM_NODES == 0 )) || { echo \"$PAI_NUM_GENERATION_SHARDS selected PAI shards cannot be distributed evenly over $PAI_NUM_NODES nodes\" >&2; exit 1; }\n", - "PAI_SHARDS_PER_NODE=$(( PAI_NUM_GENERATION_SHARDS / PAI_NUM_NODES ))\n", - "\n", - "\"$SPEC/recipes/run_multimodal_synthetic_generation.sh\" \\\n", - " pai_understanding \"$SLURM_JOB_ID\" \"$PAI_START_SHARD\" \"$PAI_SHARDS_PER_NODE\" \"$NODE_NAMES\"\n" - ] - }, - { - "cell_type": "markdown", - "id": "938c804e27f84196a10c8828c723f798", - "metadata": {}, - "source": [ - "### Visual Reasoning: VQA v2\n", - "\n", - "This component is larger than the representative sample. The default creates a deterministic sample of 20,000 train-split prompts; increase `VQA_NUM_SAMPLES` for more visual coverage.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "504fb2a444614c0babb325280ed9130a", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "# CPU-only node: fetch the VQA v2 train questions, annotations, and COCO images.\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n", - "VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2}\n", - "IMAGE_ROOT=\"$VQA_ROOT/images\"\n", - ": \"${VQA_NUM_SAMPLES:?Run the Step 2 configuration cell before this cell.}\"\n", - "VQA_SHARD_PATH=${VQA_SHARD_PATH:-$SPEC/CR3_data/vqa_v2_train_shards}\n", - "HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache}\n", - "VQA_QUESTIONS_URL=${VQA_QUESTIONS_URL:-https://cvmlp.s3.amazonaws.com/vqa/mscoco/vqa/v2_Questions_Train_mscoco.zip}\n", - "VQA_ANNOTATIONS_URL=${VQA_ANNOTATIONS_URL:-https://cvmlp.s3.amazonaws.com/vqa/mscoco/vqa/v2_Annotations_Train_mscoco.zip}\n", - "COCO_TRAIN_URL=${COCO_TRAIN_URL:-https://images.cocodataset.org/zips/train2014.zip}\n", - "\n", - "extract_zip() {\n", - " local archive=$1 destination=$2 marker=$3\n", - " if [ -f \"$marker\" ]; then\n", - " echo \"Already extracted: $archive\"\n", - " return\n", - " fi\n", - " if command -v unzip >/dev/null 2>&1; then\n", - " unzip -n \"$archive\" -d \"$destination\"\n", - " else\n", - " # Minimal CPU-only node images may omit `unzip`; Python provides a compatible fallback.\n", - " \"$PYTHON_BIN\" -m zipfile -e \"$archive\" \"$destination\"\n", - " fi\n", - " touch \"$marker\"\n", - "}\n", - "\n", - "mkdir -p \"$VQA_ROOT\" \"$IMAGE_ROOT\"\n", - "curl --proto '=https' -L --fail --retry 5 -C - -o \"$VQA_ROOT/v2_Questions_Train_mscoco.zip\" \\\n", - " \"$VQA_QUESTIONS_URL\"\n", - "curl --proto '=https' -L --fail --retry 5 -C - -o \"$VQA_ROOT/v2_Annotations_Train_mscoco.zip\" \\\n", - " \"$VQA_ANNOTATIONS_URL\"\n", - "curl --proto '=https' -L --fail --retry 5 -C - -o \"$IMAGE_ROOT/train2014.zip\" \\\n", - " \"$COCO_TRAIN_URL\"\n", - "extract_zip \"$VQA_ROOT/v2_Questions_Train_mscoco.zip\" \"$VQA_ROOT\" \"$VQA_ROOT/.questions.extracted\"\n", - "extract_zip \"$VQA_ROOT/v2_Annotations_Train_mscoco.zip\" \"$VQA_ROOT\" \"$VQA_ROOT/.annotations.extracted\"\n", - "extract_zip \"$IMAGE_ROOT/train2014.zip\" \"$IMAGE_ROOT\" \"$IMAGE_ROOT/.train2014.extracted\"\n", - "\n", - "HF_DATASETS_CACHE=\"$HF_DATASETS_CACHE\" \"$PYTHON_BIN\" \"$SPEC/recipes/prepare_multimodal_synthetic_shards.py\" \\\n", - " --dataset vqa_v2 \\\n", - " --vqa_root \"$VQA_ROOT\" \\\n", - " --image_root \"$IMAGE_ROOT\" \\\n", - " --vqa_splits train \\\n", - " --num_samples \"$VQA_NUM_SAMPLES\" \\\n", - " --shuffle_seed \"$VQA_SHUFFLE_SEED\" \\\n", - " --output_dir \"$VQA_SHARD_PATH\" \\\n", - " --overwrite\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "59bbdb311c014d738909a11f9e486628", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "# Compute node: distribute all remaining prepared VQA shards across the allocated nodes.\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${MODEL_PATH:?Set MODEL_PATH and run Step 2 before this compute-node cell.}\"\n", - "export MODEL_PATH\n", - "VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2}\n", - "IMAGE_ROOT=\"$VQA_ROOT/images\"\n", - "export VQA_ROOT IMAGE_ROOT\n", - "export SHARD_PATH=${VQA_SHARD_PATH:-$SPEC/CR3_data/vqa_v2_train_shards}\n", - "export PREPARE_SHARDS=0\nexport OUTPUT_PATH=\"${VQA_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n", - "NODE_NAMES=\"$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | paste -sd, -)\"\n", - "VQA_NUM_SHARDS=$(find \"$SHARD_PATH\" -maxdepth 1 -type f -name '*.jsonl' | wc -l)\n", - "VQA_NUM_NODES=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | wc -l)\n", - "VQA_START_SHARD=${VQA_START_SHARD:-0}\n", - "[ \"$VQA_NUM_SHARDS\" -gt 0 ] || { echo \"No VQA shards found in $SHARD_PATH\" >&2; exit 1; }\n", - "[ \"$VQA_NUM_NODES\" -gt 0 ] || { echo \"No allocated nodes found\" >&2; exit 1; }\n", - "[[ \"$VQA_START_SHARD\" =~ ^[0-9]+$ ]] || { echo \"VQA_START_SHARD must be a non-negative integer: $VQA_START_SHARD\" >&2; exit 1; }\n", - "(( VQA_START_SHARD >= 0 && VQA_START_SHARD < VQA_NUM_SHARDS )) || { echo \"VQA_START_SHARD=$VQA_START_SHARD is outside 0 through $((VQA_NUM_SHARDS - 1))\" >&2; exit 1; }\n", - "VQA_REMAINING_SHARDS=$(( VQA_NUM_SHARDS - VQA_START_SHARD ))\n", - "# The multimodal worker skips its small rounded-up tail, so all remaining shards are covered.\n", - "VQA_SHARDS_PER_NODE=$(( (VQA_REMAINING_SHARDS + VQA_NUM_NODES - 1) / VQA_NUM_NODES ))\n", - "\n", - "\"$SPEC/recipes/run_multimodal_synthetic_generation.sh\" \\\n", - " vqa_v2 \"$SLURM_JOB_ID\" \"$VQA_START_SHARD\" \"$VQA_SHARDS_PER_NODE\" \"$NODE_NAMES\"\n" - ] - }, - { - "cell_type": "markdown", - "id": "b43b363d81ae4b689946ece5c682cd59", - "metadata": {}, - "source": [ - "### High-Quality Text: Synthetic Multilingual Prompts\n", - "\n", - "Use a large text-only component unless the deployment is deliberately narrow. `nvidia/Speculative-Decoding-Multilingual-Prompt-v2` provides prompts and Cosmos3 Nano generates the assistant answers. Eight temperatures increase response diversity.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "c3933fab20d04ec698c2621248eb3be0", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "# CPU-only node: download the text prompts and create full prompt shards.\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n", - "TEXT_ROOT=${TEXT_ROOT:-$DATA_ROOT/specdec_multilingual_prompt}\n", - "TEXT_SHARDS=${TEXT_SHARDS:-$SPEC/CR3_data/specdec_multilingual_prompt_full_shards}\n", - "TEXT_REVISION_ARGS=()\n", - "[ -n \"${TEXT_PROMPT_REVISION:-}\" ] && TEXT_REVISION_ARGS+=(--revision \"$TEXT_PROMPT_REVISION\")\n", - "HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache}\n", - "\n", - "# Only default.jsonl is used below; the repository's convenience samples are intentionally not downloaded.\n", - "# If networking on the CPU-only node stalls, retry this command with HF_HUB_DISABLE_XET=1.\n", - "\"$PYTHON_BIN\" -m huggingface_hub.cli.hf download nvidia/Speculative-Decoding-Multilingual-Prompt-v2 default.jsonl \\\n", - " --repo-type dataset \\\n", - " --local-dir \"$TEXT_ROOT\" \\\n", - " \"${TEXT_REVISION_ARGS[@]}\" \\\n", - " --max-workers 8\n", - "\n", - "HF_DATASETS_CACHE=\"$HF_DATASETS_CACHE\" \"$PYTHON_BIN\" \"$SPEC/recipes/prepare_multimodal_synthetic_shards.py\" \\\n", - " --dataset specdec_multilingual_prompt \\\n", - " --text_data \"$TEXT_ROOT\" \\\n", - " --max_lines_per_shard 1024 \\\n", - " --overwrite \\\n", - " --output_dir \"$TEXT_SHARDS\"\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4dd4641cc4064e0191573fe9c69df29b", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "# Compute node: use the allocation dynamically; no fixed node list is needed.\n", - "# START_SHARD=1059 is a resume example. Set it to 0 for a new full run.\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${MODEL_PATH:?Set MODEL_PATH and run Step 2 before this compute-node cell.}\"\n", - "export MODEL_PATH\n", - "TEXT_ROOT=${TEXT_ROOT:-$DATA_ROOT/specdec_multilingual_prompt}\n", - "export TEXT_DATA=\"$TEXT_ROOT\"\n", - "export SHARD_PATH=${TEXT_SHARDS:-$SPEC/CR3_data/specdec_multilingual_prompt_full_shards}\n", - "export OUTPUT_PATH=\"${MULTILINGUAL_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n", - "export PREPARE_SHARDS=0\n", - "export BACKEND=vllm\n", - "export CONTAINER_IMAGE=vllm/vllm-openai:v0.24.0\n", - "export SGLANG_TP_SIZE=1\n", - "export NUM_TEMPERATURES=8\n", - "NODE_NAMES=\"$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | paste -sd, -)\"\n", - "START_SHARD=${START_SHARD:-0}\n", - "TEXT_NUM_SHARDS=$(find \"$SHARD_PATH\" -maxdepth 1 -type f -name 'train-*.jsonl' | wc -l)\n", - "TEXT_NUM_NODES=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | wc -l)\n", - "[ \"$TEXT_NUM_SHARDS\" -gt 0 ] || { echo \"No text shards found in $SHARD_PATH\" >&2; exit 1; }\n", - "[ \"$TEXT_NUM_NODES\" -gt 0 ] || { echo \"No allocated nodes found\" >&2; exit 1; }\n", - "[ \"$START_SHARD\" -lt \"$TEXT_NUM_SHARDS\" ] || { echo \"START_SHARD=$START_SHARD is beyond $TEXT_NUM_SHARDS text shards\" >&2; exit 1; }\n", - "TEXT_REMAINING_SHARDS=$(( TEXT_NUM_SHARDS - START_SHARD ))\n", - "# The text worker skips its small rounded-up tail, so all remaining shards are covered.\n", - "JOBS_PER_NODE=${JOBS_PER_NODE:-$(( (TEXT_REMAINING_SHARDS + TEXT_NUM_NODES - 1) / TEXT_NUM_NODES ))}\n", - "\n", - "\"$SPEC/recipes/run_multimodal_synthetic_generation.sh\" \\\n", - " specdec_multilingual_prompt \"$SLURM_JOB_ID\" \"$START_SHARD\" \"$JOBS_PER_NODE\" \"$NODE_NAMES\"\n" - ] - }, - { - "cell_type": "markdown", - "id": "8309879909854d7188b41380fd92a7c3", - "metadata": {}, - "source": [ - "### Merge, Resolve Media, and Deduplicate\n", - "\n", - "Run this cell after the PAI, VQA v2, and multilingual synthesis jobs finish and `PLAIN_TEXT_INPUT` is set. It merges all four required sources into one atomic JSONL with absolute media paths. `MERGE_WORKERS` controls parallel merge workers; temperature variants for one generated shard always stay together, and deduplication is local to each worker partition. The merger never overwrites an existing final JSONL; choose a new `TRAINING_DATA` path when intentionally rebuilding.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5d1c9edacf9a4d2fa6b2e44b73dd4c20", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "set -euo pipefail\n", - "SPEC=\"$(cd .. && pwd)\"\n", - ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n", - "PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding}\n", - "VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2}\n", - ": \"${PAI_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${VQA_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${PLAIN_TEXT_INPUT:?Set PLAIN_TEXT_INPUT in the Step 2 configuration cell before this cell.}\"\n", - ": \"${MULTILINGUAL_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${TRAINING_DATA:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${DEDUP_WORD_OVERLAP:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${DEDUP_CACHE_CONTEXTS:?Run the Step 2 configuration cell before this cell.}\"\n", - ": \"${MERGE_WORKERS:?Run the Step 2 configuration cell before this cell.}\"\n", - "\n", - "\"$PYTHON_BIN\" \"$SPEC/recipes/merge_dflash_datasets.py\" \\\n", - " --source \"pai_understanding=$PAI_OUTPUT\" \\\n", - " --source \"vqa_v2=$VQA_OUTPUT\" \\\n", - " --source \"curated_text=$PLAIN_TEXT_INPUT\" \\\n", - " --source \"specdec_multilingual_prompt=$MULTILINGUAL_OUTPUT\" \\\n", - " --media-root \"pai_understanding=$PAI_ROOT\" \\\n", - " --media-root \"vqa_v2=$VQA_ROOT/images\" \\\n", - " --output \"$TRAINING_DATA\" \\\n", - " --jobs \"$MERGE_WORKERS\" \\\n", - " --word-overlap \"$DEDUP_WORD_OVERLAP\" \\\n", - " --cache-contexts \"$DEDUP_CACHE_CONTEXTS\"\n", - "\n", - "wc -l \"$TRAINING_DATA\"\n" - ] - }, - { - "cell_type": "markdown", - "id": "cb1e1581032b452c9409d6c6813c49d1", - "metadata": {}, - "source": [ - "## Step 3 – Submit the DFlash Training Job\n", - "\n", - "This self-contained cell submits the training job directly to Slurm. Its resource request, container setup, and training command live together here; no separate `job_*.sh` file is required. The example requests eight GPUs for four hours and saves every 1,000 updates. For a single-GPU run, request one GPU and raise `training.gradient_accumulation_steps` from `2` to `16` to keep the same global batch size.\n", - "\n", - "The defaults use the JSONL built in Step 2 and write to `ckpts/cosmos3-nano-dflash-first-run`. Set `MODEL_PATH`, `TRAINING_DATA`, or `OUTPUT_DIR` before submitting to override them. Remote code is disabled by default; set `TRUST_REMOTE_CODE=true` only for a model checkpoint you trust and that requires custom code.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "379cbbc1e968416e875cc15c1202d7eb", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "set -euo pipefail\n", - "REPO_ROOT=\"$(cd ../../.. && pwd)\"\n", - "\n", - ": \"${MODEL_PATH:?Run Step 1 before this cell.}\"\n", - ": \"${TRAINING_DATA:?Run Steps 1 and 2 before this cell.}\"\n", - ": \"${OUTPUT_DIR:?Run Step 1 before this cell.}\"\n", - "[ -d \"$MODEL_PATH\" ] || { echo \"Missing MODEL_PATH: $MODEL_PATH\" >&2; exit 1; }\n", - "[ -s \"$TRAINING_DATA\" ] || { echo \"Missing training JSONL: $TRAINING_DATA\" >&2; exit 1; }\n", - "\n", - "export TRUST_REMOTE_CODE=\"${TRUST_REMOTE_CODE:-false}\"\n", - "export REPO_ROOT MODEL_PATH TRAINING_DATA OUTPUT_DIR TRUST_REMOTE_CODE\n", - "export HF_HOME=\"${HF_HOME:-/lustre/fsw/portfolios/coreai/users/skierat/cache_hf}\"\n", - "\n", - "# Submit this heredoc directly; no batch-script file is created.\n", - "sbatch --export=ALL <<'SBATCH'\n", - "#!/usr/bin/env bash\n", - "#SBATCH -p interactive\n", - "#SBATCH --nodes=1\n", - "#SBATCH --ntasks-per-node=1\n", - "#SBATCH --gpus-per-node=8\n", - "#SBATCH --time=04:00:00\n", - "#SBATCH --account=coreai_tritoninference_triton3\n", - "#SBATCH --job-name=cosmos3-dflash\n", - "#SBATCH --signal=TERM@120\n", - "\n", - "set -euo pipefail\n", - "trap \"kill -- -$$\" TERM INT HUP\n", - "\n", - "export NUM_GPUS=8\n", - "export PIP_CACHE_DIR=\"${PIP_CACHE_DIR:-/tmp/pip-cache-${SLURM_JOB_ID}}\"\n", - "export TRITON_CACHE_DIR=\"${TRITON_CACHE_DIR:-/tmp/triton-cache-${SLURM_JOB_ID}}\"\n", - "export VLM_MIN_PIXELS=\"${VLM_MIN_PIXELS:-50176}\"\n", - "export VLM_MAX_PIXELS=\"${VLM_MAX_PIXELS:-802816}\"\n", - "export VLM_MAX_ASSISTANT_TOKENS=\"${VLM_MAX_ASSISTANT_TOKENS:-2048}\"\n", - "export VLM_MAX_PROMPT_TOKENS=\"${VLM_MAX_PROMPT_TOKENS:-8192}\"\n", - "export VLM_VIDEO_MIN_PIXELS=\"${VLM_VIDEO_MIN_PIXELS:-100352}\"\n", - "export VLM_VIDEO_MAX_PIXELS=\"${VLM_VIDEO_MAX_PIXELS:-2097152}\"\n", - "\n", - "echo \"Model: ${MODEL_PATH}\"\n", - "echo \"Training data: ${TRAINING_DATA}\"\n", - "echo \"Output: ${OUTPUT_DIR}\"\n", - "\n", - "srun \\\n", - " --ntasks=1 \\\n", - " --gpus-per-task=\"${NUM_GPUS}\" \\\n", - " --kill-on-bad-exit=1 \\\n", - " --container-image=/lustre/fsw/portfolios/coreai/users/skierat/containers/recent_vllm.sqfs \\\n", - " --container-workdir=/ \\\n", - " --container-mounts=/lustre:/lustre \\\n", - " bash -lc '\n", - " set -euo pipefail\n", - " export WANDB_MODE=disabled\n", - " export TOKENIZERS_PARALLELISM=false\n", - " unset NCCL_ASYNC_ERROR_HANDLING\n", - " export TORCH_NCCL_ASYNC_ERROR_HANDLING=1\n", - " export PYTHONPATH=\"${REPO_ROOT}:${PYTHONPATH:-}\"\n", - " mkdir -p \"${PIP_CACHE_DIR}\" \"${TRITON_CACHE_DIR}\"\n", - "\n", - " cd \"${REPO_ROOT}\"\n", - " python3 -m pip install -r examples/speculative_decoding/requirements.txt\n", - " python3 -m pip install torchcodec\n", - " python3 -m pip install -e \".[hf]\"\n", - "\n", - " python3 -m torch.distributed.run \\\n", - " --standalone \\\n", - " --nproc_per_node=\"${NUM_GPUS}\" \\\n", - " examples/speculative_decoding/main.py \\\n", - " --config modelopt_recipes/general/speculative_decoding/dflash.yaml \\\n", - " model.model_name_or_path=\"${MODEL_PATH}\" \\\n", - " model.trust_remote_code=\"${TRUST_REMOTE_CODE:-false}\" \\\n", - " data.data_path=\"${TRAINING_DATA}\" \\\n", - " data.vlm_processor=\"${MODEL_PATH}\" \\\n", - " data.vlm_img_dir=/ \\\n", - " training.output_dir=\"${OUTPUT_DIR}\" \\\n", - " training.num_train_epochs=25 \\\n", - " training.per_device_train_batch_size=1 \\\n", - " training.gradient_accumulation_steps=2 \\\n", - " training.training_seq_len=16384 \\\n", - " training.answer_only_loss=true \\\n", - " training.save_steps=1000 \\\n", - " training.save_total_limit=10 \\\n", - " training.logging_steps=10 \\\n", - " training.dataloader_num_workers=2 \\\n", - " training.dataloader_prefetch_factor=2 \\\n", - " training.ddp_find_unused_parameters=false \\\n", - " training.report_to=none \\\n", - " dflash.dflash_block_size=8 \\\n", - " dflash.dflash_num_anchors=128 \\\n", - " dflash.dflash_loss_objective=decay \\\n", - " dflash.dflash_loss_decay_factor=4 \\\n", - " dflash.dflash_architecture_config.num_hidden_layers=5 \\\n", - " dflash.dflash_architecture_config.num_attention_heads=32 \\\n", - " dflash.dflash_architecture_config.num_key_value_heads=8 \\\n", - " dflash.dflash_architecture_config.head_dim=128 \\\n", - " dflash.dflash_architecture_config.intermediate_size=12288 \\\n", - " dflash.dflash_architecture_config.max_position_embeddings=262144 \\\n", - " dflash.dflash_architecture_config.rms_norm_eps=1e-06 \\\n", - " dflash.dflash_architecture_config.rope_theta=5000000 \\\n", - " dflash.dflash_mask_token_id=151669\n", - " '\n", - "SBATCH\n" - ] - }, - { - "cell_type": "markdown", - "id": "277c27b1587741f2af2001be3712ef0d", - "metadata": {}, - "source": [ - "## Step 4 – Export the DFlash Checkpoint\n", - "\n", - "Export one completed checkpoint to the compact DFlash Hugging Face format consumed by vLLM. Run this cell in the same container or Python environment used for training (or another environment with ModelOpt and the example dependencies installed).\n", - "\n", - "Set `CKPT_DIR` to a saved checkpoint directory such as `checkpoint-180000`, not the top-level run directory. The export reads `modelopt_state.pth` from that directory and writes a new deployment directory containing `config.json` and `model.safetensors`. The cell refuses to overwrite an existing export.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "db7b79bc585a40fcaf58bf750017e135", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "set -euo pipefail\n", - "REPO_ROOT=\"$(cd ../../.. && pwd)\"\n", - "RUN_DIR=${OUTPUT_DIR:-$REPO_ROOT/ckpts/cosmos3-nano-dflash-first-run}\n", - "CKPT_DIR=${CKPT_DIR:?Set CKPT_DIR to a saved checkpoint directory, for example $RUN_DIR/checkpoint-180000}\n", - "EXPORT_PATH=${EXPORT_PATH:-$RUN_DIR/export-$(basename \"$CKPT_DIR\" | sed 's/^checkpoint-//')}\n", - "\n", - "[ -f \"$CKPT_DIR/modelopt_state.pth\" ] || { echo \"Missing DFlash checkpoint: $CKPT_DIR\" >&2; exit 1; }\n", - "[ ! -e \"$EXPORT_PATH\" ] || { echo \"Export path already exists: $EXPORT_PATH\" >&2; exit 1; }\n", - "TRUST_REMOTE_CODE=\"${TRUST_REMOTE_CODE:-false}\"\n", - "EXPORT_ARGS=()\n", - "if [[ \"${TRUST_REMOTE_CODE,,}\" == \"true\" || \"$TRUST_REMOTE_CODE\" == \"1\" ]]; then\n", - " EXPORT_ARGS+=(--trust_remote_code)\n", - "fi\n", - "python3 \"$REPO_ROOT/examples/speculative_decoding/scripts/export_hf_checkpoint.py\" \\\n", - " --model_path \"$CKPT_DIR\" \\\n", - " --export_path \"$EXPORT_PATH\" \\\n", - " \"${EXPORT_ARGS[@]}\"\n", - "\n", - "[ -s \"$EXPORT_PATH/config.json\" ] && [ -s \"$EXPORT_PATH/model.safetensors\" ] \\\n", - " || { echo \"Export is incomplete: $EXPORT_PATH\" >&2; exit 1; }\n", - "echo \"DFlash deployment checkpoint: $EXPORT_PATH\"\n" - ] - }, - { - "cell_type": "markdown", - "id": "916684f9a58a4a2aa5f864670399430d", - "metadata": {}, - "source": [ - "## Deployment\n", - "\n", - "Serve the original Cosmos3 Nano target with the exported DFlash draft. A DFlash block size of eight yields seven speculative tokens: the remaining position is the context/bonus token.\n", - "\n", - "For a first deployment, use a vLLM environment compatible with this recipe. The command below is a single-GPU smoke test: it binds only to localhost and caps the context at 4096 tokens to keep startup and memory use bounded. Raise or remove `--max-model-len` only after sizing the desired context and concurrency. Add `--trust-remote-code` only when serving a model checkpoint you trust and that requires custom code.\n", - "\n", - "```bash\n", - "SERVED_MODEL_NAME=cosmos3-nano-dflash\n", - "vllm serve \"$MODEL_PATH\" \\\n", - " --host 127.0.0.1 --port 8000 \\\n", - " --max-model-len 4096 \\\n", - " --served-model-name \"$SERVED_MODEL_NAME\" \\\n", - " --speculative-config \"{\\\"method\\\":\\\"dflash\\\",\\\"model\\\":\\\"$EXPORT_PATH\\\",\\\"num_speculative_tokens\\\":7}\"\n", - "\n", - "# In another terminal, verify the server and make one request.\n", - "curl -fsS http://127.0.0.1:8000/health\n", - "curl -fsS http://127.0.0.1:8000/v1/models\n", - "curl -fsS http://127.0.0.1:8000/v1/chat/completions \\\n", - " -H 'Content-Type: application/json' \\\n", - " -d \"{\\\"model\\\":\\\"$SERVED_MODEL_NAME\\\",\\\"messages\\\":[{\\\"role\\\":\\\"user\\\",\\\"content\\\":\\\"Reply with exactly: DFlash deployment smoke test passed.\\\"}],\\\"temperature\\\":0,\\\"max_tokens\\\":16}\"\n", - "\n", - "# Stop the foreground server with Ctrl-C when the smoke test completes.\n", - "```\n" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.10.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/tools/launcher/common/specdec/merge_dflash_datasets.sh b/tools/launcher/common/specdec/merge_dflash_datasets.sh new file mode 100755 index 00000000000..6ef769c630a --- /dev/null +++ b/tools/launcher/common/specdec/merge_dflash_datasets.sh @@ -0,0 +1,52 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Merge generated sources into one DFlash training JSONL. +# +# Resolves every image/video reference to an absolute path (training then runs +# with data.vlm_img_dir=/) and drops near-duplicate completions, which the +# temperature sweep produces in bulk. +# +# All args pass through to merge_dflash_datasets.py, so sources stay in the YAML: +# script: common/specdec/merge_dflash_datasets.sh +# args: +# - --source pai_understanding=/scratchspace/data/pai_outputs +# - --media-root pai_understanding=/scratchspace/data/pai_understanding +# - --output /scratchspace/data/train.jsonl +# - --jobs 8 + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source ${SCRIPT_DIR}/../service_utils.sh + +trap 'error_handler $0 $LINENO' ERR +trap 'exit_handler' EXIT + +set -euo pipefail + +MERGE=modules/Model-Optimizer/examples/speculative_decoding/recipes/merge_dflash_datasets.py + +set -x +python3 "$MERGE" "$@" +set +x + +# Surface the merged row count; a silent drop to near-zero means the sources +# were empty or every record failed media resolution. +prev_arg="" +for arg in "$@"; do + [ "$prev_arg" = "--output" ] && wc -l "$arg" + prev_arg="$arg" +done diff --git a/tools/launcher/common/specdec/multimodal_prepare_shards.sh b/tools/launcher/common/specdec/multimodal_prepare_shards.sh new file mode 100755 index 00000000000..ee9336b1f4f --- /dev/null +++ b/tools/launcher/common/specdec/multimodal_prepare_shards.sh @@ -0,0 +1,164 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Download a multimodal source dataset and write prompt shards for synthetic +# generation. No GPU is used: this only fetches data and reshapes it into the +# train-%05d-%05d.jsonl layout the generation workers consume. +# +# One invocation prepares one DATASET. Run it once per source, then merge. +# +# Usage from YAML: +# script: common/specdec/multimodal_prepare_shards.sh +# args: +# - --dataset pai_understanding +# - --shard-path /scratchspace/data/pai_shards +# environment: +# - DATA_ROOT: /scratchspace/data +# +# Env: +# DATA_ROOT — download/extract root (default: /scratchspace/data) +# NUM_SAMPLES — cap records before sharding (default: dataset-specific) +# SHUFFLE_SEED — deterministic shuffle before slicing (default: 42) +# LINES_PER_SHARD — records per generated shard (default: 128; text: 1024) +# PAI_REPO_ID — override the PAI-Bench-U Hugging Face repo +# FORCE_DOWNLOAD — re-download PAI even if it is already materialized +# PAI_REVISION / TEXT_PROMPT_REVISION — pin a Hugging Face dataset revision + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source ${SCRIPT_DIR}/../service_utils.sh + +trap 'error_handler $0 $LINENO' ERR +trap 'exit_handler' EXIT + +set -euo pipefail + +DATASET="" +SHARD_PATH="" +while [[ $# -gt 0 ]]; do + case "$1" in + --dataset) DATASET="$2"; shift 2 ;; + --shard-path) SHARD_PATH="$2"; shift 2 ;; + *) echo "ERROR: unknown argument: $1" >&2; exit 1 ;; + esac +done +[ -n "$DATASET" ] && [ -n "$SHARD_PATH" ] \ + || { echo "ERROR: --dataset and --shard-path are required." >&2; exit 1; } + +DATA_ROOT=${DATA_ROOT:-/scratchspace/data} +SHUFFLE_SEED=${SHUFFLE_SEED:-42} +PREPARE=modules/Model-Optimizer/examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py +export HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache} + +pip install "huggingface-hub>=1.2.1" pillow +mkdir -p "$DATA_ROOT" + +# Direct `hf download` calls below go through `python3 -m huggingface_hub.cli.hf` +# rather than a bare `hf`, which may be a stale executable on PATH. If a download +# stalls, retry the task with HF_HUB_DISABLE_XET=1. +case "$DATASET" in +pai_understanding) + PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding} + LINES_PER_SHARD=${LINES_PER_SHARD:-128} + # PAI is sampled by shard count so the generation step gets a whole number + # of shards per node; see the YAML's PAI_NUM_GENERATION_SHARDS comment. + NUM_SAMPLES=${NUM_SAMPLES:-$(( ${NUM_GENERATION_SHARDS:-5} * LINES_PER_SHARD ))} + # prepare_multimodal_synthetic_shards.py downloads PAI itself via --download, + # so let it. Its snapshot_download has no revision argument, so pin a + # revision here instead when PAI_REVISION is set. + PAI_DOWNLOAD_ARGS=(--download) + if [ -n "${PAI_REVISION:-}" ]; then + python3 -m huggingface_hub.cli.hf download "${PAI_REPO_ID:-shi-labs/physical-ai-bench-understanding}" \ + --repo-type dataset --local-dir "$PAI_ROOT" --max-workers 8 \ + --revision "$PAI_REVISION" + PAI_DOWNLOAD_ARGS=() + fi + python3 "$PREPARE" \ + --dataset pai_understanding \ + --dataset_dir "$PAI_ROOT" \ + --media_root "$PAI_ROOT" \ + --output_dir "$SHARD_PATH" \ + --max_lines_per_shard "$LINES_PER_SHARD" \ + --num_samples "$NUM_SAMPLES" \ + --shuffle_seed "$SHUFFLE_SEED" \ + "${PAI_DOWNLOAD_ARGS[@]}" \ + ${PAI_REPO_ID:+--repo_id "$PAI_REPO_ID"} \ + ${FORCE_DOWNLOAD:+--force_download} \ + --overwrite + ;; +vqa_v2) + VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2} + IMAGE_ROOT="$VQA_ROOT/images" + LINES_PER_SHARD=${LINES_PER_SHARD:-128} + NUM_SAMPLES=${NUM_SAMPLES:-20000} + VQA_QUESTIONS_URL=${VQA_QUESTIONS_URL:-https://cvmlp.s3.amazonaws.com/vqa/mscoco/vqa/v2_Questions_Train_mscoco.zip} + VQA_ANNOTATIONS_URL=${VQA_ANNOTATIONS_URL:-https://cvmlp.s3.amazonaws.com/vqa/mscoco/vqa/v2_Annotations_Train_mscoco.zip} + COCO_TRAIN_URL=${COCO_TRAIN_URL:-https://images.cocodataset.org/zips/train2014.zip} + + extract_zip() { + local archive=$1 destination=$2 marker=$3 + if [ -f "$marker" ]; then + echo "Already extracted: $archive" + return + fi + if command -v unzip >/dev/null 2>&1; then + unzip -n "$archive" -d "$destination" + else + # Minimal container images may omit `unzip`. + python3 -m zipfile -e "$archive" "$destination" + fi + touch "$marker" + } + + mkdir -p "$VQA_ROOT" "$IMAGE_ROOT" + curl --proto '=https' -L --fail --retry 5 -C - -o "$VQA_ROOT/questions.zip" "$VQA_QUESTIONS_URL" + curl --proto '=https' -L --fail --retry 5 -C - -o "$VQA_ROOT/annotations.zip" "$VQA_ANNOTATIONS_URL" + curl --proto '=https' -L --fail --retry 5 -C - -o "$IMAGE_ROOT/train2014.zip" "$COCO_TRAIN_URL" + extract_zip "$VQA_ROOT/questions.zip" "$VQA_ROOT" "$VQA_ROOT/.questions.extracted" + extract_zip "$VQA_ROOT/annotations.zip" "$VQA_ROOT" "$VQA_ROOT/.annotations.extracted" + extract_zip "$IMAGE_ROOT/train2014.zip" "$IMAGE_ROOT" "$IMAGE_ROOT/.train2014.extracted" + + python3 "$PREPARE" \ + --dataset vqa_v2 \ + --vqa_root "$VQA_ROOT" \ + --image_root "$IMAGE_ROOT" \ + --vqa_splits "${VQA_SPLITS:-train}" \ + --num_samples "$NUM_SAMPLES" \ + --shuffle_seed "$SHUFFLE_SEED" \ + --output_dir "$SHARD_PATH" \ + --overwrite + ;; +specdec_multilingual_prompt) + TEXT_ROOT=${TEXT_ROOT:-$DATA_ROOT/specdec_multilingual_prompt} + # Only default.jsonl: the repo's sample-*.jsonl subsets are slices of it, so + # downloading everything would duplicate prompts. + python3 -m huggingface_hub.cli.hf download nvidia/Speculative-Decoding-Multilingual-Prompt-v2 \ + default.jsonl --repo-type dataset --local-dir "$TEXT_ROOT" --max-workers 8 \ + ${TEXT_PROMPT_REVISION:+--revision "$TEXT_PROMPT_REVISION"} + python3 "$PREPARE" \ + --dataset specdec_multilingual_prompt \ + --text_data "$TEXT_ROOT" \ + --max_lines_per_shard "${LINES_PER_SHARD:-1024}" \ + --output_dir "$SHARD_PATH" \ + --overwrite + ;; +*) + echo "ERROR: unsupported --dataset: $DATASET" >&2 + exit 1 + ;; +esac + +echo "Prepared $(find "$SHARD_PATH" -maxdepth 1 -name 'train-*.jsonl' | wc -l) shard(s) in $SHARD_PATH" diff --git a/tools/launcher/common/specdec/multimodal_synthetic_generation.sh b/tools/launcher/common/specdec/multimodal_synthetic_generation.sh new file mode 100755 index 00000000000..7b6ddf4e3d4 --- /dev/null +++ b/tools/launcher/common/specdec/multimodal_synthetic_generation.sh @@ -0,0 +1,127 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Generate target-model completions for prepared prompt shards. +# +# DFlash learns from the *target model's* own completions, not from human-written +# answers, so every prompt shard is replayed through the target and the responses +# become the training labels. +# +# Shard-to-node assignment is derived from the launcher's own allocation +# (SLURM_JOB_ID + SLURM_JOB_NODELIST); no job id or node list is passed in. +# Each node serves the target locally and processes its own slice of shards. +# +# Usage from YAML: +# script: common/specdec/multimodal_synthetic_generation.sh +# args: +# - --dataset vqa_v2 +# - --shard-path /scratchspace/data/vqa_shards +# - --output-path /scratchspace/data/vqa_outputs +# - --media-root /scratchspace/data/vqa_v2/images +# environment: +# - MODEL_PATH: /hf-local/nvidia/Cosmos3-Nano +# +# Env: +# MODEL_PATH — target checkpoint (required) +# NUM_SHARDS — shards to process; default: all remaining from START_SHARD +# START_SHARD — first shard index (default: 0; set to resume a partial run) +# SGLANG_TP_SIZE — TP per server (default 1 => one server per GPU) +# NUM_TEMPERATURES — temperature sweep width (default 8) +# BACKEND — vllm or sglang (default: sglang for media, vllm for text) + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source ${SCRIPT_DIR}/../service_utils.sh + +trap 'error_handler $0 $LINENO' ERR +trap 'exit_handler' EXIT + +set -euo pipefail + +DATASET="" +SHARD_PATH="" +OUTPUT_PATH="" +MEDIA_ROOT="" +while [[ $# -gt 0 ]]; do + case "$1" in + --dataset) DATASET="$2"; shift 2 ;; + --shard-path) SHARD_PATH="$2"; shift 2 ;; + --output-path) OUTPUT_PATH="$2"; shift 2 ;; + --media-root) MEDIA_ROOT="$2"; shift 2 ;; + *) echo "ERROR: unknown argument: $1" >&2; exit 1 ;; + esac +done + +[ -n "$DATASET" ] && [ -n "$SHARD_PATH" ] && [ -n "$OUTPUT_PATH" ] \ + || { echo "ERROR: --dataset, --shard-path and --output-path are required." >&2; exit 1; } +[ -n "${MODEL_PATH:-}" ] || { echo "ERROR: MODEL_PATH must name the target checkpoint." >&2; exit 1; } +[ -d "$SHARD_PATH" ] || { echo "ERROR: missing shard path: $SHARD_PATH" >&2; exit 1; } + +SPEC_ROOT=modules/Model-Optimizer/examples/speculative_decoding + +# The launcher allocates the nodes, so read them back instead of taking a list. +if [ -n "${SLURM_JOB_NODELIST:-}" ]; then + NODE_NAMES="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | paste -sd, -)" + NUM_NODES=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l) +else + NODE_NAMES="$(hostname)" + NUM_NODES=1 +fi +[ "$NUM_NODES" -gt 0 ] || { echo "ERROR: no allocated nodes found." >&2; exit 1; } + +AVAILABLE_SHARDS=$(find "$SHARD_PATH" -maxdepth 1 -type f -name 'train-*.jsonl' | wc -l) +[ "$AVAILABLE_SHARDS" -gt 0 ] || { echo "ERROR: no train-*.jsonl shards in $SHARD_PATH" >&2; exit 1; } + +START_SHARD=${START_SHARD:-0} +[[ "$START_SHARD" =~ ^[0-9]+$ ]] || { echo "ERROR: START_SHARD must be a non-negative integer." >&2; exit 1; } +(( START_SHARD < AVAILABLE_SHARDS )) || { echo "ERROR: START_SHARD=$START_SHARD is beyond $AVAILABLE_SHARDS shards." >&2; exit 1; } + +REMAINING=$(( AVAILABLE_SHARDS - START_SHARD )) +NUM_SHARDS=${NUM_SHARDS:-$REMAINING} +(( NUM_SHARDS > 0 && START_SHARD + NUM_SHARDS <= AVAILABLE_SHARDS )) \ + || { echo "ERROR: shard range [$START_SHARD, $((START_SHARD + NUM_SHARDS))) exceeds $AVAILABLE_SHARDS shards." >&2; exit 1; } + +# Round up so the tail shards are still assigned; workers skip a missing shard. +JOBS_PER_NODE=${JOBS_PER_NODE:-$(( (NUM_SHARDS + NUM_NODES - 1) / NUM_NODES ))} + +echo "Generating $DATASET: shards [$START_SHARD, $((START_SHARD + NUM_SHARDS))) over $NUM_NODES node(s), $JOBS_PER_NODE per node" + +export MODEL_PATH +export SHARD_PATH +export OUTPUT_PATH +export PREPARE_SHARDS=0 +export SGLANG_TP_SIZE=${SGLANG_TP_SIZE:-1} +export NUM_TEMPERATURES=${NUM_TEMPERATURES:-8} + +mkdir -p "$OUTPUT_PATH" + +if [ "$DATASET" = "specdec_multilingual_prompt" ]; then + export BACKEND=${BACKEND:-vllm} + export TEXT_DATA=${TEXT_DATA:-$SHARD_PATH} +else + # Media generation goes through SGLang's native image/video client. + [ -n "$MEDIA_ROOT" ] || { echo "ERROR: --media-root is required for $DATASET." >&2; exit 1; } + export BACKEND=sglang + export MEDIA_ROOT + export DATASET_DIR=${DATASET_DIR:-$MEDIA_ROOT} + export IMAGE_ROOT=${IMAGE_ROOT:-$MEDIA_ROOT} + export VQA_ROOT=${VQA_ROOT:-$MEDIA_ROOT} +fi + +bash "$SPEC_ROOT/recipes/run_multimodal_synthetic_generation.sh" \ + "$DATASET" "${SLURM_JOB_ID:-0}" "$START_SHARD" "$JOBS_PER_NODE" "$NODE_NAMES" + +echo "Wrote $(find "$OUTPUT_PATH" -maxdepth 1 -name '*.jsonl' | wc -l) output file(s) to $OUTPUT_PATH" diff --git a/tools/launcher/examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml b/tools/launcher/examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml new file mode 100644 index 00000000000..6e2e5461fb0 --- /dev/null +++ b/tools/launcher/examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml @@ -0,0 +1,205 @@ +# DFlash online speculative decoding training for Cosmos3 Nano (multimodal). +# +# Unlike the text-only DFlash examples, the draft here is trained on image and +# video conversations, so the pipeline first synthesizes its own training data. +# +# 6-step pipeline: +# task_0..2: Prepare prompt shards for three sources (download + reshape) +# task_3: Generate target completions for those shards +# task_4: Merge + deduplicate into one training JSONL +# task_5: Online DFlash training (exports every checkpoint) +# task_6: vLLM smoke test with DFlash speculative decoding +# +# Data sources — DFlash learns from the target model's own completions, not from +# human-written answers, so every source contributes prompts that are replayed +# through the target: +# PAI-Understanding — representative video usage +# VQA v2 — image visual reasoning +# Multilingual — high-quality text prompts +# Curated text — set CURATED_TEXT to a reviewed conversations JSONL to +# add a fourth source (see the recipes README). +# +# Cluster-neutral: node counts, container tags and /hf-local paths below are +# defaults, not requirements. Override per cluster on the command line. +# +# Usage: +# uv run launch.py --yaml examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml --yes + +job_name: Cosmos3-Nano_DFlash_online_multimodal +pipeline: + global_vars: + hf_model: /hf-local/nvidia/Cosmos3-Nano + data_root: /scratchspace/data + train_jsonl: /scratchspace/data/cosmos3_nano_dflash_train.jsonl + output_dir: /scratchspace/dflash_cosmos3_nano + + # Steps 1-3: prepare prompt shards. These only download and reshape data — no + # GPU work — so they take the smallest allocation the partition allows. + task_0: + script: common/specdec/multimodal_prepare_shards.sh + args: + - --dataset pai_understanding + - --shard-path <>/pai_shards + environment: + - DATA_ROOT: <> + # PAI is sampled by whole shards so task_3 can split them evenly per node. + - NUM_GENERATION_SHARDS: "5" + - LINES_PER_SHARD: "128" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + + task_1: + script: common/specdec/multimodal_prepare_shards.sh + args: + - --dataset vqa_v2 + - --shard-path <>/vqa_shards + environment: + - DATA_ROOT: <> + - NUM_SAMPLES: "20000" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + + task_2: + script: common/specdec/multimodal_prepare_shards.sh + args: + - --dataset specdec_multilingual_prompt + - --shard-path <>/text_shards + environment: + - DATA_ROOT: <> + - LINES_PER_SHARD: "1024" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + + # Step 4: generate target completions. Shards are split across the allocated + # nodes automatically; each node serves the target locally. + # + # SGLANG_TP_SIZE=1 runs one server per GPU, one temperature each, so a node + # sweeps NUM_TEMPERATURES in parallel. Raise TP (and drop NUM_TEMPERATURES to + # match) only if the target does not fit on a single GPU. + # + # Media sources need SGLang's native image/video client; text uses vLLM. Run + # this task once per source, overriding --dataset/--shard-path/--output-path + # (and --media-root for the media sources). + task_3: + script: common/specdec/multimodal_synthetic_generation.sh + args: + - --dataset vqa_v2 + - --shard-path <>/vqa_shards + - --output-path <>/vqa_outputs + - --media-root <>/vqa_v2/images + environment: + - MODEL_PATH: <> + - SGLANG_TP_SIZE: "1" + - NUM_TEMPERATURES: "8" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: lmsysorg/sglang:v0.5.3-cu129 + + # Step 5: merge. Media paths are resolved to absolute here, which is why + # training below can pass data.vlm_img_dir=/. The temperature sweep emits many + # near-identical completions per prompt, so dedup runs over the merged set. + task_4: + script: common/specdec/merge_dflash_datasets.sh + args: + - --source pai_understanding=<>/pai_outputs + - --source vqa_v2=<>/vqa_outputs + - --source specdec_multilingual_prompt=<>/text_outputs + - --media-root pai_understanding=<>/pai_understanding + - --media-root vqa_v2=<>/vqa_v2/images + - --output <> + - --jobs 8 + - --word-overlap 0.90 + - --overwrite + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + + # Step 6: online DFlash training. The VLM_* limits below cap text and visual + # token growth *before* tokenization; without them a high-resolution video can + # expand past training_seq_len, which the collator rejects rather than + # silently truncating. + task_5: + script: common/specdec/dflash_online_training.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml + - model.model_name_or_path=<> + # Cosmos3 Nano ships custom modeling code. + - model.trust_remote_code=true + - data.data_path=<> + # Setting vlm_processor is what selects the multimodal collator. + - data.vlm_processor=<> + # Merge already wrote absolute media paths. + - data.vlm_img_dir=/ + - training.output_dir=<> + - training.num_train_epochs=25 + - training.per_device_train_batch_size=1 + - training.gradient_accumulation_steps=2 + - training.training_seq_len=16384 + - training.answer_only_loss=true + - training.save_steps=1000 + - training.save_total_limit=10 + - training.logging_steps=10 + - training.dataloader_num_workers=2 + - training.dataloader_prefetch_factor=2 + - training.ddp_find_unused_parameters=false + # dflash.yaml defaults to tensorboard, which hard-fails if it is absent. + - training.report_to=none + # training_seq_len must be divisible by dflash_block_size (16384 / 8). + - dflash.dflash_block_size=8 + - dflash.dflash_num_anchors=128 + - dflash.dflash_loss_objective=decay + - dflash.dflash_loss_decay_factor=4 + # Qwen3-family tokenizer has no mask token; use a reserved vocab slot. + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_architecture_config.num_hidden_layers=5 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=12288 + - dflash.dflash_architecture_config.max_position_embeddings=262144 + - dflash.dflash_architecture_config.rms_norm_eps=1e-06 + - dflash.dflash_architecture_config.rope_theta=5000000 + environment: + # Per-image and per-video pixel bounds handed to the HF processor. + - VLM_MIN_PIXELS: "50176" + - VLM_MAX_PIXELS: "802816" + - VLM_VIDEO_MIN_PIXELS: "100352" + - VLM_VIDEO_MAX_PIXELS: "2097152" + # Text caps applied before media placeholders expand into tokens. + - VLM_MAX_PROMPT_TOKENS: "8192" + - VLM_MAX_ASSISTANT_TOKENS: "2048" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + + # Step 7: smoke test the exported draft under vLLM. + task_6: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_CKPT_DIR: <> + - SPEC_METHOD: "dflash" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.2" + slurm_config: + _factory_: "slurm_factory" + container: "vllm/vllm-openai:nightly" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1