fix(specdec): generate all multimodal sources and pad VLM batches

modelopt/torch/utils/plugins/transformers_dataset.py

@@ -604,6 +604,7 @@ def _apply_chat_template(self, examples):
             tokenize=True,
             return_tensors="pt",
             return_dict=True,
+            padding=True,
             add_generation_prompt=self.add_generation_prompt,
             return_assistant_tokens_mask=self.answer_only_loss and not derive_masks_from_markers,
         )

tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py

@@ -273,6 +273,7 @@ def test_vlm_collator_pads_template_output_and_builds_unshifted_labels():
         tokenize=True,
         return_tensors="pt",
         return_dict=True,
+        padding=True,
         add_generation_prompt=False,
         return_assistant_tokens_mask=True,
     )

tools/launcher/common/specdec/dflash_online_training.sh

@@ -40,6 +40,11 @@ source ${SCRIPT_DIR}/../service_utils.sh
 
 pip install -r modules/Model-Optimizer/examples/speculative_decoding/requirements.txt
 pip install huggingface-hub>=1.2.1
+# Multimodal datasets containing videos require TorchCodec. Keep this opt-in so
+# text-only launcher examples do not acquire an unnecessary runtime dependency.
+if [ "${INSTALL_TORCHCODEC:-0}" = "1" ]; then
+    pip install torchcodec
+fi
 export PATH=$PATH:/workspace/.local/bin
 
 # Some trust_remote_code MoE models pin an older transformers (e.g. MiniMax-M2.7
@@ -49,6 +54,16 @@ if [ -n "${OVERRIDE_TRANSFORMERS:-}" ]; then
     pip install "transformers==${OVERRIDE_TRANSFORMERS}"
 fi
 
+# Export must follow the training configuration: loading a checkpoint with
+# remote code is opt-in and defaults to disabled. The command-line model option
+# takes precedence so one launcher task has a single source of truth.
+TRUST_REMOTE_CODE=${TRUST_REMOTE_CODE:-false}
+for arg in "$@"; do
+    if [[ "$arg" == model.trust_remote_code=* ]]; then
+        TRUST_REMOTE_CODE=${arg#*=}
+    fi
+done
+
 ###################################################################################################
 
 trap 'error_handler $0 $LINENO' ERR
@@ -123,6 +138,11 @@ set +x
 
 # Export last checkpoint to deployment format (rank 0 only, single GPU)
 if [ "${SLURM_PROCID:-0}" = "0" ]; then
+    EXPORT_TRUST_REMOTE_CODE=()
+    if [[ "${TRUST_REMOTE_CODE,,}" == "true" || "$TRUST_REMOTE_CODE" == "1" ]]; then
+        EXPORT_TRUST_REMOTE_CODE+=(--trust_remote_code)
+    fi
+
     OUTPUT_DIR=$(python3 -c "
 import sys
 for arg in sys.argv[1:]:
@@ -145,7 +165,7 @@ for arg in sys.argv[1:]:
             CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \
                 --model_path "${CKPT}" \
                 --export_path "${EXPORT_DIR}" \
-                --trust_remote_code
+                "${EXPORT_TRUST_REMOTE_CODE[@]}"
             EXPORTED=$((EXPORTED + 1))
         done
         # Also export final model if saved directly to output_dir
@@ -154,7 +174,7 @@ for arg in sys.argv[1:]:
             CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \
                 --model_path "${OUTPUT_DIR}" \
                 --export_path "${OUTPUT_DIR}/exported-checkpoint-final" \
-                --trust_remote_code
+                "${EXPORT_TRUST_REMOTE_CODE[@]}"
             EXPORTED=$((EXPORTED + 1))
         fi
         if [ "$EXPORTED" -eq 0 ]; then

tools/launcher/examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml

@@ -3,12 +3,12 @@
 # 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:
+# 9-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
+#   task_3..5: Generate target completions for each source
+#   task_6:    Merge + deduplicate into one training JSONL
+#   task_7:    Online DFlash training (exports every checkpoint)
+#   task_8:    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
@@ -79,17 +79,33 @@ pipeline:
       ntasks_per_node: 1
       gpus_per_node: 1
 
-  # Step 4: generate target completions. Shards are split across the allocated
+  # Steps 4-6: 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).
+  # Media sources use SGLang's native image/video client; text uses vLLM.
   task_3:
+    script: common/specdec/multimodal_synthetic_generation.sh
+    args:
+      - --dataset pai_understanding
+      - --shard-path <<global_vars.data_root>>/pai_shards
+      - --output-path <<global_vars.data_root>>/pai_outputs
+      - --media-root <<global_vars.data_root>>/pai_understanding
+    environment:
+      - MODEL_PATH: <<global_vars.hf_model>>
+      - 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
+
+  task_4:
     script: common/specdec/multimodal_synthetic_generation.sh
     args:
       - --dataset vqa_v2
@@ -107,10 +123,28 @@ pipeline:
       gpus_per_node: 8
       container: lmsysorg/sglang:v0.5.3-cu129
 
-  # Step 5: merge. Media paths are resolved to absolute here, which is why
+  task_5:
+    script: common/specdec/multimodal_synthetic_generation.sh
+    args:
+      - --dataset specdec_multilingual_prompt
+      - --shard-path <<global_vars.data_root>>/text_shards
+      - --output-path <<global_vars.data_root>>/text_outputs
+    environment:
+      - MODEL_PATH: <<global_vars.hf_model>>
+      - BACKEND: vllm
+      - SGLANG_TP_SIZE: "1"
+      - NUM_TEMPERATURES: "8"
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 8
+      container: vllm/vllm-openai:v0.24.0
+
+  # Step 7: 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:
+  task_6:
     script: common/specdec/merge_dflash_datasets.sh
     args:
       - --source pai_understanding=<<global_vars.data_root>>/pai_outputs
@@ -128,11 +162,11 @@ pipeline:
       ntasks_per_node: 1
       gpus_per_node: 1
 
-  # Step 6: online DFlash training. The VLM_* limits below cap text and visual
+  # Step 8: 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:
+  task_7:
     script: common/specdec/dflash_online_training.sh
     args:
       - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml
@@ -188,8 +222,8 @@ pipeline:
       ntasks_per_node: 1
       gpus_per_node: 8
 
-  # Step 7: smoke test the exported draft under vLLM.
-  task_6:
+  # Step 9: smoke test the exported draft under vLLM.
+  task_8:
     script: common/specdec/vllm_smoke_test.sh
     environment:
       - HF_MODEL_CKPT: <<global_vars.hf_model>>