Skip to content

Add Cosmos3 Nano DFlash multimodal training recipe #2053

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
skierat wants to merge 10 commits into
base: main
Choose a base branch
from
Open

Add Cosmos3 Nano DFlash multimodal training recipe #2053

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
1ca95cc
feat(specdec): add Cosmos3 Nano DFlash multimodal training recipe
skierat Aug 3, 2026
9d839c0
Fix Cosmos3 Nano DFlash training notebook workflow
skierat Aug 3, 2026
3196515
refactor(specdec): unify DFlash dataset merge workflow
skierat Aug 3, 2026
ad12fc4
test(specdec): cover VLM DFlash collator paths
skierat Aug 3, 2026
862e611
test(specdec): cover VLM collator edge cases
skierat Aug 3, 2026
228ad4d
test(specdec): cover VLM content normalization edge cases
skierat Aug 4, 2026
5683e84
fix(specdec): harden multimodal generation and remote-code handling
skierat Aug 4, 2026
8e43d47
fix(specdec): harden multimodal media delivery
skierat Aug 4, 2026
0ebed91
fix(specdec): harden dataset preparation and parallel merge recovery
skierat Aug 4, 2026
691039c
fix typo
skierat Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .claude/skills/benchmark-model-kernels
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
../../.agents/skills/benchmark-model-kernels
32 changes: 26 additions & 6 deletions examples/speculative_decoding/distributed_generate/launch.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

set -euo pipefail

if [ $# -lt 9 ]; then
echo "Usage: $0 <job_id> <backend> <model_path> <data_path> <output_path> <scripts_path> <start_shard> <jobs_per_node> <comma_separated_node_names> <system_prompt>"
echo "Example: $0 245387 vllm /model/ /input_data/ /output_data/ /scripts/ 0 20 cluster-01,cluster-02 "\"You are a helpful assistant.\"""
Expand All @@ -32,29 +34,47 @@ NODE_NAME=$9
SYSTEM_PROMPT="${10:-}"
IFS=',' read -r -a NODE_LIST <<< "$NODE_NAME"

# Pyxis requires bind-mount sources to exist before srun creates the container.
mkdir -p "$OUTPUT_PATH"

# backend needs to be either vllm or sglang
if [ "$BACKEND" != "vllm" ] && [ "$BACKEND" != "sglang" ]; then
echo "Invalid backend: $BACKEND"
exit 1
fi

if [ "$BACKEND" == "vllm" ]; then
CONTAINER_IMAGE="vllm/vllm-openai:v0.8.5"
DEFAULT_CONTAINER_IMAGE="vllm/vllm-openai:v0.24.0"
else
CONTAINER_IMAGE="lmsysorg/sglang:v0.4.6.post2-cu124"
DEFAULT_CONTAINER_IMAGE="lmsysorg/sglang:v0.5.3-cu129"
fi
CONTAINER_IMAGE=${CONTAINER_IMAGE:-$DEFAULT_CONTAINER_IMAGE}
counter=$START_SHARD
worker_pids=()
for node in "${NODE_LIST[@]}"; do
echo "Processing node: $node"
srun --output=srun_worker_${node}.log --jobid=$JOB_ID -N 1 --ntasks=1 --ntasks-per-node=1 -w $node \
--mpi pmix --overlap --container-image=$CONTAINER_IMAGE \
--container-mounts=$MODEL_PATH:/model/,$DATA_PATH:/input_data/,$OUTPUT_PATH:/output_data/,$SCRIPTS_PATH:/scripts/ \
bash /scripts/distributed_generate/worker.sh $counter $BACKEND $JOBS_PER_NODE "$SYSTEM_PROMPT" &
srun --output="srun_worker_${node}.log" --jobid="$JOB_ID" -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
--mpi pmix --overlap --container-image="$CONTAINER_IMAGE" \
--container-mounts="$MODEL_PATH":/model/,"$DATA_PATH":/input_data/,"$OUTPUT_PATH":/output_data/,"$SCRIPTS_PATH":/scripts/ \
bash /scripts/distributed_generate/worker.sh "$counter" "$BACKEND" "$JOBS_PER_NODE" "$SYSTEM_PROMPT" &

echo "srun command for node $node started with PID $!" >> srun_launch.log
worker_pids+=("$!")

# increment counter by JOBS_PER_NODE
counter=$((counter + JOBS_PER_NODE))
done

echo "Started workers, each processing $JOBS_PER_NODE shards of data. Will process shards $START_SHARD through $((counter - 1))."

worker_status=0
for worker_pid in "${worker_pids[@]}"; do
if ! wait "$worker_pid"; then
worker_status=1
fi
done

if [ "$worker_status" -ne 0 ]; then
echo "ERROR: one or more workers failed." >&2
fi
exit "$worker_status"
96 changes: 96 additions & 0 deletions examples/speculative_decoding/distributed_generate/launch_multimodal.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
#!/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.

set -euo pipefail

if [ $# -lt 10 ]; then
echo "Usage: $0 <job_id> <backend> <model_path> <shard_path> <output_path> <scripts_path> <start_shard> <jobs_per_node> <media_path> [num_frames] <comma_separated_node_names> [system_prompt]"
echo "Also accepted: $0 <job_id> <backend> <model_path> <shard_path> <output_path> <scripts_path> <start_shard> <jobs_per_node> <comma_separated_node_names> <media_path> [num_frames] [system_prompt]"
echo "Example: $0 245387 sglang /model/ /shards/ /output/ /scripts/ 0 10 /media/ 16 cluster-01"
echo "Optional env: SGLANG_TP_SIZE=8 NUM_TEMPERATURES=8 NUM_THREADS=8 SGLANG_EXTRA_ARGS='--mem-fraction-static 0.75'"
exit 1
fi

JOB_ID=$1
BACKEND=$2
MODEL_PATH=$3
DATA_PATH=$4
OUTPUT_PATH=$5
SCRIPTS_PATH=$6
START_SHARD=$7
JOBS_PER_NODE=$8
ARG9=$9
ARG10=${10:-}
ARG11=${11:-}
ARG12=${12:-}

if [[ "$ARG9" == */* || "$ARG9" == .* ]]; then
MEDIA_PATH=$ARG9
NUM_FRAMES="${ARG10:-32}"
NODE_NAME=$ARG11
SYSTEM_PROMPT="$ARG12"
else
NODE_NAME=$ARG9
MEDIA_PATH=$ARG10
NUM_FRAMES="${ARG11:-32}"
SYSTEM_PROMPT="$ARG12"
fi

if [ -z "${NODE_NAME:-}" ] || [ -z "${MEDIA_PATH:-}" ]; then
echo "ERROR: both media_path and node_name are required." >&2
exit 1
fi

IFS=',' read -r -a NODE_LIST <<< "$NODE_NAME"

if [ "$BACKEND" != "sglang" ]; then
echo "Multimodal generation currently supports backend=sglang."
exit 1
fi

mkdir -p "$OUTPUT_PATH"

# Set CONTAINER_IMAGE to a local .sqsh image to avoid pulling from the registry.
DEFAULT_CONTAINER_IMAGE="lmsysorg/sglang:v0.5.3-cu129"
CONTAINER_IMAGE="${CONTAINER_IMAGE:-$DEFAULT_CONTAINER_IMAGE}"

counter=$START_SHARD
worker_pids=()
for node in "${NODE_LIST[@]}"; do
echo "Processing node: $node"
srun --output=srun_vlm_worker_${node}.log --jobid=$JOB_ID -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Quote the srun arguments.

--output=srun_vlm_worker_${node}.log and --jobid=$JOB_ID are unquoted, so a node name with a glob character or whitespace splits into extra arguments. launch.sh Line 56 already quotes both.

🔧 Proposed fix 
-    srun --output=srun_vlm_worker_${node}.log --jobid=$JOB_ID -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
+    srun --output="srun_vlm_worker_${node}.log" --jobid="$JOB_ID" -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
srun --output=srun_vlm_worker_${node}.log --jobid=$JOB_ID -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
srun --output="srun_vlm_worker_${node}.log" --jobid="$JOB_ID" -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
🧰 Tools
🪛 Shellcheck (0.11.0)

[info] 72-72: Double quote to prevent globbing and word splitting.

(SC2086)

[info] 72-72: Double quote to prevent globbing and word splitting.

(SC2086)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/speculative_decoding/distributed_generate/launch_multimodal.sh` at
line 72, Update the srun invocation in the distributed launch script to quote
both the --output value containing ${node} and the --jobid value containing
$JOB_ID, matching the quoting used by launch.sh while preserving the existing
argument values.

Source: Linters/SAST tools

--mpi pmix --overlap --container-image="$CONTAINER_IMAGE" \
--container-mounts="$MODEL_PATH":/model/,"$DATA_PATH":/input_data/,"$OUTPUT_PATH":/output_data/,"$SCRIPTS_PATH":/scripts/,"$MEDIA_PATH":/media_data/ \
bash /scripts/distributed_generate/worker_multimodal.sh "$counter" "$BACKEND" "$JOBS_PER_NODE" "$NUM_FRAMES" "$SYSTEM_PROMPT" &

echo "srun multimodal command for node $node started with PID $!" >> srun_launch_multimodal.log
worker_pids+=("$!")
counter=$((counter + JOBS_PER_NODE))
done

echo "Started multimodal workers, each processing $JOBS_PER_NODE shards of data. Will process shards $START_SHARD through $((counter - 1))."

worker_status=0
for worker_pid in "${worker_pids[@]}"; do
if ! wait "$worker_pid"; then
worker_status=1
fi
done

if [ "$worker_status" -ne 0 ]; then
echo "ERROR: one or more multimodal workers failed." >&2
fi
exit "$worker_status"
Loading
Loading