feat: super v3.5 posttraining with flexible length penalties

examples/nemo_gym/nemotron-3-super/super_launch.sh

@@ -70,8 +70,26 @@ fi
 # ---- Derived paths ----
 CODE_DIR=$(realpath "$PWD")
 WANDB_NAME="${EXP_NAME}"
-CHECKPOINT_DIR="results/${EXP_NAME}"
-LOG_DIR="logs/${EXP_NAME}"
+# Optional stable results layout: set RESULTS_DIR to get
+#   $RESULTS_DIR/checkpoints             (stable -> singleton auto-resume)
+#   $RESULTS_DIR/runs/<ts>/{logs,slurm}  (per-submission; runs/latest symlink)
+#   $RESULTS_DIR/ray_logs/<jobid>-logs   (ray.sub infra logs via BASE_LOG_DIR)
+# Unset: legacy snapshot-relative results/ and logs/ dirs.
+RESULTS_DIR="${RESULTS_DIR:-}"
+if [[ -n "${RESULTS_DIR}" ]]; then
+    CHECKPOINT_DIR="${CHECKPOINT_DIR:-${RESULTS_DIR}/checkpoints}"
+    RUN_DIR="${RESULTS_DIR}/runs/$(date +%Y%m%d-%H%M)"
+    LOG_DIR="${RUN_DIR}/logs"
+    SLURM_LOG_DIR="${RUN_DIR}/slurm"
+    mkdir -p "${CHECKPOINT_DIR}" "${LOG_DIR}" "${SLURM_LOG_DIR}"
+    ln -sfn "$(realpath "${RUN_DIR}")" "${RESULTS_DIR}/runs/latest"
+    export BASE_LOG_DIR="${BASE_LOG_DIR:-${RESULTS_DIR}/ray_logs}"
+    mkdir -p "${BASE_LOG_DIR}"
+else
+    CHECKPOINT_DIR="results/${EXP_NAME}"
+    LOG_DIR="logs/${EXP_NAME}"
+    SLURM_LOG_DIR=""
+fi
 
 VLLM_CACHE_DIR="${PERSISTENT_CACHE}/vllm_compile_cache"
 FLASHINFER_CUBIN_CACHE="${PERSISTENT_CACHE}/flashinfer_cubins"
@@ -134,7 +152,7 @@ export LISTEN_PORT=6000
 export NGINX_PORT=6000
 export NEMO_SKILLS_SANDBOX_PORT=6000
 export SANDBOX_CONTAINER
-export SANDBOX_COMMAND="/start-with-nginx.sh"
+export SANDBOX_COMMAND="${SANDBOX_COMMAND:-/start-with-nginx.sh}"
 export SANDBOX_ENV_VARS="NEMO_SKILLS_SANDBOX_PORT=${NEMO_SKILLS_SANDBOX_PORT}"
 
 # ---- Build the run command ----
@@ -160,8 +178,8 @@ export COMMAND="export HF_MODULES_CACHE=${HF_MODULES_CACHE_DIR} ; \
     PYTHONPATH=${SNAPSHOT_DIR}:\${PYTHONPATH:-} \
     python ./examples/nemo_gym/run_grpo_nemo_gym.py \
     --config ${CONFIG_PATH} \
-    env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR} \
-    env.nemo_gym.skip_venv_if_present=true \
+    ++env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR} \
+    ++env.nemo_gym.skip_venv_if_present=true \
     policy.model_name=${MODEL_PATH} \
     checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \
     logger.log_dir=${LOG_DIR} \
@@ -216,11 +234,27 @@ SBATCH_CMD=(
     --job-name="${WANDB_NAME}"
     --partition="${SLURM_PARTITION}"
     --time="${SLURM_TIME_LIMIT}"
-    --gres=gpu:8
+    --gres=gpu:"${GPUS_PER_NODE:-8}"
     --exclusive
-    --dependency=singleton
+    --dependency=singleton${SLURM_EXTRA_DEPENDENCY:+,${SLURM_EXTRA_DEPENDENCY}}
     ray.sub
 )
+if [[ -n "${SLURM_QOS:-}" ]]; then
+    SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --qos="${SLURM_QOS}" "${SBATCH_CMD[@]:1}")
+fi
+if [[ -n "${SLURM_COMMENT:-}" ]]; then
+    SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --comment="${SLURM_COMMENT}" "${SBATCH_CMD[@]:1}")
+fi
+if [[ -n "${SLURM_LOG_DIR}" ]]; then
+    SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --output="${SLURM_LOG_DIR}/%j.out" --error="${SLURM_LOG_DIR}/%j.err" "${SBATCH_CMD[@]:1}")
+fi
+if [[ -n "${SLURM_SEGMENT:-}" ]]; then
+    if (( SBATCH_NUM_NODES % SLURM_SEGMENT != 0 )); then
+        echo "Error: SBATCH_NUM_NODES=${SBATCH_NUM_NODES} not divisible by SLURM_SEGMENT=${SLURM_SEGMENT}" >&2
+        exit 1
+    fi
+    SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --segment="${SLURM_SEGMENT}" "${SBATCH_CMD[@]:1}")
+fi
 
 if [[ "$DRY_RUN" == true ]]; then
     echo ""

nemo_rl/algorithms/async_utils/trajectory_collector.py

@@ -1291,6 +1291,7 @@ async def _iter_rollout_groups(
                 max_rollout_turns=None,
                 greedy=False,
                 reward_penalty_config=self.master_config.reward_penalties,
+                length_penalty_config=self.master_config.grpo.model_dump(),
                 thinking_tags=get_nemo_gym_thinking_tags(self.master_config.env),
                 mask_env_flagged_samples=should_mask_flagged_samples(
                     self.master_config.env

nemo_rl/algorithms/grpo.py

@@ -3015,6 +3015,7 @@ def grpo_train(
                             greedy=False,
                             effort_config=_get_effort_config(master_config),
                             reward_penalty_config=master_config.reward_penalties,
+                            length_penalty_config=master_config.grpo.model_dump(),
                             thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
                             mask_env_flagged_samples=should_mask_flagged_samples(
                                 master_config.env
@@ -3938,6 +3939,11 @@ def validate(
                     greedy=False,
                     effort_config=_get_effort_config(master_config),
                     reward_penalty_config=master_config.reward_penalties,
+                    # No length_penalty_config here: validation metrics
+                    # (accuracy/pass_k) must reflect the raw env reward, and the
+                    # adjustment code groups by the TRAINING stride
+                    # (num_generations_per_prompt), which does not match
+                    # val_num_generations_per_prompt.
                     thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
                     mask_env_flagged_samples=should_mask_flagged_samples(
                         master_config.env

nemo_rl/experience/rollouts.py

@@ -68,6 +68,7 @@
     GenerationOutputSpec,
     GenerationSamplingParams,
 )
+from nemo_rl.utils.length_penalty import apply_group_length_penalties
 from nemo_rl.utils.multimodal_payload_metrics import (
     collect_multimodal_payload_metrics,
     print_multimodal_payload_metrics,
@@ -2271,6 +2272,7 @@ async def run_async_nemo_gym_rollout(
     greedy: bool = False,
     effort_config: Optional[EffortLevelsConfig] = None,
     reward_penalty_config: dict[str, Any] | BaseModel | None = None,
+    length_penalty_config: dict[str, Any] | BaseModel | None = None,
     thinking_tags: list[str] | tuple[str, ...] | None = None,
     mask_env_flagged_samples: bool = True,
     returns_entire_batch: bool = False,
@@ -2302,6 +2304,7 @@ async def run_async_nemo_gym_rollout(
         greedy: Must be ``False`` because this path does not support greedy mode.
         effort_config: Optional configuration for effort-based reward shaping.
         reward_penalty_config: Optional reward-penalty configuration.
+        length_penalty_config: Optional GRPO config block for length adjustments.
         thinking_tags: Optional opening and closing tags used by thinking penalties.
         mask_env_flagged_samples: Whether to carry env-driven ``mask_sample``
             flags in the rollout batch for loss masking.
@@ -2477,6 +2480,7 @@ async def run_async_nemo_gym_rollout(
                         log_full_result_tables=log_full_result_tables,
                         effort_config=effort_config,
                         reward_penalty_config=reward_penalty_config,
+                        length_penalty_config=length_penalty_config,
                         thinking_tags=thinking_tags,
                         mask_env_flagged_samples=mask_env_flagged_samples,
                     )
@@ -2514,6 +2518,7 @@ def run_nemo_gym_rollout_sync(
     greedy: bool = False,
     effort_config: Optional[EffortLevelsConfig] = None,
     reward_penalty_config: dict[str, Any] | BaseModel | None = None,
+    length_penalty_config: dict[str, Any] | BaseModel | None = None,
     thinking_tags: list[str] | tuple[str, ...] | None = None,
     sampling_params: Optional[GenerationSamplingParams] = None,
     mask_env_flagged_samples: bool = True,
@@ -2578,6 +2583,7 @@ async def _consume_rollout() -> NemoGymRolloutResult:
             greedy=greedy,
             effort_config=effort_config,
             reward_penalty_config=reward_penalty_config,
+            length_penalty_config=length_penalty_config,
             thinking_tags=thinking_tags,
             mask_env_flagged_samples=mask_env_flagged_samples,
             returns_entire_batch=True,
@@ -2604,6 +2610,7 @@ def _postprocess_single_nemo_gym_group(
     log_full_result_tables: bool,
     effort_config: Optional[EffortLevelsConfig] = None,
     reward_penalty_config: dict[str, Any] | BaseModel | None = None,
+    length_penalty_config: dict[str, Any] | BaseModel | None = None,
     thinking_tags: list[str] | tuple[str, ...] | None = None,
     mask_env_flagged_samples: bool = True,
 ) -> NemoGymRolloutResult:
@@ -2620,6 +2627,27 @@ def _postprocess_single_nemo_gym_group(
     )
     penalty_counts = apply_reward_penalties(results, resolved_reward_penalty_config)
 
+    if length_penalty_config is not None:
+        grpo_config = (
+            length_penalty_config.model_dump()
+            if isinstance(length_penalty_config, BaseModel)
+            else dict(length_penalty_config)
+        )
+        # Callers pass the whole grpo config block; runs without a
+        # grpo.length_penalty section are untouched by this block.
+        if grpo_config.get("length_penalty"):
+            # Copy the per-row fields the length adjustments consume.
+            for nemo_gym_row, result in zip(nemo_gym_rows, results):
+                result["agent_ref"] = nemo_gym_row["agent_ref"]
+                result["profiled_rewards"] = nemo_gym_row.get("profiled_rewards")
+                result["profiled_output_lengths"] = nemo_gym_row.get(
+                    "profiled_output_lengths"
+                )
+                result["profile_band"] = nemo_gym_row.get("profile_band")
+            apply_group_length_penalties(
+                results, {"grpo": grpo_config}, tokenizer=tokenizer
+            )
+
     # Prepare for the rollout metrics calculation below. Not strictly necessary here, but good to have parity with `run_async_multi_turn_rollout`
     with timer.time(f"{timer_prefix}/prepare_for_metrics_calculation"):
         batch_size = len(nemo_gym_rows)