sdg_pipeline: multi-model profiling stage replaces single-model difficulty_estimation

Introduce a new `profiling` stage that runs per-model generate->judge->aggregate
chains in parallel for N models and merges results into a single `profiling`
array per problem:

    "profiling": [
        {"model": "ModelA", "pass_rate": 0.5, "pass_at_n": "2/4"},
        {"model": "ModelB", "pass_rate": 0.8, "pass_at_n": "4/5"},
    ]

Changes:
- `run_pipeline.profiling()` orchestrates: shared prepare -> per-model chains
  (generate, judge, aggregate) in parallel -> final merge. Judge kwargs are
  copied per-iteration so args don't leak across models; `num_random_seeds`
  is inherited from generation if not explicitly set.
- `aggregate_profiling_model.py` (new): per-model aggregator over per-seed
  `output-rs*.jsonl` files. Streams inputs — keeps only the BASE_FIELDS
  projection + a small counters dict per `(id, problem)` key — so aggregation
  fits in memory at 1M+ problem scale. Falls back to `(_lineno, line_number)`
  when neither `id` nor `problem` is present on a record.
- `merge_profiling.py` (new): merges per-model result files. Asserts every
  per-model file contains the same `(id, problem)` key set so row-alignment
  mismatches fail loudly instead of silently dropping problems. After a
  successful merge, removes the per-model `result.jsonl` intermediates
  (folders — generation/, judgement/, logs/ — are retained for debugging).
- `filter_solutions.py`: replaces the scalar `difficulty_model_pass_rate`
  bounds with a per-model dict `profiling_pass_rate_ranges:
  {model_name: [min, max]}` (min exclusive, max inclusive).
- `validate_pipeline.py` and `scripts/utils/constants.py`: update stage-name
  and field-set checks (`PROFILING_FIELDS`, required `profiling` key,
  row-count equality for the new stage).
- Base + settings YAMLs: renamed stage (`difficulty_estimation` -> `profiling`)
  and directory (`step-3-difficulty-estimation` -> `step-3-profiling`); new
  `profiling.models: [...]` list with per-model `generation_kwargs` and an
  optional per-model `judge_kwargs` override.
- SLURM test: references `stages.profiling.models.0.generation_kwargs...`
  instead of the old top-level path.
- README: updates the stage list + filter-parameter description.
- `aggregate_difficulty.py` is removed (replaced by the new aggregator +
  merger).

Signed-off-by: Tatevik Ter-Hovhannisyan <tterhovhanni@nvidia.com>

Author: tatevik-t

recipes/opensciencereasoning/sdg_pipeline/README.md

@@ -5,7 +5,7 @@ This folder provides templates, prompts, and scripts for the automated pipeline
 - Deduplicate and clean incoming problems via [`filter_problems`](scripts/filter_problems.py).
 - Run contamination checks in [`decontaminate`](scripts/decontaminate.py).
 - Launch [`generate_solutions`](run_pipeline.py) to obtain model answers when no GT is supplied, then run majority voting to recover a GT answer. Will only be applied with the `without_gt` setting.
-- Score questions with [`difficulty_estimation`](run_pipeline.py) and enrich metadata with [`topics_labeling`](run_pipeline.py).
+- Profile problem difficulty with multiple models via [`profiling`](run_pipeline.py) and enrich metadata with [`topics_labeling`](run_pipeline.py).
 - Finish with [`aggregate`](scripts/aggregate_metadata.py) and [`filter_solutions`](scripts/filter_solutions.py) to produce deliverables.
 
 ## SFT Data Flow
@@ -27,10 +27,10 @@ This folder provides templates, prompts, and scripts for the automated pipeline
   - `generation/output*.jsonl`: raw generations.
   - `with_predictions/output*.jsonl`: adds `predicted_answer`, and when the majority answer is applied, also adds `expected_answer`, `majority_voting_agreement_rate`, and `majority_voting_agreement_at_n`.
   - Optional `judgement/output*.jsonl`: contains `judgement` strings when `make_judgement` is enabled. The aggregated stage output also adds `is_correct`, `generation_model_pass_rate`, `generation_model_pass_at_n`, and `generation_model` to each sample.
-- [`difficulty_estimation`](run_pipeline.py): Requires GT answers. Uses [`remove_redundant_fields.py`](scripts/remove_redundant_fields.py) to keep baseline keys, generates boxed-format solutions (`generation_kwargs`), judges them (`judge_kwargs`), and writes `final_result.jsonl` with `difficulty_model`, `difficulty_model_pass_rate`, and `difficulty_model_pass_at_n` fields (see [`aggregate_difficulty.py`](scripts/aggregate_difficulty.py)).
+- [`profiling`](run_pipeline.py): Requires GT answers. Runs multiple models in parallel, each through a generate-judge-aggregate chain. Uses [`remove_redundant_fields.py`](scripts/remove_redundant_fields.py) to keep baseline keys, generates boxed-format solutions per model, judges them, and writes `final_result.jsonl` with a `profiling` array field containing per-model `{model, pass_rate, pass_at_n}` entries (see [`aggregate_profiling_model.py`](scripts/aggregate_profiling_model.py) and [`merge_profiling.py`](scripts/merge_profiling.py)).
 - [`aggregate`](scripts/aggregate_metadata.py): Merges metadata (`metadata_files`) and optional solution glob (`solutions_path`) into `final_result.jsonl`. The resulting records combine base fields with appended metadata and solution statistics.
 - [`prepare_for_sft`](run_pipeline.py): Calls `nemo_skills.training.prepare_data` via the configured `prepare_data_kwargs` (tokenizer, prompt config, formatting toggles). Outputs an instruction-tuning JSONL file.
-- [`filter_solutions`](scripts/filter_solutions.py): Applies correctness/pass-rate/metadata filters. Parameters: `only_correct_solutions`, `generation_model_pass_rate_range`, `difficulty_model_pass_rate_range`, `metadata_values`, `only_samples_with_ground_truth_answer`. The filtered output preserves the same schema as the input `final_result.jsonl`.
+- [`filter_solutions`](scripts/filter_solutions.py): Applies correctness/pass-rate/metadata filters. Parameters: `only_correct_solutions`, `generation_model_pass_rate_range`, `profiling_pass_rate_range` (JSON dict `{model_name: [min, max]}`), `metadata_values`, `only_samples_with_ground_truth_answer`. The filtered output preserves the same schema as the input `final_result.jsonl`.
 - [`process_messages_and_bucket`](run_pipeline.py): Uses [`scripts/process_messages_and_bucket.py`](scripts/process_messages_and_bucket.py) to transform prepared rows into input/output message text, compute `input_token_length` and `output_token_length`, and optionally split into token-length buckets based on `bucket_field` and `bucket_sizes`.
 - [`validate`](scripts/validate_pipeline.py): Reuses the automated checker to verify artifacts exist, counts add up, and required metadata fields are present, so failures point directly to the problematic stage. See [What the Validation Stage Covers](#what-the-validation-stage-covers) for details and caveats.
 

recipes/opensciencereasoning/sdg_pipeline/configs/pipelines/base.yaml

@@ -21,7 +21,7 @@ pipeline_stages:
   - decontaminate               # Decontaminate problems
   - topics_labeling             # Label topics and subtopics
   - generate_solutions          # Generate solutions
-  - difficulty_estimation       # Estimate difficulty of problems
+  - profiling                   # Profile problem difficulty with multiple models
   - aggregate                   # Aggregate all the metadata into a single file
   - filter_solutions            # Filter solutions
   - prepare_for_sft             # Prepare for SFT
@@ -33,7 +33,7 @@ directories:
   step-0-filter-problems: ${base_output_dir}/solution-sdg/step-0-filter-problems
   step-1-decontaminate: ${base_output_dir}/solution-sdg/step-1-decontaminate
   step-2-topics-labeling: ${base_output_dir}/solution-sdg/step-2-topics-labeling
-  step-3-difficulty-estimation: ${base_output_dir}/solution-sdg/step-3-difficulty-estimation
+  step-3-profiling: ${base_output_dir}/solution-sdg/step-3-profiling
   step-4-generate-solutions: ${base_output_dir}/solution-sdg/step-4-generate-solutions
   step-5-aggregate: ${base_output_dir}/solution-sdg/step-5-aggregate
   step-6-filter-solutions: ${base_output_dir}/solution-sdg/step-6-filter-solutions
@@ -134,35 +134,37 @@ stages:
     dependencies:
       - decontaminate
 
-  difficulty_estimation:
-    output_dir: ${directories.step-3-difficulty-estimation}
+  profiling:
+    output_dir: ${directories.step-3-profiling}
     input_file: ${directories.step-1-decontaminate}/final_result.jsonl  # Should have expected answers
 
-    generation_kwargs:
-      args:
-        model: /hf_models/Qwen3-30B-A3B
-        server_type: vllm
-        server_gpus: 8
-        server_nodes: 1
-        dependent_jobs: 1
-        num_random_seeds: 5
-        num_chunks: 20
-      ctx_args: >-
-        ++prompt_config=generic/general-boxed
-        ++inference.tokens_to_generate=16000
-
+    # Shared judge config -- used for all models unless overridden per-model
     judge_kwargs:
       args:
         model: /hf_models/gpt-oss-20b
         server_type: vllm
         server_gpus: 8
         server_nodes: 1
-        num_random_seeds: ${stages.difficulty_estimation.generation_kwargs.args.num_random_seeds}
         dependent_jobs: 1
         num_chunks: 5
       ctx_args: >-
         ++prompt_config=judge/general-judge
 
+    models:
+      - name: Qwen3-30B-A3B
+        generation_kwargs:
+          args:
+            model: /hf_models/Qwen3-30B-A3B
+            server_type: vllm
+            server_gpus: 8
+            server_nodes: 1
+            dependent_jobs: 1
+            num_random_seeds: 5
+            num_chunks: 20
+          ctx_args: >-
+            ++prompt_config=generic/general-boxed
+            ++inference.tokens_to_generate=16000
+
     dependencies:
       - decontaminate
 
@@ -232,18 +234,18 @@ stages:
     solutions_path: ${directories.step-4-generate-solutions}/final_result.jsonl
     metadata_files:
       - ${directories.step-2-topics-labeling}/final_result.jsonl
-      - ${directories.step-3-difficulty-estimation}/final_result.jsonl
+      - ${directories.step-3-profiling}/final_result.jsonl
     dependencies:
       - topics_labeling
-      - difficulty_estimation
+      - profiling
       - generate_solutions
 
   filter_solutions:
     output_dir: ${directories.step-6-filter-solutions}
     input_file: ${directories.step-5-aggregate}/final_result.jsonl
     only_correct_solutions: True
     generation_model_pass_rate_range: [-1.0, 1.0] # minimum exclusive, maximum inclusive
-    difficulty_model_pass_rate_range: [-1.0, 1.0] # minimum exclusive, maximum inclusive
+    profiling_pass_rate_range: null  # Optional: JSON dict {"ModelName": [min, max]} for per-model filtering
     only_samples_with_ground_truth_answer: True
     metadata_values:
       topic: ["Biology", "Chemistry", "Physics", "Mathematics", "Other", "undefined"]

recipes/opensciencereasoning/sdg_pipeline/configs/settings/multiple_prompts.yaml

@@ -17,7 +17,7 @@ stages:
     enabled: False
   topics_labeling:
     enabled: False
-  difficulty_estimation:
+  profiling:
     enabled: False
   aggregate:
     enabled: False

recipes/opensciencereasoning/sdg_pipeline/configs/settings/seed_data.yaml

@@ -18,7 +18,7 @@ stages:
     solutions_path: null
     dependencies:
       - topics_labeling
-      - difficulty_estimation
+      - profiling
   filter_solutions:
     input_file: ${directories.step-4-aggregate}/final_result.jsonl
     output_dir: ${directories.step-5-filter-solutions}

recipes/opensciencereasoning/sdg_pipeline/configs/settings/seed_data_postprocess.yaml

@@ -7,7 +7,7 @@ stages:
     enabled: False
   topics_labeling:
     enabled: False
-  difficulty_estimation:
+  profiling:
     enabled: False
   generate_solutions:
     generation_kwargs:
@@ -21,5 +21,5 @@ stages:
   filter_solutions:
     only_correct_solutions: True
     is_ground_truth_answer_present: True
-    difficulty_model_pass_rate_range: null
+    profiling_pass_rate_range: null
     metadata_values: null

recipes/opensciencereasoning/sdg_pipeline/configs/settings/without_gt.yaml

@@ -2,7 +2,7 @@
 
 directories:
   step-3-generate-solutions: ${base_output_dir}/solution-sdg/step-3-generate-solutions
-  step-4-difficulty-estimation: ${base_output_dir}/solution-sdg/step-4-difficulty-estimation
+  step-4-profiling: ${base_output_dir}/solution-sdg/step-4-profiling
   step-5-aggregate: ${base_output_dir}/solution-sdg/step-5-aggregate
   step-6-filter-solutions: ${base_output_dir}/solution-sdg/step-6-filter-solutions
 
@@ -13,17 +13,17 @@ stages:
     make_majority_voting: True
     dependencies:
       - decontaminate
-  difficulty_estimation:
+  profiling:
     input_file: ${directories.step-3-generate-solutions}/final_result.jsonl
-    output_dir: ${directories.step-4-difficulty-estimation}
+    output_dir: ${directories.step-4-profiling}
     dependencies:
       - generate_solutions
   aggregate:
     output_dir: ${directories.step-5-aggregate}
     solutions_path: ${directories.step-3-generate-solutions}/final_result.jsonl
     metadata_files:
       - ${directories.step-2-topics-labeling}/final_result.jsonl
-      - ${directories.step-4-difficulty-estimation}/final_result.jsonl
+      - ${directories.step-4-profiling}/final_result.jsonl
   filter_solutions:
     input_file: ${directories.step-5-aggregate}/final_result.jsonl
     output_dir: ${directories.step-6-filter-solutions}

recipes/opensciencereasoning/sdg_pipeline/run_pipeline.py

@@ -410,32 +410,29 @@ def generate_solutions(cluster, expname, run_after, stage_config, **kwargs):
     )
 
 
-def difficulty_estimation(cluster, expname, run_after, stage_config, **kwargs):
-    """Run difficulty estimation generation, judge correctness, and postprocess metrics.
+def profiling(cluster, expname, run_after, stage_config, **kwargs):
+    """Run multi-model profiling: generate, judge, and aggregate per model in parallel, then merge.
 
-    This stage:
-      - Generates multiple solutions per problem using the provided model/prompt.
-      - Runs LLM-based judging (math_judge) over those generations to get Yes/No per sample.
-      - Postprocesses the judgements to append three keys to the final results file:
-        - difficulty_model: the model used for generation
-        - difficulty_model_pass_rate: decimal fraction of correct judgements (e.g., 0.5)
-        - difficulty_model_pass_at_n: formatted fraction "correct/total" (e.g., 2/4)
+    This stage runs N models in parallel, each through a generate -> judge -> aggregate chain.
+    All models share a common input preparation step. A final merge step combines all per-model
+    results into a single profiling array per problem.
+
+    Output format per record:
+        "profiling": [
+            {"model": "ModelA", "pass_rate": 0.5, "pass_at_n": "2/4"},
+            {"model": "ModelB", "pass_rate": 0.8, "pass_at_n": "4/5"},
+        ]
 
     Note: The judging step extracts predicted answers using the \\boxed{...} convention.
     It will only work out-of-the-box when generations include a final answer in boxed format.
     """
     output_dir = stage_config["output_dir"]
     input_file = stage_config["input_file"]
+    models = stage_config.get("models", [])
+    shared_judge_kwargs = stage_config.get("judge_kwargs", {})
 
-    generation_kwargs = stage_config.get("generation_kwargs", {})
-    judge_kwargs = stage_config.get("judge_kwargs", {})
-
-    generation_args = generation_kwargs.get("args", {})
-    generation_ctx_args = generation_kwargs.get("ctx_args", "")
-
-    judge_args = judge_kwargs.get("args", {})
-    judge_ctx_args = judge_kwargs.get("ctx_args", "")
-
+    # Step 1: Shared prepare job
+    prepare_expname = f"{expname}_prepare_profiling"
     run_cmd(
         ctx=wrap_arguments(
             f"python /nemo_run/code/recipes/opensciencereasoning/sdg_pipeline/scripts/remove_redundant_fields.py "
@@ -446,41 +443,85 @@ def difficulty_estimation(cluster, expname, run_after, stage_config, **kwargs):
         ),
         cluster=cluster,
         log_dir=f"{output_dir}/tmp/logs",
-        expname=f"{expname}_prepare_difficulty_estimation",
+        expname=prepare_expname,
         run_after=run_after,
     )
 
-    generate(
-        ctx=wrap_arguments(generation_ctx_args),
-        cluster=cluster,
-        input_file=f"{output_dir}/tmp/prepared.jsonl",
-        output_dir=f"{output_dir}/generation",
-        expname=f"{expname}-generation",
-        run_after=f"{expname}_prepare_difficulty_estimation",
-        **generation_args,
-    )
+    # Step 2: Per-model generate -> judge -> aggregate chains (all depend on prepare, run in parallel)
+    per_model_aggregate_expnames = []
+    per_model_result_files = []
 
-    generate(
-        ctx=wrap_arguments(judge_ctx_args),
-        generation_type="math_judge",
-        cluster=cluster,
-        input_dir=f"{output_dir}/generation",
-        output_dir=f"{output_dir}/judgement",
-        expname=f"{expname}-judgement",
-        run_after=f"{expname}-generation",
-        **judge_args,
-    )
+    for model_cfg in models:
+        model_name = model_cfg["name"]
+        safe_name = model_name.replace("/", "-").lower()
+        model_dir = f"{output_dir}/{safe_name}"
+
+        generation_kwargs = model_cfg.get("generation_kwargs", {})
+        generation_args = generation_kwargs.get("args", {})
+        generation_ctx_args = generation_kwargs.get("ctx_args", "")
+
+        # Per-model judge_kwargs: use model-level override if present, else shared
+        judge_kwargs = model_cfg.get("judge_kwargs", shared_judge_kwargs)
+        judge_args = judge_kwargs.get("args", {})
+        judge_ctx_args = judge_kwargs.get("ctx_args", "")
+
+        # Ensure judge num_random_seeds matches generation if not explicitly set
+        if "num_random_seeds" not in judge_args and "num_random_seeds" in generation_args:
+            judge_args = dict(judge_args)
+            judge_args["num_random_seeds"] = generation_args["num_random_seeds"]
+
+        gen_expname = f"{expname}-{safe_name}-generation"
+        judge_expname = f"{expname}-{safe_name}-judgement"
+        agg_expname = f"{expname}-{safe_name}-aggregate"
+
+        generate(
+            ctx=wrap_arguments(generation_ctx_args),
+            cluster=cluster,
+            input_file=f"{output_dir}/tmp/prepared.jsonl",
+            output_dir=f"{model_dir}/generation",
+            expname=gen_expname,
+            run_after=prepare_expname,
+            **generation_args,
+        )
+
+        generate(
+            ctx=wrap_arguments(judge_ctx_args),
+            generation_type="math_judge",
+            cluster=cluster,
+            input_dir=f"{model_dir}/generation",
+            output_dir=f"{model_dir}/judgement",
+            expname=judge_expname,
+            run_after=gen_expname,
+            **judge_args,
+        )
+
+        model_result_file = f"{model_dir}/result.jsonl"
+        run_cmd(
+            ctx=wrap_arguments(
+                f"python /nemo_run/code/recipes/opensciencereasoning/sdg_pipeline/scripts/aggregate_profiling_model.py "
+                f"    --judgement_dir '{model_dir}/judgement' "
+                f"    --output_file '{model_result_file}' "
+                f"    --model_name '{model_name}' "
+            ),
+            cluster=cluster,
+            log_dir=f"{model_dir}/logs",
+            run_after=judge_expname,
+            expname=agg_expname,
+        )
+
+        per_model_aggregate_expnames.append(agg_expname)
+        per_model_result_files.append(model_result_file)
 
+    # Step 3: Merge all per-model results into a single profiling array
     run_cmd(
         ctx=wrap_arguments(
-            f"python /nemo_run/code/recipes/opensciencereasoning/sdg_pipeline/scripts/aggregate_difficulty.py "
-            f"    --judgement_dir '{output_dir}/judgement' "
+            f"python /nemo_run/code/recipes/opensciencereasoning/sdg_pipeline/scripts/merge_profiling.py "
+            f"    --model_result_files {shlex.quote(json.dumps(per_model_result_files, ensure_ascii=False))} "
             f"    --output_file '{output_dir}/{OUTPUT_FILE}' "
-            f"    --difficulty_model '{generation_args['model'].split('/')[-1]}' "
         ),
         cluster=cluster,
         log_dir=f"{output_dir}/logs",
-        run_after=f"{expname}-judgement",
+        run_after=per_model_aggregate_expnames,
         expname=expname,
     )
 
@@ -519,7 +560,7 @@ def filter_solutions(cluster, expname, run_after, stage_config, **kwargs):
     Supported filters (see `filter_solutions.py`):
       - `only_correct_solutions`: keep only samples marked `is_correct`.
       - `generation_model_pass_rate_range`: JSON `[min, max]` range (min exclusive, max inclusive).
-      - `difficulty_model_pass_rate_range`: JSON `[min, max]` range over difficulty pass rates.
+      - `profiling_pass_rate_range`: JSON dict `{model_name: [min, max]}` for per-model profiling filtering.
       - `metadata_values`: dict of field -> allowed values.
 
     Replace `filter_solutions.py` with your own implementation if custom filtering logic is required.
@@ -528,7 +569,7 @@ def filter_solutions(cluster, expname, run_after, stage_config, **kwargs):
     input_file = stage_config["input_file"]
     only_correct_solutions = stage_config.get("only_correct_solutions", False)
     generation_model_pass_rate_range = stage_config.get("generation_model_pass_rate_range", None)
-    difficulty_model_pass_rate_range = stage_config.get("difficulty_model_pass_rate_range", None)
+    profiling_pass_rate_range = stage_config.get("profiling_pass_rate_range", None)
     metadata_values = stage_config.get("metadata_values", None)
     only_samples_with_ground_truth_answer = stage_config.get("only_samples_with_ground_truth_answer", False)
 
@@ -537,9 +578,9 @@ def filter_solutions(cluster, expname, run_after, stage_config, **kwargs):
         if generation_model_pass_rate_range
         else ""
     )
-    difficulty_model_pass_rate_range_arg = (
-        f"    --difficulty_model_pass_rate_range {shlex.quote(json.dumps(difficulty_model_pass_rate_range, ensure_ascii=False))} "
-        if difficulty_model_pass_rate_range
+    profiling_pass_rate_range_arg = (
+        f"    --profiling_pass_rate_range {shlex.quote(json.dumps(profiling_pass_rate_range, ensure_ascii=False))} "
+        if profiling_pass_rate_range
         else ""
     )
     metadata_values_arg = (
@@ -558,7 +599,7 @@ def filter_solutions(cluster, expname, run_after, stage_config, **kwargs):
             f"    --output_file '{output_dir}/{OUTPUT_FILE}' "
             f"{only_correct_arg}"
             f"{generation_model_pass_rate_range_arg} "
-            f"{difficulty_model_pass_rate_range_arg} "
+            f"{profiling_pass_rate_range_arg} "
             f"{metadata_values_arg} "
             f"{only_samples_with_ground_truth_answer_arg} "
         ),
@@ -752,7 +793,7 @@ def derive_variant_name():
     "decontaminate": decontaminate,
     "topics_labeling": topics_labeling,
     "generate_solutions": generate_solutions,
-    "difficulty_estimation": difficulty_estimation,
+    "profiling": profiling,
     "aggregate": aggregate,
     "filter_solutions": filter_solutions,
     "prepare_for_sft": prepare_for_sft,