From 74a33711586a3244e17dc0b18423f976e7805ba2 Mon Sep 17 00:00:00 2001 From: Pensieve Intern Date: Thu, 18 Jun 2026 15:43:57 +0000 Subject: [PATCH 1/5] [OMNIML-4760] synth_support: standalone hf_synth.yaml for Qwen3-8B Signed-off-by: Pensieve Intern --- tools/launcher/common/query.py | 32 ++++++++++--- .../examples/Qwen/Qwen3-8B/hf_synth.yaml | 45 +++++++++++++++++++ tools/launcher/slurm_config.py | 3 ++ 3 files changed, 74 insertions(+), 6 deletions(-) create mode 100644 tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml diff --git a/tools/launcher/common/query.py b/tools/launcher/common/query.py index b9ee58b2903..4988b557c06 100644 --- a/tools/launcher/common/query.py +++ b/tools/launcher/common/query.py @@ -94,10 +94,16 @@ def generate(self, messages, verbose=False, **chat_template_kwargs): parser.add_argument("--data-split", type=str, default="train", help="HF dataset split") parser.add_argument("--save", type=str, default=None, help="path to store the generated output.") parser.add_argument("--num-shards", type=int, default=1000, help="number of shards.") +parser.add_argument( + "--shard-id", type=int, default=None, help="single shard id to process." +) parser.add_argument("--shard-id-begin", type=int, default=0, help="the shard id to start.") parser.add_argument( "--shard-id-step", type=int, default=1, help="the step that the shard id progress." ) +parser.add_argument( + "--num-samples", "--num_samples", type=int, default=None, help="maximum samples to process." +) parser.add_argument("--num-proc", type=int, default=32, help="number of processes (concurrency).") parser.add_argument("--temperature", type=float, default=0.0, help="temperature.") parser.add_argument( @@ -207,26 +213,40 @@ def synthesize(data): else: dataset = load_dataset(args.data, split=args.data_split) -if args.num_shards * 100 > len(dataset): +if args.shard_id is None and args.num_shards * 100 > len(dataset): args.num_shards = max(1, min(16, len(dataset) // 100)) if args.save is not None: print(f"Create save dir: {args.save}") os.makedirs(args.save, exist_ok=True) -for shard_id in range(args.shard_id_begin, args.num_shards, args.shard_id_step): - file_path = args.save + f"/train-{shard_id + 1:05}-{args.num_shards:05}.jsonl" +shard_ids = [args.shard_id] if args.shard_id is not None else range( + args.shard_id_begin, args.num_shards, args.shard_id_step +) + +for shard_id in shard_ids: + if args.shard_id is None: + file_path = args.save + f"/train-{shard_id + 1:05}-{args.num_shards:05}.jsonl" + done_path = f"{file_path}.done" + else: + file_path = args.save + f"/shard_{shard_id}.jsonl" + done_path = args.save + f"/shard_{shard_id}.done" - if os.path.exists(file_path): + if os.path.exists(file_path) and os.path.exists(done_path): continue shard = dataset.shard(num_shards=args.num_shards, index=shard_id) + if args.num_samples is not None: + shard = shard.select(range(min(args.num_samples, len(shard)))) print(len(shard), file_path) + num_proc = min(args.num_proc, len(shard)) if shard_id % 2 == 0: - shard = shard.map(disable_thinking_column, num_proc=args.num_proc) - updated_shard = shard.map(synthesize, num_proc=args.num_proc) + shard = shard.map(disable_thinking_column, num_proc=num_proc) + updated_shard = shard.map(synthesize, num_proc=num_proc) updated_shard.to_json(file_path) + with open(done_path, "w") as done_file: + done_file.write("done\n") print(updated_shard[0]) if early_termination: diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml new file mode 100644 index 00000000000..0ef028a434b --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml @@ -0,0 +1,45 @@ +# Standalone vLLM data synthesis for Qwen3-8B. +# +# Usage: +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml --yes + +job_name: qwen3-8b-synth +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + output_dir: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local/modelopt/qwen3-8b-synth-v1 + + task_0: + script: common/vllm/query.sh + args: + - --model + - <> + - --tensor-parallel-size + - "8" + - --trust-remote-code + - --enforce-eager + - --gpu-memory-utilization + - "0.95" + - --max-model-len + - "4096" + - -- + - --data + - nvidia/Speculative-Decoding-Multilingual-Prompt-v2 + - --save + - <> + - --shard-id + - $SLURM_ARRAY_TASK_ID + - --num-shards + - "16" + environment: + - VLLM_STARTUP_TIMEOUT: "1800" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: vllm/vllm-openai:latest + container_mounts: + - /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local:/hf-local + - /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local:/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local + requeue: true diff --git a/tools/launcher/slurm_config.py b/tools/launcher/slurm_config.py index 9c3c853e877..d8cb8ea90dc 100644 --- a/tools/launcher/slurm_config.py +++ b/tools/launcher/slurm_config.py @@ -45,6 +45,7 @@ class SlurmConfig: container_mounts: Optional[list[str]] = None srun_args: Optional[list[str]] = None array: Optional[str] = None + requeue: bool = False nodes: int = 1 ntasks_per_node: int = 1 gpus_per_node: int = 1 @@ -74,6 +75,7 @@ def slurm_factory( ], srun_args: list[str] = ["--no-container-mount-home"], array: Optional[str] = None, + requeue: bool = False, time: str = "04:00:00", segment: Optional[int] = None, ) -> SlurmConfig: @@ -91,6 +93,7 @@ def slurm_factory( container_mounts=container_mounts, srun_args=srun_args, array=array, + requeue=requeue, time=time, segment=segment, ) From 62628945ed0c79a2782b3dccf43200e456933749 Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Thu, 18 Jun 2026 09:37:42 -0700 Subject: [PATCH 2/5] style: fix ruff format in query.py Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Chenhan Yu --- tools/launcher/common/query.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/launcher/common/query.py b/tools/launcher/common/query.py index 4988b557c06..524d3eccdd1 100644 --- a/tools/launcher/common/query.py +++ b/tools/launcher/common/query.py @@ -94,9 +94,7 @@ def generate(self, messages, verbose=False, **chat_template_kwargs): parser.add_argument("--data-split", type=str, default="train", help="HF dataset split") parser.add_argument("--save", type=str, default=None, help="path to store the generated output.") parser.add_argument("--num-shards", type=int, default=1000, help="number of shards.") -parser.add_argument( - "--shard-id", type=int, default=None, help="single shard id to process." -) +parser.add_argument("--shard-id", type=int, default=None, help="single shard id to process.") parser.add_argument("--shard-id-begin", type=int, default=0, help="the shard id to start.") parser.add_argument( "--shard-id-step", type=int, default=1, help="the step that the shard id progress." @@ -220,8 +218,10 @@ def synthesize(data): print(f"Create save dir: {args.save}") os.makedirs(args.save, exist_ok=True) -shard_ids = [args.shard_id] if args.shard_id is not None else range( - args.shard_id_begin, args.num_shards, args.shard_id_step +shard_ids = ( + [args.shard_id] + if args.shard_id is not None + else range(args.shard_id_begin, args.num_shards, args.shard_id_step) ) for shard_id in shard_ids: From 530246c5e2bc8abe7012f5f2ade2b05f32fd4d1a Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Thu, 18 Jun 2026 14:09:43 -0700 Subject: [PATCH 3/5] =?UTF-8?q?synth=5Fsupport:=20address=20review=20feedb?= =?UTF-8?q?ack=20=E2=80=94=20drop=20/lustre/fsw/portfolios/=20leak?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reviewer flagged three lines that hard-coded the NVIDIA-internal team folder path (/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/...) into the public YAML. Fixes: - output_dir: use /scratchspace/modelopt/qwen3-8b-synth-v1 (/scratchspace is a container-mounted write volume — cluster-agnostic). - Drop the explicit container_mounts list. slurm_factory's defaults already provide the right per-cluster mounts; the explicit list in the YAML duplicated that AND leaked the host path. Signed-off-by: Chenhan Yu --- tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml index 0ef028a434b..aff595bf204 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml @@ -7,7 +7,7 @@ job_name: qwen3-8b-synth pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3-8B - output_dir: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local/modelopt/qwen3-8b-synth-v1 + output_dir: /scratchspace/modelopt/qwen3-8b-synth-v1 task_0: script: common/vllm/query.sh @@ -39,7 +39,4 @@ pipeline: ntasks_per_node: 1 gpus_per_node: 8 container: vllm/vllm-openai:latest - container_mounts: - - /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local:/hf-local - - /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local:/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local requeue: true From b8b3735731b50d75b080471644468738d7d71abf Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Thu, 18 Jun 2026 14:22:26 -0700 Subject: [PATCH 4/5] synth_support: address coderabbit + claude review feedback YAML (slurm_config): - Add array: "0-15" so SLURM allocates the 16 array tasks the task expects. Without it, --shard-id $SLURM_ARRAY_TASK_ID expands to empty string and the job crashes immediately on argparse (claude[bot] yaml:42). query.py: - Apply --num-samples globally BEFORE sharding so the cap bounds total output, not per-shard output (coderabbit:query.py:241). - Validate --shard-id at the interface boundary; out-of-range ids now fail loud with a clear message instead of a confusing ValueError from dataset.shard() (coderabbit:query.py:225). Signed-off-by: Chenhan Yu --- tools/launcher/common/query.py | 15 +++++++++++++-- .../launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml | 1 + 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/tools/launcher/common/query.py b/tools/launcher/common/query.py index 524d3eccdd1..58570e7e65c 100644 --- a/tools/launcher/common/query.py +++ b/tools/launcher/common/query.py @@ -214,6 +214,19 @@ def synthesize(data): if args.shard_id is None and args.num_shards * 100 > len(dataset): args.num_shards = max(1, min(16, len(dataset) // 100)) +# Apply --num-samples globally BEFORE sharding so the cap bounds total output, +# not per-shard output (coderabbit:query.py:241). +if args.num_samples is not None: + dataset = dataset.select(range(min(args.num_samples, len(dataset)))) + +# Validate --shard-id once at the interface boundary (coderabbit:query.py:225). +# dataset.shard(index=...) raises a confusing ValueError on out-of-range ids; +# fail loud with a clear message instead. +if args.shard_id is not None and not (0 <= args.shard_id < args.num_shards): + parser.error( + f"--shard-id {args.shard_id} out of range [0, {args.num_shards})" + ) + if args.save is not None: print(f"Create save dir: {args.save}") os.makedirs(args.save, exist_ok=True) @@ -236,8 +249,6 @@ def synthesize(data): continue shard = dataset.shard(num_shards=args.num_shards, index=shard_id) - if args.num_samples is not None: - shard = shard.select(range(min(args.num_samples, len(shard)))) print(len(shard), file_path) num_proc = min(args.num_proc, len(shard)) diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml index aff595bf204..9d19d512684 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml @@ -39,4 +39,5 @@ pipeline: ntasks_per_node: 1 gpus_per_node: 8 container: vllm/vllm-openai:latest + array: "0-15" requeue: true From cf1ef82682a8bb8f1cf959701d1ce213bad571cd Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Thu, 18 Jun 2026 19:54:40 -0700 Subject: [PATCH 5/5] style: ruff-format query.py (collapse line-wrap) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pre-commit ruff-format collapses the 3-line ``parser.error(...)`` call introduced by the 530246c5e cleanup back to a single line — fits within the 100-col limit. No semantic change. Signed-off-by: Chenhan Yu --- tools/launcher/common/query.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tools/launcher/common/query.py b/tools/launcher/common/query.py index 58570e7e65c..27c41953d8a 100644 --- a/tools/launcher/common/query.py +++ b/tools/launcher/common/query.py @@ -223,9 +223,7 @@ def synthesize(data): # dataset.shard(index=...) raises a confusing ValueError on out-of-range ids; # fail loud with a clear message instead. if args.shard_id is not None and not (0 <= args.shard_id < args.num_shards): - parser.error( - f"--shard-id {args.shard_id} out of range [0, {args.num_shards})" - ) + parser.error(f"--shard-id {args.shard_id} out of range [0, {args.num_shards})") if args.save is not None: print(f"Create save dir: {args.save}")