Express the Cosmos3 Nano DFlash recipe as a launcher YAML example - #2173
Open
h-guo18 wants to merge 2 commits into
Open
Express the Cosmos3 Nano DFlash recipe as a launcher YAML example#2173h-guo18 wants to merge 2 commits into
h-guo18 wants to merge 2 commits into
Conversation
Ports the train_dflash_cosmos3_nano.ipynb workflow to tools/launcher as a 7-task pipeline, matching how the other speculative-decoding recipes (Qwen3-8B, Qwen3-30B-A3B, Kimi-K2.6, MiniMax-M3) are already expressed. Behavior is preserved: every prepare/generate/merge/train/export/serve step keeps the notebook's flags and defaults. Training and export now reuse the existing common/specdec/dflash_online_training.sh (which already exports every checkpoint, so the notebook's separate export step is no longer needed) and the vLLM smoke test reuses common/specdec/vllm_smoke_test.sh. Three new wrapper scripts cover the data-synthesis half: multimodal_prepare_shards.sh — download + shard one source multimodal_synthetic_generation.sh — generate target completions merge_dflash_datasets.sh — merge + dedup The generation wrapper derives its shard-to-node assignment from the launcher's own allocation (SLURM_JOB_ID / SLURM_JOB_NODELIST) instead of taking a job id and node list as arguments, which is what the notebook cells did by hand. Two incidental benefits: the YAML is covered by tests/test_examples_resolve.py (the notebook could only be checked with `jq empty`), and the cluster-specific paths baked into the notebook (a personal $HF_HOME, a personal .sqsh image, a specific Slurm account) become per-cluster overrides rather than committed defaults. The notebook's explanatory prose is kept in a new recipes/README_dflash_cosmos3_nano.md — why DFlash trains on the target's own completions, what each data source contributes, what the merge guarantees, and the training constraints (seq_len divisibility, the VLM_* caps, and that data.vlm_processor is what selects the multimodal collator). Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Contributor
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (3)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
The prepare wrapper reimplemented the PAI-Bench-U download that prepare_multimodal_synthetic_shards.py already performs via --download, so use the script's own path instead of a second `hf download` call. Semantics are unchanged. Its snapshot_download takes no revision argument, so the explicit CLI download is kept for the PAI_REVISION case only; both branches land the same dataset in the same directory before sharding. PAI_REPO_ID and FORCE_DOWNLOAD are now forwarded too — PAI_REPO_ID was previously accepted by the launch wrapper but never reached the prepare step. VQA's curl/unzip stays in the wrapper: the Python script has no download path for it, only --vqa_root/--image_root over already-extracted data. Also collapses the duplicated per-argument validation in both wrappers. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Targets
skierat/dflash_recipe(#2053), notmain. Merging this into thatbranch before #2053 lands keeps the recipe as a single reviewable unit.
What this changes
Ports the
train_dflash_cosmos3_nano.ipynbworkflow totools/launcheras a7-task pipeline, and replaces the notebook with a README that keeps its prose.
None of the Python, shell, or tests from #2053 are touched. The diff is
additive except for deleting the notebook:
Why
Every other speculative-decoding recipe in the repo is already a launcher YAML
(Qwen3-8B, Qwen3-30B-A3B, Kimi-K2.6, MiniMax-M3, gpt-oss-20b). This makes the
Cosmos3 Nano recipe consistent with them, and picks up three things the
notebook could not:
common/specdec/dflash_online_training.shruns the accelerate launch andexports every checkpoint, so the notebook's separate export step goes away.
The vLLM smoke test reuses
common/specdec/vllm_smoke_test.sh, which alsoasserts a minimum acceptance length rather than only printing a command.
tests/test_examples_resolve.pyvalidates the YAMLstructurally (60 passed locally, including this example). The notebook could
only be checked with
jq empty.$HF_HOME, a personal.sqshimage, and a specific Slurm account. Thosebecome per-cluster overrides. The YAML's defaults are deliberately generic.
Behavior parity
Every flag and default is carried over verbatim — the
training.*/dflash.*overrides, the six
VLM_*caps, the dedup parameters, the per-source samplecounts and shuffle seeds.
The one real rewrite is shard-to-node assignment. The notebook cells computed a
node list by hand and passed
job_id+ comma-separated node names intolaunch_multimodal.sh. The launcher owns the allocation, somultimodal_synthetic_generation.shderives both fromSLURM_JOB_ID/SLURM_JOB_NODELISTinstead. I verified the rounding covers every shard withno gaps and that the out-of-range guards fire.
Prose
The notebook's explanatory content is preserved in
recipes/README_dflash_cosmos3_nano.md: why DFlash trains on the target's owncompletions rather than human-written answers, what each data source
contributes, what the merge guarantees (absolute media paths →
vlm_img_dir=/;dedup removes redundancy, not temperature diversity), and the training
constraints —
training_seq_len % dflash_block_size == 0, why theVLM_*capsexist, and that setting
data.vlm_processoris what selects the multimodalcollator. Per-step detail lives in YAML comments next to the settings.
Not addressed here
Deliberately scoped to the example layout. The review comments on #2053 itself
(the removed
truncation=Trueturning an over-length sample into a job-fatalValueError, the shape-based tensor padding, andhash()making parallel mergenon-reproducible across processes) are unchanged and still worth a look.
🤖 Generated with Claude Code