初始化项目，由ModelHub XC社区提供模型

Model: ayh015/myLightningOPD Source: Original Platform
2026-08-27 23:50:14 +08:00
commit d4e0a1af66
368 changed files with 559583 additions and 0 deletions
--- a/data_curation/merge.py
+++ b/data_curation/merge.py
@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Merge Arrow IPC files produced by data_curation/pipeline.py into a single parquet.
+
+After multi-GPU data generation, each worker writes Arrow files into
+rank-specific subdirectories. This script merges them into one parquet
+file for downstream consumption (SFT training or Lightning OPD preparation).
+
+Usage:
+    python data_curation/merge.py \
+        --input-dir data/sft_data \
+        --output data/sft_data/merged.parquet
+
+    # With filtering: only keep samples with token count <= 16384
+    python data_curation/merge.py \
+        --input-dir data/sft_data \
+        --output data/sft_data/merged.parquet \
+        --max-tokens 16384
+"""
+
+import argparse
+from pathlib import Path
+
+import json
+import pyarrow as pa
+import pyarrow.ipc as ipc
+from tqdm import tqdm
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Merge Arrow IPC files into a single parquet file."
+    )
+    parser.add_argument(
+        "--input-dir", type=str, required=True,
+        help="Directory containing Arrow files (searched recursively).",
+    )
+    parser.add_argument(
+        "--output", type=str, required=True,
+        help="Output parquet file path.",
+    )
+    parser.add_argument(
+        "--max-tokens", type=int, default=None,
+        help="If set, discard rows with tokens > this value.",
+    )
+    return parser.parse_args()
+
+
+# def merge_arrow_files(input_dir: str, output: str, max_tokens: int | None = None):
+#     input_path = Path(input_dir)
+#     arrow_files = sorted(input_path.rglob("*.arrow"))
+
+#     if not arrow_files:
+#         print(f"No Arrow files found in {input_dir}")
+#         return
+
+#     print(f"Found {len(arrow_files)} Arrow files in {input_dir}")
+
+#     tables = []
+#     total_rows = 0
+#     for f in tqdm(arrow_files, desc="Reading Arrow files"):
+#         with pa.OSFile(str(f), "rb") as source:
+#             table = ipc.open_file(source).read_all()
+#             tables.append(table)
+#             total_rows += len(table)
+
+#     merged = pa.concat_tables(tables)
+#     print(f"Total rows before filtering: {total_rows}")
+
+#     if max_tokens is not None and "tokens" in merged.column_names:
+#         tokens = merged.column("tokens").to_pylist()
+#         mask = [t <= max_tokens for t in tokens]
+#         merged = merged.filter(mask)
+#         filtered = total_rows - len(merged)
+#         print(f"Filtered {filtered} rows with tokens > {max_tokens}")
+
+#     output_path = Path(output)
+#     output_path.parent.mkdir(parents=True, exist_ok=True)
+#     df = merged.to_pandas()
+#     df.to_parquet(output, index=False)
+
+#     print(f"Merged {len(df)} rows -> {output}")
+
+def merge_arrow_files(input_dir: str, output: str, max_tokens: int | None = None):
+    input_path = Path(input_dir)
+    arrow_files = sorted(input_path.rglob("*.arrow"))
+
+    if not arrow_files:
+        print(f"No Arrow files found in {input_dir}")
+        return
+
+    print(f"Found {len(arrow_files)} Arrow files in {input_dir}")
+
+    output_path = Path(output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Recommended path for nested conversation data
+    if output_path.suffix == ".jsonl":
+        total_rows = 0
+        kept_rows = 0
+        filtered_rows = 0
+
+        with open(output_path, "w", encoding="utf-8") as out_f:
+            for f in tqdm(arrow_files, desc="Reading Arrow files"):
+                with pa.OSFile(str(f), "rb") as source:
+                    table = ipc.open_file(source).read_all()
+
+                rows = table.to_pylist()
+                total_rows += len(rows)
+
+                for row in rows:
+                    if max_tokens is not None and "tokens" in row:
+                        if row["tokens"] > max_tokens:
+                            filtered_rows += 1
+                            continue
+
+                    out_f.write(json.dumps(row, ensure_ascii=False) + "\n")
+                    kept_rows += 1
+
+        print(f"Total rows before filtering: {total_rows}")
+        if max_tokens is not None:
+            print(f"Filtered {filtered_rows} rows with tokens > {max_tokens}")
+        print(f"Merged {kept_rows} rows -> {output}")
+        return
+
+    # Optional parquet path, not recommended for nested messages
+    elif output_path.suffix == ".parquet":
+        import pyarrow.parquet as pq
+
+        tables = []
+        total_rows = 0
+
+        for f in tqdm(arrow_files, desc="Reading Arrow files"):
+            with pa.OSFile(str(f), "rb") as source:
+                table = ipc.open_file(source).read_all()
+                tables.append(table)
+                total_rows += len(table)
+
+        merged = pa.concat_tables(tables)
+        print(f"Total rows before filtering: {total_rows}")
+
+        if max_tokens is not None and "tokens" in merged.column_names:
+            tokens = merged.column("tokens").to_pylist()
+            mask = [t <= max_tokens for t in tokens]
+            merged = merged.filter(mask)
+            filtered = total_rows - len(merged)
+            print(f"Filtered {filtered} rows with tokens > {max_tokens}")
+
+        pq.write_table(merged, output)
+        print(f"Merged {len(merged)} rows -> {output}")
+        return
+
+    else:
+        raise ValueError(f"Unsupported output format: {output_path.suffix}")
+
+if __name__ == "__main__":
+    args = parse_args()
+    merge_arrow_files(args.input_dir, args.output, args.max_tokens)
--- a/data_curation/pipeline.py
+++ b/data_curation/pipeline.py
@@ -0,0 +1,218 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Data curation pipeline: generate responses from a dataset using vLLM.
+
+Each worker (identified by --rank) processes a disjoint shard of the input
+dataset, generates responses via vLLM offline inference, and writes results
+as Arrow IPC files (one per batch) into a rank-specific output directory.
+Checkpointing allows resuming from the last completed batch.
+
+Standalone:
+    python data_curation/pipeline.py \
+        --model Qwen/Qwen3-4B \
+        --input data.jsonl \
+        --output-dir output/
+
+Multi-GPU (one model per GPU):
+    See run_curation.sh for the recommended launch pattern.
+"""
+
+import argparse
+import json
+import os
+import pickle
+from pathlib import Path
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.ipc as ipc
+from tqdm import tqdm
+from vllm import LLM, SamplingParams
+
+
+# ---------------------------------------------------------------------------
+# Data I/O
+# ---------------------------------------------------------------------------
+
+def load_dataset(path: str) -> list[dict]:
+    """Load a .jsonl or .parquet dataset into a list of dicts."""
+    if path.endswith(".parquet"):
+        df = pd.read_parquet(path)
+        records = df.to_dict("records")
+        for record in records:
+            if "prompt" in record and hasattr(record["prompt"], "tolist"):
+                record["prompt"] = record["prompt"].tolist()
+        return records
+    elif path.endswith(".jsonl"):
+        with open(path) as f:
+            return [json.loads(line) for line in f]
+    else:
+        raise ValueError(f"Unsupported format: {path}. Use .jsonl or .parquet.")
+
+
+def save_batch_arrow(rows: list[dict], path: str) -> None:
+    """Write a list of dicts as an Arrow IPC file."""
+    table = pa.Table.from_pandas(pd.DataFrame(rows))
+    with pa.OSFile(path, "wb") as sink:
+        with ipc.new_file(sink, table.schema) as writer:
+            writer.write_table(table)
+
+
+# ---------------------------------------------------------------------------
+# Core pipeline
+# ---------------------------------------------------------------------------
+
+def run_curation(args: argparse.Namespace) -> None:
+    tag = f"[Rank {args.rank}/{args.world_size}]"
+
+    # ── Load & shard dataset ──────────────────────────────────────────────
+    print(f"{tag} Loading dataset: {args.input}")
+    dataset = load_dataset(args.input)
+
+    if args.num_samples is not None:
+        dataset = dataset[: args.num_samples]
+        print(f"{tag} Debug mode: limiting to {args.num_samples} samples")
+
+    if args.world_size > 1:
+        dataset = dataset[args.rank :: args.world_size]
+    print(f"{tag} Assigned {len(dataset)} samples")
+
+    # ── Output directory ──────────────────────────────────────────────────
+    if args.world_size > 1:
+        output_dir = Path(args.output_dir) / f"rank{args.rank:05d}"
+    else:
+        output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # ── Checkpoint ────────────────────────────────────────────────────────
+    ckpt_dir = Path(args.checkpoint_dir)
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    ckpt_file = ckpt_dir / f"rank{args.rank:05d}.pkl"
+
+    start_idx = 0
+    if ckpt_file.exists():
+        with open(ckpt_file, "rb") as f:
+            start_idx = pickle.load(f)["next_idx"]
+        print(f"{tag} Resuming from index {start_idx}")
+
+    # ── Model ─────────────────────────────────────────────────────────────
+    print(f"{tag} Loading model: {args.model} (tp={args.tensor_parallel_size})")
+    llm = LLM(
+        model=args.model,
+        tensor_parallel_size=args.tensor_parallel_size,
+        trust_remote_code=True,
+    )
+
+    sampling_params = SamplingParams(
+        temperature=args.temperature,
+        top_p=args.top_p,
+        max_tokens=args.max_tokens,
+        n=args.num_responses,
+    )
+
+    # ── Batch loop ────────────────────────────────────────────────────────
+    total_batches = (len(dataset) + args.batch_size - 1) // args.batch_size
+    total_saved = 0
+
+    print(f"{tag} Processing {len(dataset)} prompts, batch_size={args.batch_size}, "
+          f"total_batches={total_batches}")
+
+    for batch_start in range(start_idx, len(dataset), args.batch_size):
+        batch_end = min(batch_start + args.batch_size, len(dataset))
+        batch = dataset[batch_start:batch_end]
+        batch_idx = batch_start // args.batch_size
+
+        prompts = [item["prompt"] for item in batch]
+        print(f"{tag} Batch {batch_idx + 1}/{total_batches} "
+              f"({batch_end - batch_start} samples) ...")
+
+        outputs = llm.chat(prompts, sampling_params)
+
+        # Build results
+        rows = []
+        for item, output in zip(batch, outputs):
+            for completion in output.outputs:
+                text = completion.text
+                # Ensure <think> tag is present
+                if "</think>" in text and not text.strip().startswith("<think>"):
+                    text = "<think>\n" + text
+                messages = item["prompt"] + [{"role": "assistant", "content": text}]
+                rows.append({
+                    "messages": messages,
+                    "tokens": len(completion.token_ids),
+                })
+
+        # Save Arrow file
+        arrow_path = output_dir / f"data-{batch_idx:05d}-of-{total_batches:05d}.arrow"
+        save_batch_arrow(rows, str(arrow_path))
+        total_saved += len(rows)
+
+        # Save checkpoint
+        with open(ckpt_file, "wb") as f:
+            pickle.dump({"next_idx": batch_end}, f)
+
+        print(f"{tag} Saved {arrow_path.name}  (total: {total_saved})")
+
+    # ── Cleanup ───────────────────────────────────────────────────────────
+    if ckpt_file.exists():
+        ckpt_file.unlink()
+    print(f"{tag} Done! {total_saved} samples → {output_dir}/")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Generate responses from a dataset using vLLM offline inference.",
+    )
+    # Required
+    p.add_argument("--model", type=str, required=True,
+                   help="HuggingFace model name or path.")
+    p.add_argument("--input", type=str, required=True,
+                   help="Input dataset (.jsonl or .parquet).")
+    p.add_argument("--output-dir", type=str, required=True,
+                   help="Root output directory. Each rank writes to a subdirectory.")
+
+    # Generation
+    p.add_argument("--max-tokens", type=int, default=16384,
+                   help="Max new tokens per response (default: 16384).")
+    p.add_argument("--temperature", type=float, default=0.7,
+                   help="Sampling temperature (default: 0.7).")
+    p.add_argument("--top-p", type=float, default=0.9,
+                   help="Nucleus sampling top-p (default: 0.9).")
+    p.add_argument("--num-responses", type=int, default=1,
+                   help="Number of responses per prompt (default: 1).")
+    p.add_argument("--batch-size", type=int, default=32,
+                   help="Prompts per vLLM batch call (default: 32).")
+
+    # Parallelism
+    p.add_argument("--tensor-parallel-size", type=int, default=1,
+                   help="vLLM tensor-parallel size (default: 1).")
+    p.add_argument("--rank", type=int, default=None,
+                   help="Worker rank (auto-detected from env if omitted).")
+    p.add_argument("--world-size", type=int, default=None,
+                   help="Total workers (auto-detected from env if omitted).")
+
+    # Misc
+    p.add_argument("--num-samples", type=int, default=None,
+                   help="Limit total samples before sharding (for debugging).")
+    p.add_argument("--checkpoint-dir", type=str, default="checkpoints",
+                   help="Directory for per-rank checkpoint files (default: checkpoints).")
+
+    args = p.parse_args()
+
+    # Auto-detect rank / world_size from environment (torchrun, etc.)
+    if args.rank is None:
+        args.rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0)))
+    if args.world_size is None:
+        args.world_size = int(os.environ.get("WORLD_SIZE", 1))
+
+    return args
+
+
+if __name__ == "__main__":
+    run_curation(parse_args())
--- a/data_curation/prepare_lightning_opd.py
+++ b/data_curation/prepare_lightning_opd.py
@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Prepare Lightning OPD parquet from student rollout data.
+
+Phase 1 – tokenize (CPU-friendly):
+    Reads student rollout parquet, builds prompt via chat template,
+    tokenizes responses, truncates to --max-response-len, writes intermediate
+    parquet WITHOUT teacher logprobs.
+
+Phase 2 – precompute teacher logprobs (requires GPU / teacher sglang server):
+    Reads the intermediate parquet produced in Phase 1, sends each
+    (prompt + response) sequence to the teacher sglang server, stores
+    per-token response logprobs back into the metadata, writes the final
+    parquet.
+
+Usage (Phase 1, CPU node):
+    python3 data_curation/prepare_lightning_opd.py \\
+        --tokenizer-path checkpoints/sft \\
+        --input-parquet data/rollouts/rollouts.parquet \\
+        --output-dir data/lightning_opd
+
+Usage (Phase 2, GPU node with teacher sglang running):
+    python3 data_curation/prepare_lightning_opd.py \\
+        --tokenizer-path checkpoints/sft \\
+        --input-parquet data/rollouts/rollouts.parquet \\
+        --output-dir data/lightning_opd \\
+        --compute-teacher-logprobs \\
+        --teacher-url http://127.0.0.1:13141/generate
+"""
+
+import argparse
+import asyncio
+from pathlib import Path
+
+import aiohttp
+import pandas as pd
+from transformers import AutoTokenizer
+from tqdm import tqdm
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Prepare Lightning OPD parquet data (tokenize + optional teacher logprobs)."
+    )
+    parser.add_argument(
+        "--tokenizer-path", type=str, required=True,
+        help="Path to HuggingFace tokenizer (e.g. the student SFT checkpoint).",
+    )
+    parser.add_argument(
+        "--input-parquet", type=str, required=True,
+        help="Path to student rollout parquet. Expected columns: messages (list[dict]), tokens (int).",
+    )
+    parser.add_argument(
+        "--output-dir", type=str, required=True,
+        help="Directory where intermediate and final parquet files are written.",
+    )
+    parser.add_argument(
+        "--max-response-len", type=int, default=4096,
+        help="Maximum response token length; longer responses are truncated (default: 4096).",
+    )
+    parser.add_argument(
+        "--compute-teacher-logprobs", action="store_true",
+        help="Run Phase 2: compute teacher logprobs via a running sglang server.",
+    )
+    parser.add_argument(
+        "--teacher-url", type=str, default="http://127.0.0.1:13141/generate",
+        help="Teacher sglang server URL (default: http://127.0.0.1:13141/generate).",
+    )
+    parser.add_argument(
+        "--concurrency", type=int, default=64,
+        help="Number of concurrent requests to teacher sglang server (default: 64).",
+    )
+    return parser.parse_args()
+
+
+# ── Phase 1: tokenize ────────────────────────────────────────────────────────
+
+def phase1_tokenize(args, intermediate_path: Path):
+    print(f"[Phase 1] Loading tokenizer from {args.tokenizer_path}")
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+
+    print(f"[Phase 1] Loading input parquet: {args.input_parquet}")
+    df = pd.read_parquet(args.input_parquet)
+    print(f"[Phase 1] Total rows: {len(df)}")
+
+    rows_out = []
+    truncated = 0
+    skipped = 0
+
+    for row in tqdm(df.itertuples(), total=len(df), desc="Tokenizing"):
+        messages = row.messages
+
+        user_messages = [m for m in messages if m["role"] != "assistant"]
+        prompt_str = tokenizer.apply_chat_template(
+            user_messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
+        )
+
+        assistant_msg = None
+        for msg in messages:
+            if msg["role"] == "assistant":
+                assistant_msg = msg["content"]
+                break
+        if assistant_msg is None:
+            skipped += 1
+            continue
+
+        response_ids = tokenizer.encode(assistant_msg, add_special_tokens=False)
+
+        if len(response_ids) > args.max_response_len:
+            truncated += 1
+            response_ids = response_ids[:args.max_response_len]
+            assistant_msg = tokenizer.decode(response_ids, skip_special_tokens=False)
+
+        rows_out.append({
+            "prompt": prompt_str,
+            "label": "0",
+            "metadata": {
+                "is_lightning_opd": True,
+                "response_tokens": response_ids,
+                "loss_mask": [1] * len(response_ids),
+                "response": assistant_msg,
+            },
+        })
+
+    print(f"[Phase 1] Rows written: {len(rows_out)}, "
+          f"truncated to {args.max_response_len}: {truncated}, skipped: {skipped}")
+    df_out = pd.DataFrame(rows_out)
+    intermediate_path.parent.mkdir(parents=True, exist_ok=True)
+    df_out.to_parquet(intermediate_path, index=False)
+    print(f"[Phase 1] Saved to {intermediate_path}")
+
+
+# ── Phase 2: precompute teacher logprobs ─────────────────────────────────────
+
+async def _fetch_logprobs(
+    session: aiohttp.ClientSession,
+    teacher_url: str,
+    full_ids: list[int],
+    response_len: int,
+) -> list[float]:
+    """Call teacher sglang server and return per-token logprobs for the response portion."""
+    payload = {
+        "input_ids": full_ids,
+        "sampling_params": {
+            "temperature": 0,
+            "max_new_tokens": 0,
+            "skip_special_tokens": False,
+        },
+        "return_logprob": True,
+        "logprob_start_len": 0,
+    }
+    async with session.post(teacher_url, json=payload) as resp:
+        resp.raise_for_status()
+        ret = await resp.json()
+
+    all_lps = ret["meta_info"]["input_token_logprobs"]
+    response_lps = [float(item[0]) for item in all_lps[1:]][-response_len:]
+    assert len(response_lps) == response_len, (
+        f"Expected {response_len} logprobs, got {len(response_lps)}"
+    )
+    return response_lps
+
+
+async def _process_all(args, tokenizer, rows: list[dict]) -> list[list[float]]:
+    """Process all rows concurrently with a live progress bar, preserving order."""
+    semaphore = asyncio.Semaphore(args.concurrency)
+    connector = aiohttp.TCPConnector(limit=args.concurrency)
+    results = [None] * len(rows)
+
+    async def bounded_fetch(idx: int, full_ids: list[int], response_len: int):
+        async with semaphore:
+            result = await _fetch_logprobs(session, args.teacher_url, full_ids, response_len)
+        results[idx] = result
+        pbar.update(1)
+
+    async with aiohttp.ClientSession(connector=connector) as session:
+        with tqdm(total=len(rows), desc="[Phase 2] Teacher logprobs") as pbar:
+            tasks = []
+            for idx, row in enumerate(rows):
+                meta = row["metadata"]
+                prompt_ids = tokenizer.encode(row["prompt"], add_special_tokens=False)
+                response_ids = [int(x) for x in meta["response_tokens"]]
+                full_ids = prompt_ids + response_ids
+                tasks.append(bounded_fetch(idx, full_ids, len(response_ids)))
+            await asyncio.gather(*tasks)
+
+    return results
+
+
+def phase2_logprobs(args, intermediate_path: Path, output_path: Path):
+    print(f"[Phase 2] Loading intermediate parquet: {intermediate_path}")
+    df = pd.read_parquet(intermediate_path)
+    rows = df.to_dict(orient="records")
+    print(f"[Phase 2] Total rows: {len(rows)}")
+
+    print(f"[Phase 2] Loading tokenizer from {args.tokenizer_path}")
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+
+    print(f"[Phase 2] Computing teacher logprobs via {args.teacher_url} "
+          f"(concurrency={args.concurrency})")
+    all_logprobs = asyncio.run(_process_all(args, tokenizer, rows))
+
+    for row, lps in zip(rows, all_logprobs):
+        row["metadata"]["teacher_log_probs"] = lps
+
+    df_out = pd.DataFrame(rows)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    df_out.to_parquet(output_path, index=False)
+    print(f"[Phase 2] Saved to {output_path}")
+
+    # Sanity check
+    df_check = pd.read_parquet(output_path)
+    row0 = df_check.iloc[0]
+    meta = row0["metadata"]
+    print("\n[Phase 2] Sanity check row 0:")
+    print(f"  prompt[:80]:              {row0['prompt'][:80]}")
+    print(f"  label:                    {row0['label']}")
+    print(f"  len(response_tokens):     {len(meta['response_tokens'])}")
+    print(f"  len(teacher_log_probs):   {len(meta['teacher_log_probs'])}")
+    print(f"  teacher_log_probs[:5]:    {meta['teacher_log_probs'][:5]}")
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+def main():
+    args = parse_args()
+
+    output_dir = Path(args.output_dir)
+    input_stem = Path(args.input_parquet).stem
+    intermediate_path = output_dir / f"{input_stem}-lightning-opd.parquet"
+    output_path = output_dir / f"{input_stem}-lightning-opd-precomputed.parquet"
+
+    if args.compute_teacher_logprobs:
+        if not intermediate_path.exists():
+            print("[INFO] Intermediate parquet not found, running Phase 1 first.")
+            phase1_tokenize(args, intermediate_path)
+        phase2_logprobs(args, intermediate_path, output_path)
+    else:
+        phase1_tokenize(args, intermediate_path)
+        print(f"\n[INFO] To add teacher logprobs, re-run with --compute-teacher-logprobs "
+              f"after starting the teacher sglang server.")
+
+
+if __name__ == "__main__":
+    main()
--- a/data_curation/run_curation.sh
+++ b/data_curation/run_curation.sh
@@ -0,0 +1,126 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+#
+# Launch data curation across multiple GPUs / nodes.
+#
+# Each GPU runs one independent vLLM worker that processes a disjoint shard
+# of the dataset. No torch.distributed communication is needed — each worker
+# is a standalone process with its own rank derived from environment variables.
+#
+# ── Single node, 8 GPUs (tp=1, 8 workers) ────────────────────────────────
+#   bash data_curation/run_curation.sh \
+#       --model Qwen/Qwen3-4B \
+#       --input data.jsonl \
+#       --output-dir output/ \
+#       --num-gpus 8
+#
+# ── Single node, 2 GPUs (tp=2, 1 worker) ─────────────────────────────────
+#   bash data_curation/run_curation.sh \
+#       --model Qwen/Qwen3-8B \
+#       --input data.jsonl \
+#       --output-dir output/ \
+#       --num-gpus 2 \
+#       --tensor-parallel-size 2
+#
+# ── Multi-node (2 nodes × 8 GPUs, tp=1, 16 workers) ─────────────────────
+#   # On node 0:
+#   NODE_RANK=0 NUM_NODES=2 bash data_curation/run_curation.sh \
+#       --model Qwen/Qwen3-4B \
+#       --input data.jsonl \
+#       --output-dir output/ \
+#       --num-gpus 8
+#
+#   # On node 1:
+#   NODE_RANK=1 NUM_NODES=2 bash data_curation/run_curation.sh \
+#       --model Qwen/Qwen3-4B \
+#       --input data.jsonl \
+#       --output-dir output/ \
+#       --num-gpus 8
+#
+# Environment variables (optional):
+#   NUM_NODES   – total number of nodes (default: 1)
+#   NODE_RANK   – rank of this node (default: 0)
+# ──────────────────────────────────────────────────────────────────────────
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# ── Parse --num-gpus and --tensor-parallel-size from args ─────────────────
+NUM_GPUS=1
+TP=1
+PIPELINE_ARGS=()
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --num-gpus)
+            NUM_GPUS="$2"; shift 2 ;;
+        --tensor-parallel-size)
+            TP="$2"; PIPELINE_ARGS+=("--tensor-parallel-size" "$2"); shift 2 ;;
+        *)
+            PIPELINE_ARGS+=("$1"); shift ;;
+    esac
+done
+
+# ── Compute worker layout ────────────────────────────────────────────────
+NUM_NODES="${NUM_NODES:-1}"
+NODE_RANK="${NODE_RANK:-0}"
+WORKERS_PER_NODE=$(( NUM_GPUS / TP ))
+WORLD_SIZE=$(( WORKERS_PER_NODE * NUM_NODES ))
+
+echo "=== Data Curation Launch ==="
+echo "  Nodes:            ${NUM_NODES} (this node: ${NODE_RANK})"
+echo "  GPUs per node:    ${NUM_GPUS}"
+echo "  TP size:          ${TP}"
+echo "  Workers per node: ${WORKERS_PER_NODE}"
+echo "  World size:       ${WORLD_SIZE}"
+echo "  Pipeline args:    ${PIPELINE_ARGS[*]}"
+echo "============================"
+
+# ── Launch workers ───────────────────────────────────────────────────────
+PIDS=()
+for (( LOCAL=0; LOCAL<WORKERS_PER_NODE; LOCAL++ )); do
+    GLOBAL_RANK=$(( NODE_RANK * WORKERS_PER_NODE + LOCAL ))
+    GPU_START=$(( LOCAL * TP ))
+    GPU_END=$(( GPU_START + TP - 1 ))
+
+    # Build CUDA_VISIBLE_DEVICES string, e.g. "0" or "2,3"
+    GPUS=""
+    for (( g=GPU_START; g<=GPU_END; g++ )); do
+        [[ -n "$GPUS" ]] && GPUS="${GPUS},"
+        GPUS="${GPUS}${g}"
+    done
+
+    echo "[Node ${NODE_RANK}] Launching worker rank=${GLOBAL_RANK} on GPU(s) ${GPUS}"
+
+    CUDA_VISIBLE_DEVICES="${GPUS}" \
+    RANK="${GLOBAL_RANK}" \
+    WORLD_SIZE="${WORLD_SIZE}" \
+    python "${SCRIPT_DIR}/pipeline.py" \
+        --rank "${GLOBAL_RANK}" \
+        --world-size "${WORLD_SIZE}" \
+        "${PIPELINE_ARGS[@]}" \
+        > >(sed "s/^/[rank${GLOBAL_RANK}] /") \
+        2>&1 &
+
+    PIDS+=($!)
+done
+
+# ── Wait for all workers ─────────────────────────────────────────────────
+echo "Waiting for ${#PIDS[@]} workers to finish..."
+FAILED=0
+for PID in "${PIDS[@]}"; do
+    if ! wait "$PID"; then
+        echo "Worker PID ${PID} failed!"
+        FAILED=1
+    fi
+done
+
+if [[ $FAILED -eq 1 ]]; then
+    echo "Some workers failed. Check logs above."
+    exit 1
+fi
+
+echo "All workers finished successfully."