Add together_prime/: B200-capable PRIME harness

README.md

@@ -22,6 +22,11 @@
 
 </div>
 
+> **Running on NVIDIA Blackwell (B200)?** Use [`together_prime/`](./together_prime) — a
+> minimal, B200-validated harness (pinned upstream veRL `recipe/prime` + one recipe patch +
+> launch config). The legacy `training/` tree vendors an old verl that has no `sm_100`
+> kernels and will not run on Blackwell.
+
 
 # 🎉News
 

together_prime/.gitignore

@@ -0,0 +1,4 @@
+# Created by setup.sh (a fresh upstream verl checkout) — never commit it.
+/verl/
+# Data + model cache + checkpoints produced by prep.sh / smoke_run.sh.
+/data/

together_prime/README.md

@@ -0,0 +1,92 @@
+# Together Prime — B200-capable PRIME harness
+
+This directory is the **Blackwell (B200) capable** way to run the PRIME RL workload
+(Eurus-2-7B policy + implicit-PRM + vLLM rollout, RLOO). The legacy top-level `training/`
+tree vendors an old verl (`vllm<=0.6.3`) that has no `sm_100` kernels and cannot run on
+Blackwell — it is left untouched for historical reference. **Use this directory on B200.**
+
+PRIME has since been upstreamed into veRL as `recipe/prime`. Rather than fork or vendor
+verl, this harness pins an upstream verl commit, applies **one** small patch to the recipe,
+and drives it with launch config. All pins are in [`VERSIONS`](./VERSIONS).
+
+## The fix set (deliberately minimal)
+
+We validated a 3-step smoke test end-to-end on 8xB200, then ablated every workaround to
+keep only what is actually required. The result is small:
+
+| # | What | Kind | Why |
+|---|------|------|-----|
+| 1 | **recipe/prime port** (`patches/recipe-prime-port.patch`) | **code** | The upstream recipe had drifted from current verl. Fixes: `need_reference_policy(config)`; set `reward_fn`/`val_reward_fn` as attributes (dropped from `RayPPOTrainer.__init__`); rewrite `fit()` rollout to `_get_gen_batch` + `async_rollout_manager` (current verl tokenizes server-side from raw messages instead of popping pre-tokenized tensors). |
+| 2 | **`rollout.free_cache_engine=False`** | config | **The one Blackwell fix.** vLLM's sleep/wake `cumem` allocator (used to share GPU memory with the colocated FSDP actor) corrupts GPU memory on B200, surfacing as either a FlashInfer worker crash or a cuBLAS `CUBLAS_STATUS_EXECUTION_FAILED` / illegal-memory-access in the model GEMM. Disabling it fixes both. |
+| 3 | cu12.9 image `verlai/verl:vllm018.dev1` | config | The cu13.0 images (`vllm020`/`vllm023`) fault in the model cuBLAS GEMM on B200. |
+| 4 | `rollout.mode=async`, `rollout.enforce_eager=True` | config | `sync` rollout was removed upstream; eager avoids cudagraph capture during colocation. |
+| 5 | `resume_mode=disable`, `default_local_dir=/data/...` | config | Auto-resume trips on PyTorch 2.6 `weights_only=True` un-pickling a leftover dataloader checkpoint; and checkpoints shouldn't land inside the repo. |
+| 6 | HF offline (or set `HF_TOKEN`) | config | Repeated runs without a token get HF Hub 429s; `prep.sh` pre-caches the model, the run stays offline. |
+
+**verl core: 0 patches. vLLM: 0 patches.** The only code change is the recipe port (#1);
+everything else is a launch flag.
+
+### Workarounds and their removal triggers
+
+- **#2 `free_cache_engine=False`** is a real workaround for a vLLM-on-Blackwell memory bug,
+  not a preference — it costs the memory savings of sleep mode. Revisit when vLLM's cumem
+  sleep/wake is fixed for `sm_100`; a tracking issue should be filed and linked here.
+- **#3 image pin**: move to a newer cu12.x app image once one is validated on B200.
+- **#1 recipe port** is a genuine fix against current verl-recipe and is worth upstreaming
+  to `verl-project/verl-recipe`; once merged at the pinned commit, drop the patch.
+
+## Run it
+
+```bash
+# 0. one-time: build the pinned verl checkout + apply the recipe port
+bash together_prime/setup.sh                 # creates together_prime/verl
+
+# 1. prep data + model cache (needs network; downloads ~15GB once)
+bash together_prime/run/in_container.sh prep.sh
+
+# 2. smoke test (8xB200, 3 steps) — proves the pipeline runs
+bash together_prime/run/in_container.sh smoke_run.sh
+
+# 3. throughput capture (8xB200, 6 steps, realistic batch) — records tokens/sec
+bash together_prime/run/in_container.sh throughput_run.sh
+```
+
+Expect exit 0, `SMOKE_DONE`, three `step:N` lines with `actor/pg_loss`, `actor/grad_norm`,
+`actor/entropy`, `perf/mfu/actor` logged, and zero CUBLAS / illegal-memory / EngineDead
+errors. (`acc:0.0` is expected — the smoke config disables the accuracy filter and 3
+untrained steps on gsm8k under the reasoning-prompt format won't score; the goal is that
+the full pipeline runs.)
+
+Diagnostics: `run/gpu_check.sh` (bare GEMM sanity), `run/vllm_smoke.sh` (standalone vLLM,
+no verl) — both via `in_container.sh`.
+
+## Throughput baseline (8xB200, captured 2026-07-20)
+
+`run/throughput_run.sh` runs PRIME at a realistic shape (64 prompts × `rollout.n=4` =
+256 sequences, `max_response_length=3072`, 6 steps) and reaches steady state by step 3.
+Re-run 2026-07-20 reproduced the original 2026-07 capture within noise (step time
+53.9 → 53.1 s, all shares unchanged), confirming the baseline is stable.
+
+| Metric | Value |
+| --- | --- |
+| Step time (steady state) | ~53.1 s |
+| ↳ generation (vLLM rollout) | ~35.4 s (67%) |
+| ↳ actor update (FSDP) | ~11.9 s (22%) |
+| ↳ ref / old_log_prob / verify | ~5.7 s (11%) |
+| Generation throughput | ~21.6k tok/s total (~2.7k/GPU) |
+| End-to-end response throughput | ~14.4k tok/s (~1.8k/GPU) |
+| Sequences/sec | ~4.8 |
+| Actor MFU | ~19% |
+| Memory | ~33 GB alloc / ~55 GB reserved of 183 GB |
+
+The profile is healthy and rollout-dominated, as expected for this shape. It is **not
+optimized**: `enforce_eager=True` (our Blackwell-safety flag) disables the vLLM CUDA graph
+and leaves decode throughput on the table, and there is large memory headroom for a bigger
+batch.
+
+## Next
+
+- **Optimization levers:** test `enforce_eager=False` (risks the vLLM cumem/Blackwell
+  crash → needs its own validation) and a 2–3× larger batch to use the memory headroom.
+- **Reference bar:** no hard reference throughput is available yet; the numbers above are a
+  functional baseline, not a target-relative result.

together_prime/VERSIONS

@@ -0,0 +1,19 @@
+# Pinned versions for the B200-capable Together Prime harness.
+# These are the exact revisions the smoke test was validated against.
+
+# Container image (cu12.9 line; NOT the cu13.0 vllm020/vllm023 images — those fault in
+# the model cuBLAS GEMM on Blackwell).
+IMAGE=verlai/verl:vllm018.dev1
+
+# Upstream verl core (volcengine/verl == verl-project/verl).
+VERL_REPO=https://github.com/volcengine/verl.git
+VERL_COMMIT=bcb638649a50e58494a8ddd92085ad1174f674b8
+
+# recipe/prime submodule (verl-project/verl-recipe). The recipe-prime-port.patch is
+# diffed against this commit; setup.sh checks it out explicitly.
+RECIPE_REPO=https://github.com/verl-project/verl-recipe.git
+RECIPE_COMMIT=e7f889574b8301cc0f0fc1d57c6d67f31ffeb689
+
+# Base model (policy + implicit-PRM init) and eval dataset.
+MODEL=PRIME-RL/Eurus-2-7B-SFT
+DATA=openai/gsm8k

together_prime/patches/recipe-prime-port.patch

@@ -0,0 +1,83 @@
+diff --git a/prime/main_prime.py b/prime/main_prime.py
+index 39d20de..ad23395 100644
+--- a/prime/main_prime.py
++++ b/prime/main_prime.py
+@@ -117,7 +117,7 @@ def main_task(config, compute_score=None):
+     # TODO: Additional config checks can be added with proper function under prime recipe
+     validate_config(
+         config=config,
+-        use_reference_policy=need_reference_policy(role_worker_mapping),
++        use_reference_policy=need_reference_policy(config),
+         use_critic=False,
+     )
+
+diff --git a/prime/prime_ray_trainer.py b/prime/prime_ray_trainer.py
+index df53b0e..0372698 100644
+--- a/prime/prime_ray_trainer.py
++++ b/prime/prime_ray_trainer.py
+@@ -171,10 +171,12 @@ class RayPRIMETrainer(RayPPOTrainer):
+             role_worker_mapping,
+             resource_pool_manager,
+             ray_worker_group_cls,
+-            reward_fn=reward_fn,
+-            val_reward_fn=val_reward_fn,
+             device_name=device_name,
+         )
++        # verl core dropped reward_fn/val_reward_fn from RayPPOTrainer.__init__;
++        # PRIME's own fit()/verify path still uses these attributes, so set them here.
++        self.reward_fn = reward_fn
++        self.val_reward_fn = val_reward_fn
+
+         self.use_critic = False
+
+@@ -414,9 +416,18 @@ class RayPRIMETrainer(RayPPOTrainer):
+                 timing_raw = {}
+
+                 batch: DataProto = DataProto.from_single_dict(batch_dict)
++                batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature
+
+-                # pop those keys for generation
+-                gen_batch = batch.pop(batch_keys=["input_ids", "attention_mask", "position_ids"])
++                # add uid before splitting off the gen batch, so the reward keys
++                # (uid/data_source/reward_model/extra_info) are retained on `batch`.
++                batch.non_tensor_batch["uid"] = np.array(
++                    [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object
++                )
++
++                # current verl tokenizes server-side from raw messages; select the gen
++                # batch via the core helper instead of popping pre-tokenized tensors.
++                gen_batch = self._get_gen_batch(batch)
++                gen_batch.meta_info["global_steps"] = self.global_steps
+                 gen_batch_output = gen_batch.repeat(
+                     repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True
+                 )
+@@ -424,7 +435,7 @@ class RayPRIMETrainer(RayPPOTrainer):
+                 with simple_timer("step", timing_raw):
+                     # generate a batch
+                     with simple_timer("gen", timing_raw):
+-                        gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch_output)
++                        gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch_output)
+                         timing_raw.update(gen_batch_output.meta_info["timing"])
+                         gen_batch_output.meta_info.pop("timing", None)
+
+@@ -432,7 +443,7 @@ class RayPRIMETrainer(RayPPOTrainer):
+                         with simple_timer("gen_max", timing_raw):
+                             gen_baseline_batch = deepcopy(gen_batch)
+                             gen_baseline_batch.meta_info["do_sample"] = False
+-                            gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
++                            gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)
+
+                             batch = batch.union(gen_baseline_output)
+                             rm_scores, _ = self.compute_reward(batch, 1)
+@@ -453,10 +464,7 @@ class RayPRIMETrainer(RayPPOTrainer):
+
+                             del gen_baseline_batch, gen_baseline_output
+
+-                    batch.non_tensor_batch["uid"] = np.array(
+-                        [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object
+-                    )
+-                    # repeat to align with repeated responses in rollout
++                    # repeat to align with repeated responses in rollout (uid added above)
+                     batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+                     batch = batch.union(gen_batch_output)
+

together_prime/run/gpu_check.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+# Runs INSIDE the container. Minimal GPU sanity: device info + a bf16/fp16 GEMM on B200.
+# Useful to confirm the node/driver/cuBLAS are healthy independent of vLLM/verl.
+set -uxo pipefail
+python3 - <<'PY'
+import torch
+print("torch", torch.__version__, "cuda", torch.version.cuda)
+print("device", torch.cuda.get_device_name(0), "cap", torch.cuda.get_device_capability(0),
+      "count", torch.cuda.device_count())
+a = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+b = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
+torch.cuda.synchronize(); print("bf16 GEMM ok, sum=", float((a @ b).float().sum()))
+a16, b16 = a.half(), b.half()
+torch.cuda.synchronize(); print("fp16 GEMM ok, sum=", float((a16 @ b16).float().sum()))
+print("GPU_CHECK_OK")
+PY

together_prime/run/in_container.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Launch the pinned verl container with GPUs + persistent mounts, run the given script
+# inside it. Paths derive from this file's location so the harness is relocatable.
+#
+# Usage: bash run/in_container.sh prep.sh   |   bash run/in_container.sh smoke_run.sh
+# Env overrides: VERL_IMAGE, VERL_DIR (the setup.sh checkout), DATA_DIR.
+set -euo pipefail
+
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TOGETHER_PRIME="$(cd "$HERE/.." && pwd)"
+# shellcheck disable=SC1091
+source "$TOGETHER_PRIME/VERSIONS"
+
+VERL_IMAGE="${VERL_IMAGE:-$IMAGE}"
+VERL_DIR="${VERL_DIR:-$TOGETHER_PRIME/verl}"
+DATA_DIR="${DATA_DIR:-$TOGETHER_PRIME/data}"
+SCRIPT="${1:?usage: in_container.sh <script.sh>}"
+
+if [ ! -d "$VERL_DIR" ]; then
+    echo "verl checkout not found at $VERL_DIR — run: bash together_prime/setup.sh" >&2
+    exit 1
+fi
+mkdir -p "$DATA_DIR/hf_cache"
+
+docker run --rm --gpus all \
+    --ipc=host --shm-size=32g \
+    --ulimit memlock=-1 --ulimit stack=67108864 \
+    -e HF_HUB_ENABLE_HF_TRANSFER=1 \
+    -e HF_HUB_OFFLINE=1 \
+    -e TRANSFORMERS_OFFLINE=1 \
+    -v "$VERL_DIR:/workspace/verl" \
+    -v "$DATA_DIR:/data" \
+    -v "$HERE:/scripts" \
+    -w /workspace/verl \
+    "$VERL_IMAGE" \
+    bash "/scripts/$SCRIPT"

together_prime/run/prep.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Runs INSIDE the container. Prepares gsm8k parquet + pre-downloads the base model so the
+# smoke run can stay offline (avoids HF Hub 429s on repeated runs without an HF_TOKEN).
+set -euxo pipefail
+
+export HF_HOME=/data/hf_cache
+# prep must reach the HF Hub to download; override the container-level offline flags.
+export HF_HUB_OFFLINE=0
+export TRANSFORMERS_OFFLINE=0
+export PYTHONPATH=/workspace/verl:${PYTHONPATH:-}
+mkdir -p /data/gsm8k
+
+# gsm8k -> train.parquet / test.parquet (data_source="openai/gsm8k")
+python3 /workspace/verl/examples/data_preprocess/gsm8k.py --local_save_dir /data/gsm8k
+
+# pre-fetch policy / implicit-PRM init model (public, no token needed)
+python3 - <<'PY'
+from huggingface_hub import snapshot_download
+print("model at", snapshot_download("PRIME-RL/Eurus-2-7B-SFT"))
+PY
+
+ls -la /data/gsm8k
+echo "PREP_DONE"

together_prime/run/smoke_run.sh

@@ -0,0 +1,74 @@
+#!/bin/bash
+# Runs INSIDE the container. Scaled-down PRIME smoke test on 8xB200: 3 steps, gsm8k only,
+# no wandb. Validates the full loop (rollout -> verify -> implicit-PRM update -> RLOO ->
+# FSDP policy update) with Eurus-2-7B-SFT.
+#
+# The ONLY Blackwell-specific fix needed is actor_rollout_ref.rollout.free_cache_engine=False
+# (see below). Everything else here is ordinary launch config.
+set -euxo pipefail
+
+export HF_HOME=/data/hf_cache
+export PYTHONPATH=/workspace/verl:${PYTHONPATH:-}
+export TOKENIZERS_PARALLELISM=true
+export NCCL_DEBUG=WARN
+# NOTE: do NOT set PYTORCH_CUDA_ALLOC_CONF=expandable_segments — vLLM's sleep-mode memory
+# pool asserts against it (pytorch/pytorch#147851). We disable sleep mode anyway (below).
+
+cd /workspace/verl
+
+train_path=/data/gsm8k/train.parquet
+test_path=/data/gsm8k/test.parquet
+model_path=PRIME-RL/Eurus-2-7B-SFT
+
+python3 -m recipe.prime.main_prime \
+    data.train_files="['$train_path']" \
+    data.val_files="['$test_path']" \
+    data.train_batch_size=16 \
+    data.val_batch_size=64 \
+    data.max_prompt_length=1024 \
+    data.max_response_length=1024 \
+    data.filter_overlong_prompts=True \
+    data.filter_accuracy=False \
+    data.oversample_factor=1 \
+    actor_rollout_ref.model.path=$model_path \
+    actor_rollout_ref.actor.optim.lr=5e-7 \
+    actor_rollout_ref.model.use_remove_padding=True \
+    actor_rollout_ref.actor.ppo_mini_batch_size=16 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    actor_rollout_ref.actor.fsdp_config.param_offload=True \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+    actor_rollout_ref.actor.use_kl_loss=False \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.mode=async \
+    actor_rollout_ref.rollout.enforce_eager=True \
+    actor_rollout_ref.rollout.free_cache_engine=False \
+    actor_rollout_ref.rollout.n=4 \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
+    algorithm.adv_estimator=rloo \
+    algorithm.use_kl_in_reward=True \
+    algorithm.kl_penalty=kl \
+    algorithm.kl_ctrl.kl_coef=0.001 \
+    reward_model.model.path=$model_path \
+    reward_model.model.update=before \
+    reward_model.model.beta_train=0.05 \
+    reward_model.model.optim.lr=1e-6 \
+    reward_model.model.optim.grad_clip=10.0 \
+    reward_model.mini_batch_size=16 \
+    trainer.val_before_train=False \
+    trainer.logger='["console"]' \
+    trainer.project_name='prime_smoke' \
+    trainer.experiment_name='eurus2-7b-b200-smoke' \
+    trainer.n_gpus_per_node=8 \
+    trainer.nnodes=1 \
+    trainer.default_local_dir=/data/ckpts/prime_smoke \
+    trainer.resume_mode=disable \
+    trainer.save_freq=999 \
+    trainer.test_freq=999 \
+    trainer.total_training_steps=3 \
+    trainer.total_epochs=1 "$@"
+
+echo "SMOKE_DONE"