Skip to content

[OMNIML-4760] synth_support #1696

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

Merged
ChenhanYu merged 5 commits into from
Jun 19, 2026
Merged

[OMNIML-4760] synth_support #1696

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
74a3371
[OMNIML-4760] synth_support: standalone hf_synth.yaml for Qwen3-8B
Jun 18, 2026
6262894
style: fix ruff format in query.py
ChenhanYu Jun 18, 2026
530246c
synth_support: address review feedback — drop /lustre/fsw/portfolios/…
ChenhanYu Jun 18, 2026
b8b3735
synth_support: address coderabbit + claude review feedback
ChenhanYu Jun 18, 2026
cf1ef82
style: ruff-format query.py (collapse line-wrap)
ChenhanYu Jun 19, 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
41 changes: 35 additions & 6 deletions tools/launcher/common/query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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))

@claude claude Bot Jun 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[SUGGESTION] num_proc = min(args.num_proc, len(shard)) evaluates to 0 when a shard is empty (len(shard) == 0), and Dataset.map(num_proc=0) raises ValueError: num_proc must be an integer > 0.

The previous code passed args.num_proc (default 32, always > 0) directly, so this is a newly-introduced crash path. It won't trigger for the large Speculative-Decoding-Multilingual-Prompt-v2 dataset in hf_synth.yaml, but query.py is a general tool: a small/sample dataset (e.g. the existing sample-1K.jsonl) with a high --num-shards produces empty tail shards and would crash here instead of writing an empty output + .done marker.

Why it matters: empty shards are a normal outcome of over-sharding, and crashing on them defeats the idempotent .done design — the array job for that shard fails and (with requeue: true) may retry indefinitely.

Suggested fix — clamp to at least 1, or skip empty shards entirely:

    if len(shard) == 0:
        # nothing to synthesize; still write the .done marker so the shard isn't retried
        shard.to_json(file_path)
        with open(done_path, "w") as done_file:
            done_file.write("done\n")
        continue

    num_proc = max(1, 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:
Expand Down
43 changes: 43 additions & 0 deletions tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
- <<global_vars.hf_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
- <<global_vars.output_dir>>
- --shard-id
- $SLURM_ARRAY_TASK_ID

@claude claude Bot Jun 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[IMPORTANT Compatibility] This line passes the literal string $SLURM_ARRAY_TASK_ID as an arg value, and its correctness depends on shell expansion happening when nemo_run renders run.Script(args=...). If nemo_run quotes args individually (e.g. via shlex.quote), query.py receives the literal "$SLURM_ARRAY_TASK_ID" and crashes on --shard-id's type=int parse — failing every array task, not just one.

Two reasons this is risky:

  1. No precedent — this is the only example in tools/launcher/examples/ that passes a $VAR as an arg value. Every other sharded run relies on common/vllm/query.sh reading SLURM_ARRAY_TASK_ID from the environment and auto-injecting --shard-id-begin ${TASK_ID} --shard-id-step ${TASK_COUNT} (see query.sh:66-78). query.sh deliberately never passes --shard-id as an arg — strongly suggesting arg-level $VAR expansion isn't the established/working mechanism here.
  2. Appears unverified — the PR's verification narration is for hf_offline_eagle3.yaml task_0 (a file not in this diff), not this new hf_synth.yaml. I don't see evidence this YAML was run on-cluster.

Note you do need --shard-id to be set per-task to get query.py's new shard_{id}.jsonl naming + .done idempotency path (query.sh's begin/step injection alone keeps the old train-NNNNN-NNNNN.jsonl naming). So this isn't simply redundant — it's load-bearing.

Suggestion: confirm a dry-run/cluster run actually expands $SLURM_ARRAY_TASK_ID into an integer reaching query.py (e.g. --dryrun -v to inspect the resolved command). If nemo_run does not expand it, query.sh will need to forward --shard-id "$SLURM_ARRAY_TASK_ID" itself (where it's a real shell variable), rather than threading it through the YAML arg list.

- --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
Comment on lines +36 to +43

@claude claude Bot Jun 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[IMPORTANT Compatibility] This task fans work out over 16 shards via --num-shards 16 and --shard-id $SLURM_ARRAY_TASK_ID, but the slurm_config declares no array: field. build_slurm_executor passes array=slurm_config.array straight through (core.py:313), so with the default array=None SLURM submits a single, non-array job — $SLURM_ARRAY_TASK_ID is then unset.

Concretely, that breaks this job two ways:

  1. The literal $SLURM_ARRAY_TASK_ID arg shell-expands to an empty string at runtime, so query.py receives --shard-id "", and argparse's type=int raises ValueError: invalid int value: '' — the job crashes immediately.
  2. Even if it didn't crash, only one shard (index 0) of the 16 would ever run, so 15/16 of the dataset is never synthesized.

The existing eagle3 examples avoid this because they let query.sh read SLURM_ARRAY_TASK_ID from the environment (defaulting TASK_ID=0 when unset) rather than passing it as a CLI arg — but they also aren't trying to fan out across an array here.

Add the array spec so SLURM actually allocates the 16 array tasks:

    slurm_config:
      _factory_: "slurm_factory"
      nodes: 1
      ntasks_per_node: 1
      gpus_per_node: 8
      array: "0-15"
      container: vllm/vllm-openai:latest

3 changes: 3 additions & 0 deletions tools/launcher/slurm_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -91,6 +93,7 @@ def slurm_factory(
container_mounts=container_mounts,
srun_args=srun_args,
array=array,
requeue=requeue,
time=time,
segment=segment,
)
Loading