Add Cosmos3 Nano DFlash multimodal training recipe - #2053
Conversation
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe changes add multimodal dataset preparation, distributed SGLang generation, JSONL merging and deduplication, a Cosmos3 Nano training workflow, VLM collation controls, truncation reporting, and a shared benchmark skill link. ChangesSpeculative decoding pipeline
Repository skill link
Estimated code review effort: 5 (Critical) | ~120 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant RecipeLauncher
participant Slurm
participant MultimodalWorker
participant VLMGenerator
participant SGLang
RecipeLauncher->>Slurm: submit distributed generation
Slurm->>MultimodalWorker: start shard worker
MultimodalWorker->>SGLang: launch and health-check servers
MultimodalWorker->>VLMGenerator: process assigned JSONL shards
VLMGenerator->>SGLang: send image or video generation requests
SGLang-->>VLMGenerator: return assistant responses
VLMGenerator-->>MultimodalWorker: write JSONL output records
Caution Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional.
❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2053 +/- ##
===========================================
+ Coverage 66.94% 77.76% +10.81%
===========================================
Files 519 521 +2
Lines 59401 61264 +1863
===========================================
+ Hits 39767 47640 +7873
+ Misses 19634 13624 -6010
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 15
🧹 Nitpick comments (13)
examples/speculative_decoding/distributed_generate/worker_multimodal.sh (1)
234-240: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueEmit the
SYSTEM_PROMPTwarning once, not per shard.The check sits inside the shard loop and inside every temperature process. A run with 8 temperatures and 50 shards prints the same warning 400 times. Move the check above the loop, near the argument parsing at Line 23.
🤖 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/worker_multimodal.sh` around lines 234 - 240, Move the SYSTEM_PROMPT non-empty check and warning out of the shard and temperature process loops in the distributed generation script, placing it near the argument-parsing setup before those loops begin. Preserve the existing warning text and ensure it is emitted once per script execution.examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py (2)
159-181: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCompute the prompt string once per record.
prompt(question, args.prompt_style, options)runs twice for every PAI record andprompt(question, args.prompt_style)runs twice for every VQA record. Assign the value to a local variable and reuse it in the message content and the"prompt"field. This also guarantees that both fields stay identical.Also applies to: 220-221
🤖 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/recipes/prepare_multimodal_synthetic_shards.py` around lines 159 - 181, Compute the formatted prompt once per record in the generator before constructing user_content, storing it in a local variable. Reuse that variable for the text content and the yielded "prompt" field, including the corresponding VQA record path, so both fields remain identical and prompt generation runs only once.
39-78: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
__all__and docstrings for the module-level functions.The repository guidelines require an explicit public API and docstrings for public functions.
prompt,load_json,find_json,image_path, andnormalize_optionshave no docstrings, and the module has no__all__. The recipe is a CLI entry point, so a short__all__ = ["main"]plus one-line docstrings is enough.As per coding guidelines: "Define each module's public API with
__all__ = [...]" and "Document public and higher-level APIs with docstrings".🤖 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/recipes/prepare_multimodal_synthetic_shards.py` around lines 39 - 78, Add an explicit module-level __all__ containing the public CLI entry point main, and add concise one-line docstrings to prompt, load_json, find_json, image_path, and normalize_options describing their behavior. Keep the existing implementations and signatures unchanged.Source: Coding guidelines
examples/speculative_decoding/distributed_generate/worker.sh (1)
67-90: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
cleanupdoes not stop the generation processes.
GENERATION_PIDSholds the backgroundrun_temperaturejobs, butcleanupterminates onlySERVER_PIDS. If the script exits on a signal, the generation children keep running against dead servers and keep appending to the output files. KillGENERATION_PIDSfirst, then the servers. The same gap exists inworker_multimodal.shLines 64-87.🔧 Proposed fix
cleanup() { local attempt pid still_running + for pid in "${GENERATION_PIDS[@]}"; do + kill "$pid" 2>/dev/null || true + done for pid in "${SERVER_PIDS[@]}"; do kill "$pid" 2>/dev/null || true done🤖 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/worker.sh` around lines 67 - 90, Update cleanup in worker.sh to terminate and reap all GENERATION_PIDS before processing SERVER_PIDS, applying the same ordering and shutdown behavior to worker_multimodal.sh. Preserve the existing graceful wait, forced kill, and final wait handling for both process groups.examples/speculative_decoding/recipes/merge_dflash_datasets.py (2)
318-326: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
--cache-contextsbounds the number of keys, not the retained word arrays.Each key stores a list of word arrays that grows without a limit. A dataset with many records that share one prompt context keeps every word array for that context in memory, and the
any(...)scan at Line 319 becomes linear in that list for every new record. Consider a per-key cap, for example keeping only the most recent N arrays per context.🤖 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/recipes/merge_dflash_datasets.py` around lines 318 - 326, The candidates cache limits only the number of context keys, allowing each key’s retained word-array list to grow without bound and making the has_required_overlap scan increasingly expensive. Update the candidates insertion logic near has_required_overlap and candidates.setdefault so each key retains at most a fixed recent-array cap, removing the oldest arrays when that per-key limit is exceeded while preserving the existing key-level cache_contexts eviction.
183-206: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueUse
math.ceilfor the overlap thresholds.The current expression cannot round above
math.ceil; it can undercount when the product is just above an integer but within1e-9. Usemath.ceilor exact integer arithmetic.🤖 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/recipes/merge_dflash_datasets.py` around lines 183 - 206, Update has_required_overlap to use math.ceil for both threshold products instead of int(... + 0.999999999), including the minimum calculation and final shared-count comparison; preserve the existing overlap scan and boundary checks.examples/speculative_decoding/distributed_generate/launch_multimodal.sh (1)
35-55: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe positional-argument heuristic silently misreads valid inputs.
The script decides the argument order from
ARG9containing/or starting with.. A relative media path such asmediafails that test, so the script assigns it toNODE_NAMEand assigns the node list toMEDIA_PATH. The 10-argument form... /media/ cluster-01assignscluster-01toNUM_FRAMESand then aborts with a message about missing arguments, which does not describe the real cause.Prefer named flags, for example
--media-pathand--nodes, or accept only one documented order. If the dual order must stay, validate thatMEDIA_PATHis a directory and thatNUM_FRAMESmatches^[0-9]+$, and report the specific mismatch.🤖 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` around lines 35 - 55, Replace the ARG9 path-character heuristic in the launch argument parsing with a single documented positional order or explicit named flags such as --media-path and --nodes. If both orders remain supported, validate the resolved MEDIA_PATH as a directory and NUM_FRAMES as numeric, then emit a specific error describing the invalid argument arrangement before execution.examples/speculative_decoding/recipes/merge_dflash_datasets_parallel.py (1)
149-153: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNote the deduplication boundary in the module docstring or
--help.Records are deduplicated only inside one partition. Two near-duplicate records that carry different numeric output prefixes land on different workers and both survive. The single-worker merger deduplicates globally, so the parallel and serial paths give different results for the same inputs. State this trade-off in the docstring so users choose the mode deliberately.
🤖 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/recipes/merge_dflash_datasets_parallel.py` around lines 149 - 153, Update the module docstring and the argparse help text for the parallel merge workflow to state that deduplication occurs only within each partition, so records with different numeric output prefixes may survive on separate workers; note that the single-worker merger deduplicates globally and can therefore produce different results.modelopt/torch/utils/plugins/transformers_dataset.py (3)
544-598: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate
_truncate_assistant_contentand_truncate_prompt_content.The two methods are structurally identical: iterate messages, filter by role, tokenize text content, truncate to a prefix/suffix window, decode, and print a one-time warning. Only the role filter, the token-limit attribute, the warned-flag attribute, and the warning text differ. Extract a shared private helper parameterized by role predicate, limit, and warned-flag name to avoid maintaining two copies of the same truncation logic.
♻️ Proposed refactor to share the truncation logic
+ def _truncate_content_by_role(self, messages, roles, max_tokens, warned_attr, label): + if max_tokens is None: + return + for message in messages: + if message.get("role") not in roles: + continue + for content in message.get("content", []): + if content.get("type") != "text" or not isinstance(content.get("text"), str): + continue + token_ids = self.tokenizer(content["text"], add_special_tokens=False).input_ids + if len(token_ids) <= max_tokens: + continue + suffix_tokens = min(128, max(1, max_tokens // 4)) + prefix_tokens = max_tokens - suffix_tokens + content["text"] = self.tokenizer.decode( + token_ids[:prefix_tokens] + token_ids[-suffix_tokens:], + skip_special_tokens=True, + ) + if not getattr(self, warned_attr): + print_rank_0(f"Truncating {label} content to {max_tokens} tokens before VLM tokenization.") + setattr(self, warned_attr, True) + + def _truncate_assistant_content(self, messages): + self._truncate_content_by_role( + messages, {"assistant"}, self.max_assistant_tokens, + "_assistant_content_truncation_warned", "assistant" + ) + + def _truncate_prompt_content(self, messages): + self._truncate_content_by_role( + messages, {"system", "user"}, self.max_prompt_tokens, + "_prompt_content_truncation_warned", "user/system" + )🤖 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 `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 544 - 598, Consolidate the duplicated logic in `_truncate_assistant_content` and `_truncate_prompt_content` into one private helper parameterized by the allowed-role predicate, token limit, warned-flag attribute, and warning text. Update both methods to delegate to that helper while preserving their existing role filtering, truncation behavior, decoded content, and one-time warnings.
335-349: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd min/max order validation for image pixel bounds.
VLM_MIN_PIXELS/VLM_MAX_PIXELSare parsed and passed directly toAutoProcessor.from_pretrainedwithout checking thatmin_pixels <= max_pixels._configure_video_processorvalidates this same relationship for the video processor at Line 421-422 ("VLM_VIDEO_MIN_PIXELS must not exceed VLM_VIDEO_MAX_PIXELS."). Apply the same check here so a misconfigured environment produces a clear error instead of an unpredictable failure inside the processor constructor.🛡️ Proposed fix to validate image pixel bounds order
self.processor = transformers.AutoProcessor.from_pretrained(processor, **processor_kwargs) + if ( + "min_pixels" in processor_kwargs + and "max_pixels" in processor_kwargs + and processor_kwargs["min_pixels"] > processor_kwargs["max_pixels"] + ): + raise ValueError("VLM_MIN_PIXELS must not exceed VLM_MAX_PIXELS.") if processor_kwargs: print_rank_0(f"Loaded VLM processor with {processor_kwargs}")🤖 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 `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 335 - 349, After parsing VLM_MIN_PIXELS and VLM_MAX_PIXELS in the processor initialization flow, validate that min_pixels does not exceed max_pixels before calling AutoProcessor.from_pretrained. Raise a clear ValueError matching the existing video-bound validation behavior, while preserving processing when either value is unset or the bounds are valid.
350-375: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate the
VLM_MAX_PROMPT_TOKENS/VLM_MAX_ASSISTANT_TOKENSparsing blocks.These two blocks repeat the same parse-and-validate pattern already abstracted into a loop for the pixel-bound env vars at Line 335-344. Extract a small helper, for example
_parse_positive_int_env(env_name), and reuse it for both attributes to avoid future drift between the two nearly identical error messages and validation paths.♻️ Proposed refactor to share the parsing logic
+ def _parse_positive_int_env(env_name): + value = os.environ.get(env_name) + if value is None: + return None + try: + parsed = int(value) + except ValueError as exc: + raise ValueError(f"{env_name} must be an integer, got {value!r}") from exc + if parsed <= 0: + raise ValueError(f"{env_name} must be a positive integer.") + return parsed + - max_prompt_tokens = os.environ.get("VLM_MAX_PROMPT_TOKENS") - self.max_prompt_tokens = None - if max_prompt_tokens is not None: - try: - self.max_prompt_tokens = int(max_prompt_tokens) - except ValueError as exc: - raise ValueError( - "VLM_MAX_PROMPT_TOKENS must be a positive integer, got " - f"{max_prompt_tokens!r}" - ) from exc - if self.max_prompt_tokens <= 0: - raise ValueError("VLM_MAX_PROMPT_TOKENS must be a positive integer.") + self.max_prompt_tokens = _parse_positive_int_env("VLM_MAX_PROMPT_TOKENS") self._prompt_content_truncation_warned = False - max_assistant_tokens = os.environ.get("VLM_MAX_ASSISTANT_TOKENS") - self.max_assistant_tokens = None - if max_assistant_tokens is not None: - try: - self.max_assistant_tokens = int(max_assistant_tokens) - except ValueError as exc: - raise ValueError( - "VLM_MAX_ASSISTANT_TOKENS must be a positive integer, got " - f"{max_assistant_tokens!r}" - ) from exc - if self.max_assistant_tokens <= 0: - raise ValueError("VLM_MAX_ASSISTANT_TOKENS must be a positive integer.") + self.max_assistant_tokens = _parse_positive_int_env("VLM_MAX_ASSISTANT_TOKENS") self._assistant_content_truncation_warned = False🤖 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 `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 350 - 375, Deduplicate the positive-integer environment parsing in the initializer by extracting a helper such as _parse_positive_int_env that reads, converts, and validates an environment variable while preserving the existing error messages and None behavior. Replace the separate parsing blocks for VLM_MAX_PROMPT_TOKENS and VLM_MAX_ASSISTANT_TOKENS with calls to the helper, while keeping their attribute assignments and truncation-warning initialization unchanged.examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb (2)
586-588: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffReconsider installing dependencies inside the training job.
Each job start runs three
pip installcommands on the compute node. This requires network access from the allocation, adds startup latency to every retry, and can change the resolved dependency set between runs.python3 -m pip install torchcodecis also unpinned, so a new release changes the environment silently.Prefer a prebuilt container that already contains the requirements. If the installs must stay, pin
torchcodecand add--no-inputplus a failure message.🤖 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/recipes/train_dflash_cosmos3_nano.ipynb` around lines 586 - 588, Update the training setup commands in the notebook to avoid installing dependencies at job startup by using a prebuilt container containing the requirements. If the pip installs must remain, pin torchcodec to a fixed version, add --no-input to each install command, and emit a clear failure message when installation fails.
121-132: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate
VQA_NUM_SAMPLESwith the neighbouring checks.The cell validates
PAI_NUM_GENERATION_SHARDS,PAI_LINES_PER_SHARD,DEDUP_WORD_OVERLAP,DEDUP_CACHE_CONTEXTS, andMERGE_WORKERS. It does not validateVQA_NUM_SAMPLES. A zero or negative value reaches--num_samplesand yields an empty or failed VQA shard set later.♻️ Proposed change
if PAI_NUM_GENERATION_SHARDS <= 0 or PAI_LINES_PER_SHARD <= 0: raise ValueError("PAI_NUM_GENERATION_SHARDS and PAI_LINES_PER_SHARD must be positive.") +if VQA_NUM_SAMPLES <= 0: + raise ValueError("VQA_NUM_SAMPLES must be positive.")🤖 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/recipes/train_dflash_cosmos3_nano.ipynb` around lines 121 - 132, Update the configuration validation block near VQA_NUM_SAMPLES to require VQA_NUM_SAMPLES to be positive, alongside the existing checks for PAI_NUM_GENERATION_SHARDS, PAI_LINES_PER_SHARD, DEDUP_CACHE_CONTEXTS, and MERGE_WORKERS, and raise a clear ValueError when it is zero or negative.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@examples/speculative_decoding/distributed_generate/launch_multimodal.sh`:
- Line 66: Update the default value of CONTAINER_IMAGE in launch_multimodal.sh
to use the same public registry image reference as launch.sh instead of the
hardcoded personal cluster path. Add documentation describing how to override
CONTAINER_IMAGE when using a pre-built squashfs image.
- 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.
In
`@examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py`:
- Around line 86-111: Update _resolve_media_path to return None instead of the
original path when no candidate exists, and log the unresolved path once before
returning. Preserve all existing candidate-resolution checks so main’s not
resolved_image branch can fall back to the video path or skip the sample.
In `@examples/speculative_decoding/distributed_generate/worker_multimodal.sh`:
- Around line 218-243: Replace the string-based cmd construction and eval
invocation in the multimodal worker with a Bash argument array, following the
existing worker.sh array pattern. Append optional --overwrite directly to the
array, invoke the command without eval, and retain the existing command
arguments and logging behavior.
- Around line 93-110: Update the media HTTP server setup in the API_MODE=openai
block to serve /media_data instead of the container root, and adjust
MEDIA_URL_BASE and related media path construction consistently so generated
URLs remain valid. Keep /input_data unexposed unless an existing client
explicitly requires it, and preserve the startup health check behavior.
In `@examples/speculative_decoding/distributed_generate/worker.sh`:
- Around line 100-121: Make remote-code execution caller-controlled with an
off-by-default TRUST_REMOTE_CODE setting: in
examples/speculative_decoding/distributed_generate/worker.sh lines 100-121,
remove the hardcoded flag from all four server commands and append a
TRUST_REMOTE_CODE_ARGS array only when opted in; apply the same change to both
sglang.launch_server invocations in
examples/speculative_decoding/distributed_generate/worker_multimodal.sh lines
113-141, and document TRUST_REMOTE_CODE alongside the optional environment
variables in launch_multimodal.sh line 23.
In `@examples/speculative_decoding/recipes/merge_dflash_datasets_parallel.py`:
- Around line 193-195: Update the cleanup logic in the merge workflow’s finally
block to track whether processing completed successfully, rather than checking
args.output.exists(). Initialize a success flag before the work begins, set it
only after the output is successfully produced, and remove work_dir only when
that flag is true and args.keep_work_dir is false; preserve failed-run artifacts
for diagnosis.
In
`@examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py`:
- Around line 97-104: Update the ffmpeg invocation in the frame-extraction
function to sample across the full video duration, using a duration-aware rate
such as num_frames divided by the video duration or an equivalent
thumbnail-based filter, while still limiting output to num_frames. Capture the
subprocess result and stderr instead of discarding failures, and log or surface
the ffmpeg error before returning None when extraction fails.
In
`@examples/speculative_decoding/recipes/run_multimodal_synthetic_generation.sh`:
- Around line 72-74: Define and apply one prepared-shard filename contract
across both sites: in
examples/speculative_decoding/recipes/run_multimodal_synthetic_generation.sh
lines 72-74, confirm the emitted filenames, use the notebook’s matching
train-*.jsonl glob, and ensure PREPARE_SHARDS=0 explicitly skips shard
preparation; in
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb line 361,
change the discovery pattern from *.jsonl to train-*.jsonl so counting matches
the PAI cell and excludes unrelated JSONL files.
In `@examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb`:
- Line 52: Remove user-specific cluster defaults from the notebook: at
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb lines
52-52 and 97-97, require MODEL_PATH and raise a clear error when it is unset; at
line 538-538, derive HF_HOME from DATA_ROOT or the user home directory; at line
548-548, replace --account with a documented ${SLURM_ACCOUNT} placeholder; and
at line 573-573, make --container-image overridable using a public image
reference or documented placeholder.
- Line 596: Disable remote-code execution by default across the recipe: at
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb:596-596,
use the caller-controlled TRUST_REMOTE_CODE value with a false default and
export that default before sbatch; at :661-661, conditionally add the
--trust_remote_code argument based on TRUST_REMOTE_CODE defaulting to 0; at
:682-682, remove --trust-remote-code from the documented vllm serve command and
instruct users to add it only for trusted checkpoints.
- Line 234: Add explicit Step 2 configuration guards in both CPU-only cells: at
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb lines
234-234, guard PAI_SHUFFLE_SEED alongside the existing DATA_ROOT and PYTHON_BIN
checks; at lines 336-336, add the equivalent VQA_SHUFFLE_SEED guard. Use the
same actionable missing-configuration message so set -u failures identify the
required setup cell.
- Around line 300-328: Validate each downloaded archive after the three curl
calls and before extract_zip, using known SHA256 values for the expected VQA
questions, VQA annotations, and COCO train artifacts; fail clearly on
mismatches, including caller-overridden URLs rather than extracting unchecked
content. Reuse the archive paths and existing shell flow, and ensure extraction
only proceeds after all integrity checks pass.
- Around line 259-266: Validate PAI_START_SHARD before the arithmetic range
check in the shard setup block, requiring a non-negative integer and emitting an
actionable error message before exiting on invalid input. Match the existing
validation behavior used for the VQA counterpart, while preserving the
subsequent PAI shard range and distribution checks.
- Around line 219-225: Guard both PAI_SHUFFLE_SEED and VQA_SHUFFLE_SEED before
they are referenced in the notebook’s shell setup, using explicit defaults or
required-value checks compatible with set -u. Ensure missing variables produce
intentional behavior or a clear validation error instead of an unbound-variable
failure.
---
Nitpick comments:
In `@examples/speculative_decoding/distributed_generate/launch_multimodal.sh`:
- Around line 35-55: Replace the ARG9 path-character heuristic in the launch
argument parsing with a single documented positional order or explicit named
flags such as --media-path and --nodes. If both orders remain supported,
validate the resolved MEDIA_PATH as a directory and NUM_FRAMES as numeric, then
emit a specific error describing the invalid argument arrangement before
execution.
In `@examples/speculative_decoding/distributed_generate/worker_multimodal.sh`:
- Around line 234-240: Move the SYSTEM_PROMPT non-empty check and warning out of
the shard and temperature process loops in the distributed generation script,
placing it near the argument-parsing setup before those loops begin. Preserve
the existing warning text and ensure it is emitted once per script execution.
In `@examples/speculative_decoding/distributed_generate/worker.sh`:
- Around line 67-90: Update cleanup in worker.sh to terminate and reap all
GENERATION_PIDS before processing SERVER_PIDS, applying the same ordering and
shutdown behavior to worker_multimodal.sh. Preserve the existing graceful wait,
forced kill, and final wait handling for both process groups.
In `@examples/speculative_decoding/recipes/merge_dflash_datasets_parallel.py`:
- Around line 149-153: Update the module docstring and the argparse help text
for the parallel merge workflow to state that deduplication occurs only within
each partition, so records with different numeric output prefixes may survive on
separate workers; note that the single-worker merger deduplicates globally and
can therefore produce different results.
In `@examples/speculative_decoding/recipes/merge_dflash_datasets.py`:
- Around line 318-326: The candidates cache limits only the number of context
keys, allowing each key’s retained word-array list to grow without bound and
making the has_required_overlap scan increasingly expensive. Update the
candidates insertion logic near has_required_overlap and candidates.setdefault
so each key retains at most a fixed recent-array cap, removing the oldest arrays
when that per-key limit is exceeded while preserving the existing key-level
cache_contexts eviction.
- Around line 183-206: Update has_required_overlap to use math.ceil for both
threshold products instead of int(... + 0.999999999), including the minimum
calculation and final shared-count comparison; preserve the existing overlap
scan and boundary checks.
In
`@examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py`:
- Around line 159-181: Compute the formatted prompt once per record in the
generator before constructing user_content, storing it in a local variable.
Reuse that variable for the text content and the yielded "prompt" field,
including the corresponding VQA record path, so both fields remain identical and
prompt generation runs only once.
- Around line 39-78: Add an explicit module-level __all__ containing the public
CLI entry point main, and add concise one-line docstrings to prompt, load_json,
find_json, image_path, and normalize_options describing their behavior. Keep the
existing implementations and signatures unchanged.
In `@examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb`:
- Around line 586-588: Update the training setup commands in the notebook to
avoid installing dependencies at job startup by using a prebuilt container
containing the requirements. If the pip installs must remain, pin torchcodec to
a fixed version, add --no-input to each install command, and emit a clear
failure message when installation fails.
- Around line 121-132: Update the configuration validation block near
VQA_NUM_SAMPLES to require VQA_NUM_SAMPLES to be positive, alongside the
existing checks for PAI_NUM_GENERATION_SHARDS, PAI_LINES_PER_SHARD,
DEDUP_CACHE_CONTEXTS, and MERGE_WORKERS, and raise a clear ValueError when it is
zero or negative.
In `@modelopt/torch/utils/plugins/transformers_dataset.py`:
- Around line 544-598: Consolidate the duplicated logic in
`_truncate_assistant_content` and `_truncate_prompt_content` into one private
helper parameterized by the allowed-role predicate, token limit, warned-flag
attribute, and warning text. Update both methods to delegate to that helper
while preserving their existing role filtering, truncation behavior, decoded
content, and one-time warnings.
- Around line 335-349: After parsing VLM_MIN_PIXELS and VLM_MAX_PIXELS in the
processor initialization flow, validate that min_pixels does not exceed
max_pixels before calling AutoProcessor.from_pretrained. Raise a clear
ValueError matching the existing video-bound validation behavior, while
preserving processing when either value is unset or the bounds are valid.
- Around line 350-375: Deduplicate the positive-integer environment parsing in
the initializer by extracting a helper such as _parse_positive_int_env that
reads, converts, and validates an environment variable while preserving the
existing error messages and None behavior. Replace the separate parsing blocks
for VLM_MAX_PROMPT_TOKENS and VLM_MAX_ASSISTANT_TOKENS with calls to the helper,
while keeping their attribute assignments and truncation-warning initialization
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6af90900-aa04-4ce8-8cec-7bdad53496ab
📒 Files selected for processing (12)
examples/speculative_decoding/distributed_generate/launch.shexamples/speculative_decoding/distributed_generate/launch_multimodal.shexamples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.pyexamples/speculative_decoding/distributed_generate/worker.shexamples/speculative_decoding/distributed_generate/worker_multimodal.shexamples/speculative_decoding/recipes/merge_dflash_datasets.pyexamples/speculative_decoding/recipes/merge_dflash_datasets_parallel.pyexamples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.pyexamples/speculative_decoding/recipes/run_multimodal_synthetic_generation.shexamples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynbexamples/speculative_decoding/scripts/server_generate.pymodelopt/torch/utils/plugins/transformers_dataset.py
| 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" \ |
There was a problem hiding this comment.
🩺 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.
| 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
| "\"$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", |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
huggingface_hub python -m module path for hf CLI download command
💡 Result:
The hf download command provided by the huggingface_hub Python package is implemented within the huggingface_hub module's command-line interface structure [1][2]. Specifically, the logic for the download subcommand is located in the file src/huggingface_hub/commands/download.py [3]. This file contains the DownloadCommand class, which handles the execution of the download process [3]. When you run hf download from your terminal, the command internally utilizes the same Python helper functions provided by the library, primarily hf_hub_download (for individual files) and snapshot_download (for entire repositories) [1][3][4]. The command's entry point is managed within src/huggingface_hub/cli/hf.py [2], which registers the download command as part of the hf CLI application [2]. In summary, while you interact with it via the hf CLI command, the underlying Python source code path for this functionality within the package is huggingface_hub.commands.download [3].
Citations:
- 1: https://huggingface.co/docs/huggingface_hub/guides/cli
- 2: https://github.com/huggingface/huggingface_hub/blob/0b55fb46/src/huggingface_hub/cli/hf.py
- 3: https://github.com/huggingface/huggingface_hub/blob/9e46a06f/src/huggingface_hub/commands/download.py
- 4: https://huggingface.co/docs/huggingface_hub/main/guides/download
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- relevant notebook cells ---'
python3 - <<'PY'
import json
from pathlib import Path
p = Path("examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb")
nb = json.loads(p.read_text())
for i, cell in enumerate(nb["cells"]):
text = "".join(cell.get("source", []))
if any(x in text for x in ("PAI_SHUFFLE_SEED", "huggingface_hub.cli.hf", "PAI_REVISION_ARGS", "VQA_SHUFFLE_SEED")):
print(f"\nCELL {i} ({cell.get('cell_type')}):")
for n, line in enumerate(text.splitlines(), 1):
print(f"{n:4}: {line}")
PY
printf '%s\n' '--- repository references and dependency pins ---'
rg -n --hidden -g '!*.ipynb' 'huggingface.?hub|PAI_SHUFFLE_SEED|VQA_SHUFFLE_SEED' .Repository: NVIDIA/Model-Optimizer
Length of output: 2035
🌐 Web query:
site:github.com/huggingface/huggingface_hub "python -m huggingface_hub.cli.hf" "__main__"
💡 Result:
The command python -m huggingface_hub.cli.hf is an entry point to the Hugging Face Hub Command Line Interface (CLI) [1][2]. The huggingface_hub package provides a built-in CLI tool typically invoked simply as hf [3][4]. When you run python -m huggingface_hub.cli.hf, you are executing the main.py equivalent logic (or the main entry point function) defined within the huggingface_hub.cli.hf module [1][2]. This allows you to access the same functionality as the hf command, such as managing authentication, downloading and uploading files, and managing repositories, without relying on the installed hf wrapper script [3][4][5]. The command serves as a direct way to trigger the CLI application when the hf command might not be in your system's PATH or when you want to ensure you are executing the CLI from a specific Python environment [3][4]. The hf CLI itself is built using libraries like Typer (in newer versions) or standard argparse (in older versions) to manage its command structure [1][6][2]. Usage example: python -m huggingface_hub.cli.hf --help This will display the help documentation, listing available subcommands like auth, download, upload, repo, and others, consistent with the standard hf command usage [3][5][2].
Citations:
- 1: https://github.com/huggingface/huggingface_hub/blob/0b55fb46/src/huggingface_hub/cli/hf.py
- 2: https://github.com/huggingface/huggingface_hub/blob/9e46a06f/src/huggingface_hub/cli/hf.py
- 3: https://github.com/huggingface/huggingface_hub/blob/main/docs/source/en/guides/cli.md
- 4: https://github.com/huggingface/huggingface_hub
- 5: https://github.com/huggingface/huggingface_hub/blob/9e46a06f/docs/source/en/guides/cli.md
- 6: https://github.com/huggingface/huggingface_hub/blob/main/docs/source/en/package_reference/cli.md
🌐 Web query:
site:huggingface.co/docs/huggingface_hub "hf download" "--repo-type" "--local-dir" "--max-workers"
💡 Result:
The hf download command is part of the Hugging Face Hub CLI used to download files from the Hugging Face Hub [1]. Below is a breakdown of the arguments you specified: --repo-type [model|dataset|space] Specifies the type of repository to download from [2][3]. The default value is model [2][3]. --local-dir TEXT If this option is set, the downloaded files will be placed into the specified directory rather than the default cache directory [2][3]. --max-workers INTEGER Determines the maximum number of concurrent workers used for downloading files [2][3]. The default value is 8 [2][3]. In addition to these options, hf download supports various others, such as --revision (to specify a branch, tag, or commit hash), --include and --exclude (for glob pattern filtering), and --token (for authentication) [2][3]. Note that hf_transfer is deprecated and has been removed from the library; it is recommended to use the default high-performance download mechanisms provided by the Hub [4][5]. You can view the full list of options at any time by running hf download --help in your terminal [1].
Citations:
- 1: https://huggingface.co/docs/huggingface_hub/guides/cli
- 2: https://huggingface.co/docs/huggingface_hub/v1.2.2/package_reference/cli
- 3: https://huggingface.co/docs/huggingface_hub/v1.3.3/package_reference/cli
- 4: https://huggingface.co/docs/huggingface_hub/concepts/migration
- 5: https://huggingface.co/docs/huggingface_hub/guides/download
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer train_dflash_cosmos3_nano.ipynb PAI_SHUFFLE_SEED VQA_SHUFFLE_SEED
💡 Result:
The identifiers PAI_SHUFFLE_SEED and VQA_SHUFFLE_SEED are environment variables or configuration parameters associated with data shuffling in training workflows within the NVIDIA Model Optimizer repository, specifically for speculative decoding or VLM (Vision Language Model) training pipelines [1][2]. These parameters control the random number generator (RNG) seeds used to shuffle datasets when performing data loading or streaming in training scripts, such as those found in the examples/speculative_decoding/ directory (e.g., train_dflash_cosmos3_nano.ipynb or associated training scripts) [3][4][5][1]. While specific documentation for a single file named train_dflash_cosmos3_nano.ipynb may vary depending on the repository version, the Model Optimizer library typically utilizes seed parameters in dataset utilities to ensure reproducibility during training [1][2]. In the context of NVIDIA's VLM and speculative decoding workflows: - PAI and VQA often refer to specific dataset subsets or task categories within these pipelines (e.g., VLM or Assistant-related datasets) [3][6]. - Setting these seed variables allows users to fix the order of shuffled samples, which is critical for consistent training results [1][2]. - You can typically set these in your environment or configuration files to ensure that data shuffling is deterministic across different training runs [1][2]. For the most accurate implementation, check the specific cell definitions within the train_dflash_cosmos3_nano.ipynb notebook or the launch_train.sh scripts referenced in the repository, as they define how these environment variables are passed to the training processes [3][4][7].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/modelopt/torch/utils/vlm_dataset_utils.py
- 2: https://github.com/NVIDIA/Model-Optimizer/blob/7f5fd650/modelopt/torch/utils/vlm_dataset_utils.py
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/speculative_decoding/recipes/train_eagle_head_cosmos_reason2.ipynb
- 4: https://github.com/NVIDIA/Model-Optimizer/blob/7f5fd650/examples/speculative_decoding/recipes/train_eagle_head_cosmos_reason2.ipynb
- 5: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/speculative_decoding/doc/dflash.md
- 6: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/ABLATIONS.md
- 7: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/examples/speculative_decoding/README.md
Guard both shuffle-seed variables before use. With set -u, missing PAI_SHUFFLE_SEED or VQA_SHUFFLE_SEED aborts the cell with an unhelpful unbound-variable error. Add explicit defaults or :? checks.
🤖 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/recipes/train_dflash_cosmos3_nano.ipynb` around
lines 219 - 225, Guard both PAI_SHUFFLE_SEED and VQA_SHUFFLE_SEED before they
are referenced in the notebook’s shell setup, using explicit defaults or
required-value checks compatible with set -u. Ensure missing variables produce
intentional behavior or a clear validation error instead of an unbound-variable
failure.
| " --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", |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Two shuffle-seed variables are used without the guard applied to their neighbours. The shared root cause is an incomplete set of :? preconditions in the CPU-only cells. Each cell guards DATA_ROOT and PYTHON_BIN, but reads the seed directly. Under set -u, a reader who skips the Step 2 configuration cell gets a bare "unbound variable" message instead of the actionable instruction.
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb#L234-L234: add: "${PAI_SHUFFLE_SEED:?Run the Step 2 configuration cell before this cell.}"near the other guards in the PAI cell.examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb#L336-L336: add the equivalent guard forVQA_SHUFFLE_SEEDin the VQA cell.
📍 Affects 1 file
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb#L234-L234(this comment)examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb#L336-L336
🤖 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/recipes/train_dflash_cosmos3_nano.ipynb` at
line 234, Add explicit Step 2 configuration guards in both CPU-only cells: at
examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb lines
234-234, guard PAI_SHUFFLE_SEED alongside the existing DATA_ROOT and PYTHON_BIN
checks; at lines 336-336, add the equivalent VQA_SHUFFLE_SEED guard. Use the
same actionable missing-configuration message so set -u failures identify the
required setup cell.
| "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", |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Validate PAI_START_SHARD as a non-negative integer.
Line 264 evaluates PAI_START_SHARD inside (( ... )). If the caller exports a non-numeric value, the arithmetic evaluation fails under set -e with no actionable message. The VQA compute cell already validates its counterpart at line 366. Apply the same check here.
🛠️ Proposed fix
PAI_START_SHARD=${PAI_START_SHARD:-0}
+[[ "$PAI_START_SHARD" =~ ^[0-9]+$ ]] || { echo "PAI_START_SHARD must be a non-negative integer: $PAI_START_SHARD" >&2; exit 1; }📝 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.
| "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", | |
| "PAI_START_SHARD=${PAI_START_SHARD:-0}\n", | |
| "[[ \"$PAI_START_SHARD\" =~ ^[0-9]+$ ]] || { echo \"PAI_START_SHARD must be a non-negative integer: $PAI_START_SHARD\" >&2; exit 1; }\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", |
🤖 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/recipes/train_dflash_cosmos3_nano.ipynb` around
lines 259 - 266, Validate PAI_START_SHARD before the arithmetic range check in
the shard setup block, requiring a non-negative integer and emitting an
actionable error message before exiting on invalid input. Match the existing
validation behavior used for the VQA counterpart, while preserving the
subsequent PAI shard range and distribution checks.
| "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", |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Consider verifying the downloaded archives before extraction.
The cell downloads three third-party archives and extracts them without integrity verification. VQA_QUESTIONS_URL, VQA_ANNOTATIONS_URL, and COCO_TRAIN_URL are also caller-overridable, so a substituted URL is extracted into $VQA_ROOT unchecked. SECURITY.md asks that downloaded artifacts be validated. Add a known SHA256 check between the curl calls and extract_zip.
Zip extraction of an untrusted archive can also write outside the destination through crafted entry names. python -m zipfile -e sanitizes paths, and unzip refuses absolute paths, so the current risk is limited to content substitution.
As per path instructions: "validate external paths, JSON/media inputs, subprocess arguments, resource limits, and downloaded artifacts".
🤖 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/recipes/train_dflash_cosmos3_nano.ipynb` around
lines 300 - 328, Validate each downloaded archive after the three curl calls and
before extract_zip, using known SHA256 values for the expected VQA questions,
VQA annotations, and COCO train artifacts; fail clearly on mismatches, including
caller-overridden URLs rather than extracting unchecked content. Reuse the
archive paths and existing shell flow, and ensure extraction only proceeds after
all integrity checks pass.
Source: Path instructions
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/speculative_decoding/recipes/merge_dflash_datasets.py (1)
141-189: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winEnforce
media_rootcontainment inresolve_media_path. The function accepts../paths and absolute paths outsidemedia_rootbecause it only callsresolve()andis_file(). Reject resolved paths outside the configured root.🤖 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/recipes/merge_dflash_datasets.py` around lines 141 - 189, Update resolve_media_path, used by normalize_messages, to resolve the configured media_root and candidate media path, then reject any candidate that is not contained within media_root before checking or returning the file. Preserve valid files inside the root, including nested paths, while rejecting traversal and absolute paths outside it.Source: Path instructions
🤖 Prompt for all review comments with 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.
Inline comments:
In `@examples/speculative_decoding/recipes/merge_dflash_datasets.py`:
- Around line 496-544: Update merge_in_parallel to track completion with an
explicit success flag initialized before the try block and set only after
concatenate_parts succeeds. In the finally block, remove work_dir only when that
flag is true and args.keep_work_dir is false; do not infer success from
args.output.exists(), so failed runs preserve manifests and partial outputs.
In
`@examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py`:
- Around line 203-207: Update the video-path handling around relative_video and
media_root to resolve the candidate path, then require
video.is_relative_to(media_root) before calling is_file() or passing it to
make_mosaic(). Treat paths outside media_root as missing media and emit only the
safely resolved, validated path for downstream use.
---
Outside diff comments:
In `@examples/speculative_decoding/recipes/merge_dflash_datasets.py`:
- Around line 141-189: Update resolve_media_path, used by normalize_messages, to
resolve the configured media_root and candidate media path, then reject any
candidate that is not contained within media_root before checking or returning
the file. Preserve valid files inside the root, including nested paths, while
rejecting traversal and absolute paths outside it.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 26a3c726-a678-4841-b92c-168970744ca1
📒 Files selected for processing (7)
.claude/skills/benchmark-model-kernelsexamples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.pyexamples/speculative_decoding/recipes/merge_dflash_datasets.pyexamples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.pyexamples/speculative_decoding/recipes/run_multimodal_synthetic_generation.shexamples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynbmodelopt/torch/utils/plugins/transformers_dataset.py
🚧 Files skipped from review as they are similar to previous changes (4)
- examples/speculative_decoding/recipes/run_multimodal_synthetic_generation.sh
- examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py
- examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb
- modelopt/torch/utils/plugins/transformers_dataset.py
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with 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.
Inline comments:
In `@tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py`:
- Around line 292-329: Add coverage in
test_vlm_collator_normalizes_and_truncates_structured_messages for a non-empty
video content value, and assert the corresponding normalized content part has
type "video" with the expected processed video value, while preserving the
existing image assertion.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d70c49c2-66b7-4871-8403-15c28b90161e
📒 Files selected for processing (1)
tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
| def test_vlm_collator_normalizes_and_truncates_structured_messages(tmp_path): | ||
| """Structured contents become processor-compatible text, image, and video parts.""" | ||
|
|
||
| collator = _bare_vlm_collator() | ||
| collator.local_image_path = str(tmp_path) | ||
| collator.max_prompt_tokens = 3 | ||
| collator.max_assistant_tokens = 3 | ||
| captured = {} | ||
| collator._process_multimodal_sample = lambda batch: captured.setdefault("batch", batch) | ||
|
|
||
| result = collator( | ||
| [ | ||
| { | ||
| "messages": [ | ||
| {"role": "system", "content": {"format": "json"}}, | ||
| { | ||
| "role": "user", | ||
| "content": [ | ||
| "one two three four", | ||
| {"text": "five six seven eight"}, | ||
| {"image": "image.png", "text": "", "video": "", "fps": 0}, | ||
| ], | ||
| }, | ||
| {"role": "assistant", "content": 42}, | ||
| {"role": "assistant", "content": "nine ten eleven twelve"}, | ||
| ] | ||
| } | ||
| ] | ||
| ) | ||
|
|
||
| assert result == captured["batch"] | ||
| messages = captured["batch"][0] | ||
| assert messages[0]["content"] == [{"type": "text", "text": '{"format": "json"}'}] | ||
| assert messages[1]["content"][0]["text"] == "token-0 token-1 token-3" | ||
| assert messages[1]["content"][1]["text"] == "token-0 token-1 token-3" | ||
| assert messages[1]["content"][2] == {"type": "image", "image": str(tmp_path / "image.png")} | ||
| assert messages[2]["content"] == [{"type": "text", "text": "42"}] | ||
| assert messages[3]["content"] == [{"type": "text", "text": "token-0 token-1 token-3"}] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add coverage for a non-empty video content part.
The test docstring claims image and video normalization. The input has video: "", and the assertions only verify an image part. Add a non-empty video entry and assert the normalized {"type": "video", ...} output. This will detect regressions in the new video path.
Based on learnings from the coding guidelines, tests must exercise the behavior they claim to validate.
🤖 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 `@tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py` around
lines 292 - 329, Add coverage in
test_vlm_collator_normalizes_and_truncates_structured_messages for a non-empty
video content value, and assert the corresponding normalized content part has
type "video" with the expected processed video value, while preserving the
existing image assertion.
Sources: Coding guidelines, Path instructions
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
1 participant
What does this PR do?
Type of change: new example
Adds an end-to-end Cosmos3 Nano DFlash training recipe for multimodal speculative decoding.
Usage