diff --git a/tools/launcher/common/query.py b/tools/launcher/common/query.py index b9ee58b2903..27c41953d8a 100644 --- a/tools/launcher/common/query.py +++ b/tools/launcher/common/query.py @@ -94,10 +94,14 @@ 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 +211,51 @@ 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)) +# 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) -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) 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..9d19d512684 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml @@ -0,0 +1,43 @@ +# 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: /scratchspace/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 + array: "0-15" + 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, )