From 31a90d04532a34225582137d6344dab60cb25c5e Mon Sep 17 00:00:00 2001 From: Slawomir Kierat Date: Fri, 14 Aug 2026 09:37:05 -0700 Subject: [PATCH] fix(specdec): generate all multimodal sources and pad VLM batches Signed-off-by: Slawomir Kierat --- .../utils/plugins/transformers_dataset.py | 1 + .../plugins/test_hf_speculative_offline.py | 1 + .../common/specdec/dflash_online_training.sh | 24 ++++++- .../hf_online_dflash_multimodal.yaml | 64 ++++++++++++++----- 4 files changed, 73 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/utils/plugins/transformers_dataset.py b/modelopt/torch/utils/plugins/transformers_dataset.py index 6a6a7a4e249..9d28d7bd66a 100644 --- a/modelopt/torch/utils/plugins/transformers_dataset.py +++ b/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, ) diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py index 16408e2c3e8..8c8c454d907 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py +++ b/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, ) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index efc38be3b17..10fec646baa 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/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 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 index 6e2e5461fb0..eaf3f70c38a 100644 --- a/tools/launcher/examples/nvidia/Cosmos3-Nano/hf_online_dflash_multimodal.yaml +++ b/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 <>/pai_shards + - --output-path <>/pai_outputs + - --media-root <>/pai_understanding + 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 + + 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 <>/text_shards + - --output-path <>/text_outputs + environment: + - MODEL_PATH: <> + - 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=<>/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: <>