[OMNIML-4760] synth_support (#1696)

Draft PR opened by **pensieve-intern** for [OMNIML-4760](https://jirasw.nvidia.com/browse/OMNIML-4760). Stage `synth_support` of Epic `OMNIML-4755`. The agent ran from the SPEC on the ticket description; review every change before marking ready. _Always-draft is enforced — the bot never auto-merges._ --- **Agent's self-narration** (stripped from PR diff; surfaced here for context): `VERIFICATION_COMMENT.txt`: ``` OMNIML-4760 synth_support verification for qwen3-8b found and fixed one task_0 issue in `tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml`. What changed: - Updated `task_0` data input from missing `/hf-local/modelopt/Speculative-Decoding-Prompt-Samples` to existing `/hf-local/modelopt/Speculative-Decoding-Dataset-v1-Qwen3-8B/default-openai.jsonl`. - Left the rest of the 4-task monolithic `hf_offline_eagle3.yaml` unchanged. Verification: - `task_0.script` is `common/tensorrt_llm/query.sh`. - `--model <<global_vars.hf_model>>` resolves to `/hf-local/Qwen/Qwen3-8B`. - `task_0.slurm_config.container` is `nvcr.io/nvidia/tensorrt-llm/release:1.2.0`. - Cluster validation on cw_dfw succeeded for the fixed task_0 data path: experiment `cicd_1781221901`, Slurm job `12739762`, remote directory `/lustre/fsw/portfolios/coreai/users/chenhany/experiments/cicd/cicd_1781221901/Qwen3-8B_EAGLE3_offline_task0_verify_jsonl_0`. - Log evidence included successful SSH tunnel/authentication, TensorRT-LLM 1.2.0 container import, Qwen3-8B server health checks, a successful `/v1/chat/completions` request, and loading `1393367` train examples from the Qwen3-8B JSONL data file. Next: - Runner should open the single-file PR for human review because this was a task_0 config fix, not verification-only. ``` _Pollution-strip removed `VERIFICATION_COMMENT.txt` from this commit (sidecar narration and/or incidental lockfile regeneration are never part of the agent's intended deliverable)._  ## Summary by CodeRabbit * **New Features** * Enhanced TE grouped MoE weight quantization with optional per-GEMM/per-expert quantizers (environment-controlled) and updated calibration flow * Expanded launcher CLI with `--shard-id` and per-run sample limiting via `--num-samples` * Added a Qwen3-8B standalone vLLM synthesis job example (`hf_synth.yaml`) * Added Slurm job requeue support * **Bug Fixes** * Improved sharded dataset synthesis with idempotent `.done` markers and smarter shard sizing/capping * **Tests** * Added coverage validating TEGrouped vs sequential MoE default amax behavior and output divergence/accuracy  --------- Signed-off-by: Pensieve Intern <pensieve-intern@nvidia.com> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Pensieve Intern <pensieve-intern@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-10-02 03:14:52 +08:00 · 2026-06-18 20:35:10 -07:00
parent e0125294f9
commit 93dd08f429
3 changed files with 81 additions and 6 deletions
@@ -94,10 +94,14 @@ parser.add_argument(
 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 @@ if os.path.isfile(args.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)
+)

-    if os.path.exists(file_path):
+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) 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:
@@ -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
+      - --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
@@ -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,
    )