From 082f0239070d5d4a771213a0251e023272d8da2d Mon Sep 17 00:00:00 2001 From: Rockdu Date: Sun, 6 Sep 2026 11:41:15 -0700 Subject: [PATCH] revert: feat(reward): per-reward colocate, mixture RM, verified SD3.5 HPS recipe (#234) --- docker/Dockerfile | 5 - docs/advanced/streaming-reward.md | 6 +- docs/models/sd3/sd3.md | 71 ++------- docs/user-guide/cli-reference.md | 16 +- docs/user-guide/customization.md | 14 +- docs/user-guide/launch-script.md | 2 +- docs/user-guide/recipe-verification.md | 1 - docs/user-guide/rewards.md | 92 +---------- miles/ray/placement_group.py | 3 +- miles/ray/rollout.py | 33 +--- miles/ray/utils.py | 5 - miles/rollout/rm_hub/__init__.py | 28 +--- miles/rollout/rm_hub/core.py | 127 ++++----------- miles/rollout/rm_hub/hps.py | 150 ------------------ miles/rollout/rm_hub/ocr.py | 61 ++++--- miles/rollout/rm_hub/pickscore.py | 85 ++++++---- miles/rollout/rm_hub/weighted_mixture_rm.py | 51 ------ miles/rollout/sft_rollout.py | 35 ++-- miles/utils/arguments.py | 93 ++--------- miles/utils/processing_utils.py | 26 +-- miles/utils/types.py | 3 +- requirements.txt | 1 - ...ffusion_grpo_cosmos3_pickscore_t2i_4gpu.py | 4 +- scripts/run_diffusion_grpo_h3_t2va_2gpu.py | 2 +- scripts/run_diffusion_grpo_sd3_hps_sglang.py | 130 --------------- ...on_grpo_wan22_pickscore_17gpu_multinode.py | 2 +- tests/fast-gpu/test_hpsv2_alignment.py | 80 ---------- tests/fast/rollout/test_hps_reward.py | 122 -------------- tests/fast/rollout/test_reward_pool.py | 124 --------------- .../rollout/test_reward_pool_placement.py | 74 --------- .../fast/rollout/test_weighted_mixture_rm.py | 63 -------- tests/fast/utils/test_processing_utils.py | 25 +-- 32 files changed, 199 insertions(+), 1335 deletions(-) delete mode 100644 miles/rollout/rm_hub/hps.py delete mode 100644 miles/rollout/rm_hub/weighted_mixture_rm.py delete mode 100644 scripts/run_diffusion_grpo_sd3_hps_sglang.py delete mode 100644 tests/fast-gpu/test_hpsv2_alignment.py delete mode 100644 tests/fast/rollout/test_hps_reward.py delete mode 100644 tests/fast/rollout/test_reward_pool.py delete mode 100644 tests/fast/rollout/test_reward_pool_placement.py delete mode 100644 tests/fast/rollout/test_weighted_mixture_rm.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 06cdec94c..4846266b9 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -65,11 +65,6 @@ COPY requirements.txt /tmp/requirements.txt RUN pip install --ignore-installed PyJWT RUN pip install -r /tmp/requirements.txt -# Official HPSv2 reference for the alignment test; its protobuf<4 pin conflicts with wandb. -# Pinned to git: the PyPI 1.2.0 wheel omits the BPE vocab, so `import hpsv2.src.open_clip` fails. -RUN pip install --no-deps args==0.1.0 clint==0.5.1 "hpsv2 @ git+https://github.com/tgxs002/HPSv2@866735ecaae9" && \ - python3 -c "import hpsv2.src.open_clip" - # ====================================== Install main package ============================================ # Install miles_diffusion from a pinned ref (like radixark/miles does for miles), diff --git a/docs/advanced/streaming-reward.md b/docs/advanced/streaming-reward.md index cab663b30..fbda01220 100644 --- a/docs/advanced/streaming-reward.md +++ b/docs/advanced/streaming-reward.md @@ -76,7 +76,7 @@ seen over the rollout: | Metric | Meaning | |---|---| | `perf/parser_max_queue_depth` | Deepest backlog on a parser actor when a response was handed to it | -| `perf/reward_max_queue_depth_` | Deepest backlog on that reward pool's actors when a batch was handed to them, one metric per pool (`hps`, `pickscore`, `ocr`) | +| `perf/reward_max_queue_depth` | Deepest backlog on a reward actor when a batch was handed to it | `0` means every dispatch found its worker idle β€” that stage never made anything wait. A number that climbs with the number of concurrent microgroups means the stage is saturated and requests are lining up behind it. @@ -86,9 +86,9 @@ Read them together with `perf/rollout_time`: | Symptom | Look at | Likely fix | |---|---|---| | `perf/rollout_time` high | `perf/parser_max_queue_depth` > 0 | Deserialization is the bottleneck; raise `--rollout-parser-num-workers` | -| `perf/rollout_time` high | `perf/reward_max_queue_depth_` > 0 | That pool's scoring is the bottleneck; add its workers, or give it a dedicated GPU | +| `perf/rollout_time` high | `perf/reward_max_queue_depth` > 0 | Scoring is the bottleneck; add reward workers, or give them a dedicated GPU | | `perf/rollout_time` high, both depths `0` | Neither pool is holding anything up | The engines themselves are the limit β€” check `--rollout-microgroup-size` and `--sglang-server-concurrency` | Both metrics are emitted only when the corresponding pool ran, so a reward with no actor pool leaves -its `perf/reward_max_queue_depth_` absent rather than zero. +`perf/reward_max_queue_depth` absent rather than zero. diff --git a/docs/models/sd3/sd3.md b/docs/models/sd3/sd3.md index cb0f51272..b678b904d 100644 --- a/docs/models/sd3/sd3.md +++ b/docs/models/sd3/sd3.md @@ -13,7 +13,7 @@ DiT transformer with dual text-encoder conditioning (`encoder_hidden_states` + - Single DiT component β€” weight sync targets `--update-weight-target-module transformer` (default). - LoRA on all attention projections (self-attn + cross-attn add projections). -- Supports **Flow-GRPO** (OCR or HPS) and **DiffusionNFT** (PickScore) objectives. +- Supports **Flow-GRPO** (OCR) and **DiffusionNFT** (PickScore) objectives. - Gated Hugging Face model β€” requires `HF_TOKEN`. ## 2. Supported variants @@ -51,7 +51,6 @@ Prompt datasets live under | Recipe | Subset | Train path | |---|---|---| | GRPO + OCR | `flowgrpo_ocr` | `.../flowgrpo_ocr/train.jsonl` | -| GRPO + HPS | `hpdv2` | `.../hpdv2/train.jsonl` | | NFT + PickScore | `flowgrpo_pickscore` | `.../flowgrpo_pickscore/train.jsonl` | Launch scripts download the matching subset automatically via @@ -103,7 +102,6 @@ All recipes are Python modules under `scripts/`. Each exposes a Typer CLI | Script | Reward | GPUs | Algorithm | |---|---|---|---| | `run_diffusion_grpo_sd3_ocr_sglang.py` | OCR (CPU) | 2 colocate | Flow-GRPO | -| `run_diffusion_grpo_sd3_hps_sglang.py` | HPS | 2 colocate | Flow-GRPO | | `run_diffusion_nft_sd3_pickscore.py` | PickScore | 3 (2+1) | DiffusionNFT | ### 5.2 Flow-GRPO + OCR (2 GPU colocate) @@ -122,27 +120,7 @@ Walkthrough: [Quick Start](../../getting-started/quick-start.md). E2E test: `tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py`. -### 5.3 Flow-GRPO + HPS (2 GPU colocate) - -Canonical script: `scripts/run_diffusion_grpo_sd3_hps_sglang.py` - -**Status:** [πŸ“ˆ V β€” Verified](../../user-guide/recipe-verification.md#v) β€” 600 rollouts -(2 optimizer steps each, 1,200 in total) on 2Γ—H200; `rollout/reward/raw_mean` 0.284 β†’ 0.349 -(mean of the last 100 rollouts, peak 10-rollout moving average 0.362). - -```bash -export HF_TOKEN=... -python3 scripts/run_diffusion_grpo_sd3_hps_sglang.py \ - --cuda-visible-devices 6,7 -``` - -The recipe uses the SD3 Flow-GRPO SDE, LoRA and precision configuration of the OCR -recipe, swaps in the deduplicated `hpdv2` prompts and `--rm-type hps`, and colocates one -HPS reward actor with the train and rollout workers. It keeps Flow-GRPO's own KL weight -(`--diffusion-kl-beta 0.01`) and group-wise advantage std instead of the OCR recipe's -`--diffusion-kl-beta 0.04 --globalize-reward-std`. - -### 5.4 DiffusionNFT + PickScore (3 GPU) +### 5.3 DiffusionNFT + PickScore (3 GPU) Script: `scripts/run_diffusion_nft_sd3_pickscore.py` @@ -162,14 +140,14 @@ MILES_SCRIPT_SMOKE=1 python3 scripts/run_diffusion_nft_sd3_pickscore.py ### Recipe comparison -| | GRPO + OCR | GRPO + HPS | NFT + PickScore | -|---|---|---|---| -| Script | `run_diffusion_grpo_sd3_ocr_sglang.py` | `run_diffusion_grpo_sd3_hps_sglang.py` | `run_diffusion_nft_sd3_pickscore.py` | -| `--loss-type` | `policy_loss` (default) | `policy_loss` (default) | `nft` | -| SDE | Full window, noise=0.7, CFG=4.5 | Full window, noise=0.7, CFG=4.5 | ODE, noise=0 | -| Reference | LoRA base KL (Ξ² 0.04) | LoRA base KL (Ξ² 0.01) | EMA (`--use-ema`) | -| Reward placement | CPU OCR | Colocated HPS actor | Dedicated PickScore GPU | -| Verification | FG | V | FG | +| | GRPO + OCR | NFT + PickScore | +|---|---|---| +| Script | `run_diffusion_grpo_sd3_ocr_sglang.py` | `run_diffusion_nft_sd3_pickscore.py` | +| `--loss-type` | `policy_loss` (default) | `nft` | +| SDE | Full window, noise=0.7, CFG=4.5 | ODE, noise=0 | +| Reference | LoRA base KL | EMA (`--use-ema`) | +| Reward GPU | None (CPU OCR) | Dedicated (3 GPU total) | +| Deterministic e2e | `test_sd3_ocr_grpo_2xGPU` | `test_sd3_nft_pickscore_3xGPU` | ## 6. Recipe configuration @@ -177,9 +155,6 @@ MILES_SCRIPT_SMOKE=1 python3 scripts/run_diffusion_nft_sd3_pickscore.py **GRPO + OCR (default script):** 2 GPUs colocated (`--colocate`); OCR on CPU Ray actors. -**GRPO + HPS:** 2 GPUs colocated (`--colocate --hps-reward-colocate`); one HPS -actor shares the rollout GPUs. - **NFT + PickScore:** | GPU role | Count | Flags | @@ -188,9 +163,9 @@ actor shares the rollout GPUs. | PickScore reward | 1 | `--pickscore-num-workers 1`, `--pickscore-num-gpus-per-worker 1.0` | | **Total** | **3** | `--num-gpus-per-node 3` | -PickScore runs as a Ray actor pool on a dedicated GPU. With `--pickscore-reward-colocate` -(not used in the default script), reward workers share the rollout GPUs -instead β€” useful only when GPU count is tight. +PickScore runs as a Ray actor pool on a dedicated GPU. With `--colocate-reward` +(not used in the default script), reward workers share rollout GPUs at 0.05 GPU +per worker β€” useful only when GPU count is tight. ### Algorithm flags @@ -206,18 +181,6 @@ instead β€” useful only when GPU count is tight. | Weight sync | `--lora-ipc-weight-sync` (colocate IPC merge) | | Determinism | `--deterministic-mode` (CI / e2e parity) | -**Flow-GRPO + HPS:** uses the same SDE, LoRA and precision flags as the OCR recipe, -with these settings of its own: - -| Setting | Value | -|---|---| -| Reference | `--diffusion-kl-beta 0.01` | -| Advantage | Group-wise std (no `--globalize-reward-std`) | -| Reward | `--rm-type hps --hps-version v2.1` | -| Reward worker | One actor, batch size 8 | -| Placement | `--hps-reward-colocate` | -| Prompt subset | `hpdv2` | - **DiffusionNFT + PickScore:** | Setting | Value | @@ -231,7 +194,7 @@ with these settings of its own: ## 7. LoRA and weight sync -All SD3 recipes use LoRA with IPC weight sync: +Both SD3 recipes use LoRA with IPC weight sync: ```bash --use-lora \ @@ -247,7 +210,7 @@ merge internals. ## 8. Precision notes -All SD3 launch scripts use fp16 DiT forward: +Both SD3 launch scripts use fp16 DiT forward: ```bash --diffusion-forward-dtype fp16 \ @@ -290,7 +253,7 @@ by the E2E fixture. ## 10. Pairs well with - [Quick Start](../../getting-started/quick-start.md) β€” SD3.5 Flow-GRPO OCR walkthrough. -- [Rewards](../../user-guide/rewards.md) β€” OCR, HPS and PickScore scoring. +- [Rewards](../../user-guide/rewards.md) β€” OCR and PickScore scoring. - [Customization](../../user-guide/customization.md) β€” `--*-path` plug-points. - [SDE step backend](../../advanced/sde-backend.md) β€” SDE window (GRPO) vs ODE (NFT). -- [LoRA weight sync](../../advanced/lora.md) β€” IPC merge used by all SD3 recipes. +- [LoRA weight sync](../../advanced/lora.md) β€” IPC merge used by both recipes. diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index c1476d82a..fb1433d53 100644 --- a/docs/user-guide/cli-reference.md +++ b/docs/user-guide/cli-reference.md @@ -265,26 +265,18 @@ See [Dtype Control](../advanced/dtype-control.md). | Flag | Type | Default | Notes | |---|---|---|---| -| `--rm-type` | enum | – | `pickscore` / `hps` / `ocr`. Overridable per sample via `metadata.rm_type`. Ignored when `--custom-rm-path` is set. | +| `--rm-type` | enum | – | `pickscore` / `ocr`. Overridable per sample via `metadata.rm_type`. | | `--reward-key` | str | – | When the reward is a dict. | | `--group-rm` | flag | off | Score a whole prompt group at once. | -| `--custom-rm-path` | str | – | `async def rm(args, samples) -> list[float]`. Batched only; replaces the `--rm-type` dispatch entirely. Shipped: `miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm` (weighted sum of built-in rewards). | -| `--custom-rm-args` | str | – | Opaque config string for the custom RM, read as `args.custom_rm_args`; e.g. `"hps=0.7,pickscore=0.3"` for `rm_hub.weighted_mixture_rm`. | -| `--reward-key` | str | – | For dict-valued rewards: the entry GRPO trains on. Every entry is also logged as `rollout/reward/_mean` and `eval//`. | +| `--custom-rm-path` | str | – | `async def rm(args, samples) -> list[float]`. Batched only. | | `--custom-reward-post-process-path` | str | – | Replace advantage normalisation. | +| `--colocate-reward` | flag | off | Reward actors onto rollout GPUs (train 0.7 + rollout 0.25 + reward 0.05). Requires `--colocate`. | | `--pickscore-model-path` | str | – | Required for `--rm-type pickscore`. | | `--pickscore-processor-path` | str | – | Required for `--rm-type pickscore`. | | `--pickscore-num-workers` | int | `1` | | -| `--pickscore-num-gpus-per-worker` | float | `1.0` | Standalone workers only; fractional values allowed. | -| `--pickscore-reward-colocate` | flag | off | One worker per rollout GPU, sharing it with the train actor and rollout engine. Requires `--colocate`. | +| `--pickscore-num-gpus-per-worker` | float | `1.0` | Fractional values allowed. | | `--pickscore-batch-size` | int | `8` | | | `--pickscore-num-frames` | int | – | Frames scored per video; unset = all. | -| `--hps-num-workers` | int | `1` | | -| `--hps-num-gpus-per-worker` | float | `1.0` | Standalone workers only; fractional values allowed. | -| `--hps-reward-colocate` | flag | off | One worker per rollout GPU, sharing it with the train actor and rollout engine. Requires `--colocate`. | -| `--hps-batch-size` | int | `8` | | -| `--hps-version` | enum | `v2.1` | `v2.0` / `v2.1`. | -| `--hps-checkpoint-path` | str | – | Local checkpoint; unset downloads from Hugging Face. | | `--ocr-num-workers` | int | `4` | | | `--rollout-parser-num-workers` | int | `1` | Ray actors deserializing rollout responses. Raise when trajectory tensors are large. | diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md index 81da27572..c40e07dca 100644 --- a/docs/user-guide/customization.md +++ b/docs/user-guide/customization.md @@ -126,25 +126,15 @@ async def custom_rm(args, samples: list[Sample], **kwargs) -> list[float]: ``` Wired only through `batched_async_rm` β€” implement per-sample routing inside your -batched function if needed. Once set, `--rm-type` and `metadata.rm_type` are ignored: the -custom function is the whole dispatch. +batched function if needed. ```bash --custom-rm-path my_project.rewards.aesthetic_rm ``` -`--custom-rm-args` is an opaque string passed through as `args.custom_rm_args` for the -custom RM to parse. - -Shipped custom RMs: - -| Path | What | -|---|---| -| `miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm` | Weighted sum of built-in rewards (`hps`, `pickscore`, `ocr`), weights from `--custom-rm-args "hps=0.7,pickscore=0.3"`; returns a dict per sample, train on it with `--reward-key weighted`. See [Rewards](rewards.md) Β§ Combining rewards. | - HTTP / remote scoring: implement a batched custom RM and read `args.rm_url` (or your own flags). Encode images from `sample.generated_output` (see -`generated_output_to_rgb_hwc_uint8_frames` in `miles/utils/processing_utils.py`): +`_sample_to_rgb_hwc_uint8_frames` in `miles/rollout/rm_hub/pickscore.py`): ```python import aiohttp diff --git a/docs/user-guide/launch-script.md b/docs/user-guide/launch-script.md index 3770130e6..7724a6d14 100644 --- a/docs/user-guide/launch-script.md +++ b/docs/user-guide/launch-script.md @@ -237,7 +237,7 @@ MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_grpo_wan22_pickscore_1 ``` Reward workers (`--pickscore-num-workers 4 --pickscore-num-gpus-per-worker 0.25`, no -`--pickscore-reward-colocate`) are default-scheduled and land on the only free GPU: the reward node. +`--colocate-reward`) are default-scheduled and land on the only free GPU: the reward node. ### Verify the run is healthy diff --git a/docs/user-guide/recipe-verification.md b/docs/user-guide/recipe-verification.md index cec7d8e5e..a25bd4cc4 100644 --- a/docs/user-guide/recipe-verification.md +++ b/docs/user-guide/recipe-verification.md @@ -48,7 +48,6 @@ count as verification. full-finetune Flow-GRPO + PickScore. - **πŸ“ˆ V** - `run_diffusion_sft_h3_t2va.py` β€” MiniMax H3 8-GPU LoRA SFT. - - `run_diffusion_grpo_sd3_hps_sglang.py` β€” SD3.5 Flow-GRPO + HPSv2.1. - **β—‹ NV** - `run_diffusion_grpo_wan22_pickscore_5gpu.py` β€” Wan2.2 5-GPU LoRA Flow-GRPO + PickScore. diff --git a/docs/user-guide/rewards.md b/docs/user-guide/rewards.md index 55e0b5ac6..561be08a3 100644 --- a/docs/user-guide/rewards.md +++ b/docs/user-guide/rewards.md @@ -1,6 +1,6 @@ --- title: Rewards -description: Built-in reward models (PickScore, HPS, OCR), rm_hub dispatch, and prompt data format. +description: Built-in reward models (PickScore, OCR), rm_hub dispatch, and prompt data format. --- Miles-diffusion scores generated images (or video frames) after each rollout microgroup. Reward computation lives in `miles/rollout/rm_hub/` and is invoked @@ -13,7 +13,7 @@ For `--custom-rm-path`, `--custom-reward-post-process-path`, and other | Stage | Flag | Role | |---|---|---| -| Reward type | `--rm-type` | Selects built-in scorer (`pickscore`, `hps`, `ocr`); ignored when `--custom-rm-path` is set | +| Reward type | `--rm-type` | Selects built-in scorer (`pickscore`, `ocr`) | | Per-sample override | `metadata.rm_type` in JSONL | Overrides global `--rm-type` | | Custom reward / norm | see [Customization](customization.md) | `--custom-rm-path`, `--custom-reward-post-process-path` | @@ -49,7 +49,7 @@ batching. For video outputs, frames are uniformly sampled | `--pickscore-processor-path` | β€” | Required for pickscore | | `--pickscore-model-path` | β€” | Required for pickscore | | `--pickscore-num-frames` | None | Video frame sampling count | -| `--pickscore-reward-colocate` | False | One worker per rollout GPU (requires `--colocate`) | +| `--colocate-reward` | False | Share rollout GPUs (0.05 GPU/worker) | Example from `scripts/run_diffusion_nft_sd3_pickscore.py`: @@ -62,85 +62,6 @@ Example from `scripts/run_diffusion_nft_sd3_pickscore.py`: --pickscore-model-path yuvalkirstain/PickScore_v1 ``` -### HPS (`--rm-type hps`) - -Implementation: `miles/rollout/rm_hub/hps.py`. - -HPSv2 scores text–image alignment using a preference-tuned CLIP model: - -- Model: `ViT-H-14` -- Version: `--hps-version` (`v2.0` or `v2.1`) -- Checkpoint: `--hps-checkpoint-path` (optional; defaults to `xswu/HPSv2`) - -Scoring formula: - -``` -score = diagonal(image_features @ text_features.T) -``` - -The diagonal pairs each image with its corresponding prompt in the batch. - -HPS runs as a **Ray actor pool** (`HPSRewardActor`) with round-robin batching. -It currently accepts image outputs only (`generated_output` with one frame). -Use a custom RM when defining video frame aggregation semantics. - -| Flag | Default | Description | -|---|---|---| -| `--hps-num-workers` | 1 | Ray actor count | -| `--hps-num-gpus-per-worker` | 1.0 | GPU per worker (non-colocate) | -| `--hps-batch-size` | 8 | Batch size per actor | -| `--hps-version` | `v2.1` | `v2.0` or `v2.1` checkpoint | -| `--hps-checkpoint-path` | None | Local checkpoint; unset downloads from `xswu/HPSv2` | -| `--hps-reward-colocate` | False | One worker per rollout GPU (requires `--colocate`) | - -Example from `scripts/run_diffusion_grpo_sd3_hps_sglang.py`: - -```bash ---rm-type hps \ ---hps-num-workers 1 \ ---hps-batch-size 8 \ ---hps-version v2.1 \ ---hps-reward-colocate -``` - -### Reward placement - -Every GPU reward pool is placed one of two ways: - -- `---reward-colocate`: every worker takes one **slot** on a rollout placement-group bundle, - sharing that GPU with the train actor and the rollout engine. Colocated pools share one slot ledger, so `hps` and `pickscore` can both colocate without - overlapping; more workers than bundles is rejected at parse time. Requires `--colocate`. -- Otherwise the pool is **standalone**: default-scheduled at `---num-gpus-per-worker`, - which only lands on GPUs outside every placement group (placement groups reserve their - GPUs, so Ray never packs these onto rollout GPUs). - -`RolloutManager` seats the colocated pools before the first rollout; standalone pools are -built on first use. - -### Combining rewards - -`--custom-rm-path` receives `(args, samples)` and can call the built-in scorers -directly; `--custom-rm-args` is an opaque string the framework hands to that function -through `args`, so the function owns its own config grammar. The shipped example -`miles/rollout/rm_hub/weighted_mixture_rm.py` reads `name=weight,name=weight`: - -```bash ---custom-rm-path miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm \ ---custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weighted \ ---pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K \ ---pickscore-model-path yuvalkirstain/PickScore_v1 \ ---hps-reward-colocate --pickscore-reward-colocate # each reward keeps its own placement flags -``` - -The example returns a dict per sample (`{"hps": ..., "pickscore": ..., "weighted": ...}`). -`--reward-key` picks the entry GRPO trains on, and every entry of a dict reward gets its own -`rollout/reward/_mean` panel, so the components stay visible while the sum is optimized. - -Weights apply to raw scores (HPSv2.1 β‰ˆ 0.25–0.35, PickScore/26 β‰ˆ 0.8–0.9, OCR ∈ [0, 1]), so -pick them with the scales in mind. Colocated pools share one slot ledger, so several rewards -can colocate without overlapping. Rewards receive `generated_output` itself, and every reward actor -quantises it to uint8 on its own terms. - ### OCR (`--rm-type ocr`) Implementation: `miles/rollout/rm_hub/ocr.py`. @@ -174,9 +95,7 @@ generate_and_rm_microgroup() β†’ batched_async_rm(args, microgroup) # rm_hub/__init__.py β†’ custom_rm_path? user batched function β†’ all pickscore? pickscore_rm (batched) - β†’ all hps? hps_rm (batched) - β†’ all ocr? ocr_rm (batched, one image per actor call) - β†’ else per-sample async_rm β†’ ocr / pickscore / hps / NotImplementedError + β†’ else per-sample async_rm β†’ ocr / pickscore / NotImplementedError β†’ sample.reward = score β†’ RolloutManager._post_process_rewards() # GRPO advantage normalization ``` @@ -208,7 +127,6 @@ Dataset repo: [`rockdu/miles-diffusion-datasets`](https://huggingface.co/dataset | Subset | Used by | |---|---| -| `hpdv2/` | SD3 HPS Flow-GRPO | | `flowgrpo_pickscore/` | PickScore recipes (SD3 NFT, Qwen-Image, Wan2.2, LTX) | | `flowgrpo_ocr/` | SD3 OCR Flow-GRPO, NFT smoke test | @@ -222,4 +140,4 @@ metadata.get("rm_type") or args.rm_type ``` Mixed rm_types within one microgroup fall back to per-sample dispatch (no -batched PickScore/HPS fast path). +batched PickScore fast path). diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index b87de71d7..82bae912e 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -7,7 +7,6 @@ from .actor_group import RayTrainGroup from .rollout import RolloutManager -from .utils import TRAIN_ACTOR_GPU logger = logging.getLogger(__name__) @@ -125,7 +124,7 @@ def allocate_train_group(args, num_nodes, num_gpus_per_node, pg): num_nodes=num_nodes, num_gpus_per_node=num_gpus_per_node, pg=pg, - num_gpus_per_actor=TRAIN_ACTOR_GPU, + num_gpus_per_actor=0.7, ) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 2da1f07a4..cef95d3f0 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -18,9 +18,7 @@ expand_samples_to_train_pairs as flow_grpo_expand_samples_to_train_pairs, ) from miles.rollout.base_types import call_rollout_fn -from miles.rollout.rm_hub import create_colocated_reward_pools -from miles.rollout.rm_hub.core import ColocatedRewardSlots, bundle_deal_order -from miles.rollout.sft_rollout import SftEncodePool +from miles.rollout.rm_hub.core import set_manager_placement_group from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor from miles.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client @@ -36,7 +34,7 @@ from miles.utils.train_metric_utils import log_perf_data_raw from miles.utils.types import Sample -from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, ROLLOUT_ENGINE_GPU, Lock +from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("httpcore").setLevel(logging.WARNING) @@ -57,6 +55,7 @@ def __init__(self, args, pg): from miles.dashboard import hooks hooks.register_rollout_manager(args) + set_manager_placement_group(pg) if not args.train_only: logger.info("RolloutManager: starting router...") _start_router(args) @@ -104,14 +103,6 @@ def __init__(self, args, pg): self.all_rollout_engines = [None] * num_engines self.num_new_engines = init_rollout_engines(args, pg, self.all_rollout_engines) logger.info("RolloutManager started %s rollout engines", len(self.all_rollout_engines)) - # pools outlive every rollout call, so the manager seats them and the rm functions find them as singletons - _, bundle_indices, gpu_ids = pg - self.reward_slots = ColocatedRewardSlots( - bundle_deal_order(bundle_indices, gpu_ids, args.num_gpus_per_node, args.rollout_num_gpus_per_engine) - ) - self.reward_pools = create_colocated_reward_pools(args, pg, self.reward_slots) - logger.info("RolloutManager colocated reward slots: %s", self.reward_slots) - self.encode_pool = SftEncodePool(args, pg) if args.loss_type == "sft_loss" else None logger.info("RolloutManager: creating lock...") self.nodes_per_engine = max(1, args.rollout_num_gpus_per_engine // args.num_gpus_per_node) self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() @@ -397,11 +388,6 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl reward_stats["rollout/reward/group_mean_avg"] = float(groups_raw.mean(dim=-1).mean()) if groups_raw.shape[-1] > 1: reward_stats["rollout/reward/group_std_avg"] = float(groups_raw.std(dim=-1, unbiased=False).mean()) - # a dict reward (--reward-key) carries components; a per-sample custom RM may give each sample different ones - if self.args.reward_key: - for key in sorted({key for sample in samples for key in sample.reward}): - values = [sample.reward[key] for sample in samples if key in sample.reward] - reward_stats[f"rollout/reward/{key}_mean"] = float(np.mean(values)) print( f"[reward stats] raw mean={raw_t.mean():.4f} std={raw_t.std():.4f} min={raw_t.min():.4f} max={raw_t.max():.4f} | " @@ -507,7 +493,8 @@ def init_rollout_engines(args, pg, all_rollout_engines): if all_rollout_engines[i] is not None: continue - num_gpus = ROLLOUT_ENGINE_GPU + # Leave 0.05 for a colocated reward actor when --colocate-reward is set. + num_gpus = 0.25 if args.colocate_reward else 0.3 num_cpus = num_gpus # Get the base GPU ID from placement group @@ -672,10 +659,6 @@ def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] log_dict[f"eval/{key}"] = sum(rewards) / len(rewards) if (samples := data[key].get("samples")) is not None: log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), f"eval/{key}/") - if args.eval_reward_key: - for name in sorted({name for sample in samples for name in sample.reward}): - values = [sample.reward[name] for sample in samples if name in sample.reward] - log_dict[f"eval/{key}/{name}"] = sum(values) / len(values) logger.info(f"eval {rollout_id}: {log_dict}") step = compute_rollout_step(args, rollout_id) @@ -734,9 +717,9 @@ def compute_perf_metrics_from_samples(args, samples, rollout_time): parser_depths = [s.parser_max_queue_depth for s in samples if s.parser_max_queue_depth is not None] if parser_depths: log_dict["parser_max_queue_depth"] = max(parser_depths) - reward_depths = [s.reward_max_queue_depth for s in samples if s.reward_max_queue_depth] - for name in sorted({name for depths in reward_depths for name in depths}): - log_dict[f"reward_max_queue_depth_{name}"] = max(depths[name] for depths in reward_depths if name in depths) + reward_depths = [s.reward_max_queue_depth for s in samples if s.reward_max_queue_depth is not None] + if reward_depths: + log_dict["reward_max_queue_depth"] = max(reward_depths) return log_dict diff --git a/miles/ray/utils.py b/miles/ray/utils.py index f09fe0d5b..b1e2b9fc1 100644 --- a/miles/ray/utils.py +++ b/miles/ray/utils.py @@ -12,11 +12,6 @@ # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/neuron.py#L108-L109 # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/tpu.py#L171-L172 # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/intel_gpu.py#L97-L98 -# GPU shares of one placement-group bundle; Ray admits an actor only while they still fit, so they sum to 1.0 -TRAIN_ACTOR_GPU = 0.7 -ROLLOUT_ENGINE_GPU = 0.25 -COLOCATED_REWARD_GPU = 0.05 - NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py index d3477835a..e2c981686 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -15,33 +15,15 @@ async def async_rm(args, sample: Sample, **kwargs): if rm_type == "ocr": from .ocr import ocr_rm - return (await ocr_rm(args, [sample]))[0] + return await ocr_rm(args, sample) elif rm_type == "pickscore": from .pickscore import pickscore_rm return (await pickscore_rm(args, [sample]))[0] - elif rm_type == "hps": - from .hps import hps_rm - - return (await hps_rm(args, [sample]))[0] else: raise NotImplementedError(f"Rule-based RM for {rm_type!r} is not implemented.") -def create_colocated_reward_pools(args, placement_group, slots) -> list: - """Seat every colocated pool; the rm functions' singleton lookup then finds them.""" - pools = [] - if args.pickscore_reward_colocate: - from .pickscore import AsyncPickScorePool - - pools.append(AsyncPickScorePool(args, placement_group=placement_group, slots=slots)) - if args.hps_reward_colocate: - from .hps import AsyncHPSPool - - pools.append(AsyncHPSPool(args, placement_group=placement_group, slots=slots)) - return pools - - async def batched_async_rm( args, samples: list[Sample], @@ -57,14 +39,6 @@ async def batched_async_rm( from .pickscore import pickscore_rm return await pickscore_rm(args, samples) - if all(rm_type == "hps" for rm_type in rm_types): - from .hps import hps_rm - - return await hps_rm(args, samples) - if all(rm_type == "ocr" for rm_type in rm_types): - from .ocr import ocr_rm - - return await ocr_rm(args, samples) tasks = [async_rm(args, sample, **kwargs) for sample in samples] rewards = await asyncio.gather(*tasks) diff --git a/miles/rollout/rm_hub/core.py b/miles/rollout/rm_hub/core.py index 65f8e0dbb..b443a1292 100644 --- a/miles/rollout/rm_hub/core.py +++ b/miles/rollout/rm_hub/core.py @@ -8,85 +8,26 @@ import ray from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from miles.ray.utils import COLOCATED_REWARD_GPU - +_manager_placement_group = None logger = logging.getLogger(__name__) -def bundle_deal_order( - bundle_indices: list[int], - gpu_ids: list[int], - num_gpus_per_node: int, - num_gpus_per_engine: int, -) -> list[int]: - """Order bundles for reward actors: those without a rollout engine's own claim first, spread across GPUs.""" - span = min(num_gpus_per_engine, num_gpus_per_node) - covered = len(bundle_indices) // span * span - - def key(item: tuple[int, int, int]) -> tuple[int, int, int, int]: - position, _bundle_index, gpu_id = item - if position >= covered: - return (0, -1, gpu_id, position) - position_in_actor = position % span - # an engine declares its claim on the first bundle of its span, so deal those last - if position_in_actor == 0: - return (1, 0, gpu_id, position) - return (0, position_in_actor, gpu_id, position) - - slots = zip(range(len(bundle_indices)), bundle_indices, gpu_ids, strict=True) - return [bundle_index for _, bundle_index, _ in sorted(slots, key=key)] - - -class ColocatedRewardSlots: - """Deal placement-group bundles to colocated reward pools, one long-lived actor per bundle. - - Owned by ``RolloutManager`` and shared by every colocated pool. - """ - - def __init__(self, order: list[int]) -> None: - self._order = order - self._owners: dict[int, str] = {} - self._pool_names: set[str] = set() - - @property - def remaining(self) -> int: - return len(self._order) - len(self._owners) - - def allocate(self, name: str, num_workers: int) -> list[int]: - if name in self._pool_names: - raise RuntimeError(f"--{name}-reward-colocate: {name} already owns reward slots") - if num_workers > self.remaining: - raise RuntimeError( - f"--{name}-reward-colocate: {num_workers} slots requested, but only " - f"{self.remaining}/{len(self._order)} remain ({self}). " - f"Reduce --{name}-num-workers or run the pool on dedicated GPUs." - ) - start = len(self._owners) - slots = list(self._order[start : start + num_workers]) - self._owners.update({slot: name for slot in slots}) - self._pool_names.add(name) - return slots - - def __str__(self) -> str: - return ", ".join(f"bundle {bundle}: {self._owners.get(bundle)}" for bundle in sorted(self._order)) - - -def record_reward_queue_depth(samples, name: str, max_queue_depth: int) -> None: - """Store the backlog per pool so the perf log shows which reward lags when several score one batch.""" - for sample in samples: - depths = sample.reward_max_queue_depth or {} - depths[name] = float(max_queue_depth) - sample.reward_max_queue_depth = depths +def set_manager_placement_group(pg) -> None: + """Publish the manager's (pg, bundle_indices, gpu_ids) for colocated actor pools.""" + global _manager_placement_group + _manager_placement_group = pg -class AsyncRewardActorPool: - """Round-robin pool for Ray reward actors exposing ``score_batch(outputs, prompts)``. +def get_manager_placement_group(): + return _manager_placement_group + - Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves. - Colocated pools take one slot per rollout bundle from the manager's - ``slots``; standalone pools are default-scheduled at ``num_gpus_per_worker``, which - only lands on GPUs outside every placement group. - """ +set_reward_placement_group = set_manager_placement_group +get_reward_placement_group = get_manager_placement_group + + +class AsyncRewardActorPool: + """Round-robin pool for Ray reward actors exposing ``score_batch``.""" def __init__( self, @@ -98,30 +39,32 @@ def __init__( num_gpus_per_worker: float, colocate: bool, name: str, - placement_group=None, - slots: ColocatedRewardSlots | None = None, ) -> None: if colocate: - if placement_group is None or slots is None: - raise RuntimeError(f"--{name}-reward-colocate: the {name} pool was not seated by RolloutManager.") - pg, _, _ = placement_group + pg, bundle_indices, _ = get_reward_placement_group() + # bundle_indices is sorted by (node, gpu); stride so workers spread across nodes + # instead of stacking onto the first node's GPUs. + stride = max(1, len(bundle_indices) // num_workers) strategies = [ - PlacementGroupSchedulingStrategy(placement_group=pg, placement_group_bundle_index=bundle) - for bundle in slots.allocate(name, num_workers) + PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=bundle_indices[w * stride], + ) + for w in range(num_workers) ] - num_gpus_per_worker = COLOCATED_REWARD_GPU + num_gpus_per_worker = 0.05 + num_cpus_per_worker = 0.05 else: strategies = ["DEFAULT"] * num_workers + num_cpus_per_worker = 1 self._actors = [ - ray.remote(actor_cls) - .options( - num_cpus=num_gpus_per_worker, + actor_cls.options( + num_cpus=num_cpus_per_worker, num_gpus=num_gpus_per_worker, - scheduling_strategy=strategy, - ) - .remote(**actor_kwargs) - for strategy in strategies + scheduling_strategy=strategies[i], + ).remote(**actor_kwargs) + for i in range(num_workers) ] self._batch_size = batch_size self._round_robin_index = 0 @@ -139,16 +82,16 @@ def _next_actor_idx(self) -> int: self._round_robin_index += 1 return i - async def score(self, outputs: list, prompts: list[str]) -> tuple[list[float], int]: - """Score samples in batches; also report the deepest dispatch-time backlog this call saw.""" + async def score(self, images: list, prompts: list[str]) -> tuple[list[float], int]: + """Score in batches; also report the deepest dispatch-time backlog this call saw.""" refs, idxs, max_queue_depth = [], [], 0 - for start in range(0, len(outputs), self._batch_size): + for start in range(0, len(images), self._batch_size): end = start + self._batch_size i = self._next_actor_idx() max_queue_depth = max(max_queue_depth, self._inflight[i]) self._inflight[i] += 1 idxs.append(i) - refs.append(self._actors[i].score_batch.remote(outputs[start:end], prompts[start:end])) + refs.append(self._actors[i].score_batch.remote(images[start:end], prompts[start:end])) loop = asyncio.get_running_loop() try: diff --git a/miles/rollout/rm_hub/hps.py b/miles/rollout/rm_hub/hps.py deleted file mode 100644 index 0311f102f..000000000 --- a/miles/rollout/rm_hub/hps.py +++ /dev/null @@ -1,150 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence - -import ray -import torch -from PIL import Image -from torchvision.transforms import InterpolationMode -from torchvision.transforms import functional as vision_functional - -from miles.utils.misc import SingletonMeta -from miles.utils.processing_utils import generated_output_to_rgb_hwc_uint8_frames -from miles.utils.types import Sample - -from .core import AsyncRewardActorPool, record_reward_queue_depth - -_HPS_VERSION_TO_FILENAME = { - "v2.0": "HPS_v2_compressed.pt", - "v2.1": "HPS_v2.1_compressed.pt", -} - - -class _HPSImageTransform: - """HPSv2's inference resize: fit the longest side, then zero-pad.""" - - def __init__(self, image_size: tuple[int, int], mean: Sequence[float], std: Sequence[float]) -> None: - self.image_size = image_size[0] - self.mean = list(mean) - self.std = list(std) - - def __call__(self, image: Image.Image) -> torch.Tensor: - tensor = vision_functional.to_tensor(image) - height, width = tensor.shape[-2:] - scale = self.image_size / float(max(height, width)) - new_height, new_width = (round(height * scale), round(width * scale)) - if (new_height, new_width) != (height, width): - tensor = vision_functional.resize( - tensor, - [new_height, new_width], - interpolation=InterpolationMode.BICUBIC, - ) - - pad_height = self.image_size - new_height - pad_width = self.image_size - new_width - tensor = vision_functional.pad( - tensor, - [ - pad_width // 2, - pad_height // 2, - pad_width - pad_width // 2, - pad_height - pad_height // 2, - ], - fill=0, - ) - return vision_functional.normalize(tensor, mean=self.mean, std=self.std) - - -class HPSScorer(torch.nn.Module): - """HPSv2 scorer for aligned prompt/image batches.""" - - def __init__( - self, - *, - device: str = "cuda", - hps_version: str = "v2.1", - checkpoint_path: str | None = None, - ) -> None: - super().__init__() - import huggingface_hub - from open_clip import create_model, get_tokenizer - from open_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD - - self.device = torch.device(device) - model = create_model( - "ViT-H-14", - pretrained=None, - precision="amp", - device=str(self.device), - jit=False, - force_quick_gelu=False, - force_custom_text=False, - force_patch_dropout=False, - force_image_size=None, - pretrained_image=False, - output_dict=True, - ) - - if checkpoint_path is None: - checkpoint_path = huggingface_hub.hf_hub_download("xswu/HPSv2", _HPS_VERSION_TO_FILENAME[hps_version]) - checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) - model.load_state_dict(checkpoint["state_dict"]) - self.model = model.to(self.device).eval() - self.preprocess = _HPSImageTransform(model.visual.image_size, OPENAI_DATASET_MEAN, OPENAI_DATASET_STD) - self.tokenizer = get_tokenizer("ViT-H-14") - - @torch.no_grad() - def forward(self, prompts: Sequence[str], images: Sequence[Image.Image]) -> list[float]: - image_batch = torch.stack([self.preprocess(image) for image in images]).to(self.device, non_blocking=True) - text_batch = self.tokenizer(list(prompts)).to(self.device, non_blocking=True) - with torch.amp.autocast(self.device.type, enabled=self.device.type == "cuda"): - outputs = self.model(image_batch, text_batch) - scores = torch.diagonal(outputs["image_features"] @ outputs["text_features"].T) - return [float(score) for score in scores.detach().float().cpu()] - - -class HPSRewardActor: - def __init__(self, *, hps_version: str, checkpoint_path: str | None = None) -> None: - use_cuda = bool(ray.get_gpu_ids()) and torch.cuda.is_available() - if use_cuda: - torch.cuda.set_device(0) - self.scorer = HPSScorer( - device="cuda" if use_cuda else "cpu", - hps_version=hps_version, - checkpoint_path=checkpoint_path, - ) - - def score_batch(self, outputs: list[torch.Tensor], prompts: list[str]) -> list[float]: - # HPSv2 rounds when quantising to uint8; matching it keeps scores comparable with the reference - images = [] - for output in outputs: - (image,) = generated_output_to_rgb_hwc_uint8_frames(output, None, round_normalized=True) - images.append(Image.fromarray(image)) - return self.scorer(prompts, images) - - -class AsyncHPSPool(AsyncRewardActorPool, metaclass=SingletonMeta): - """Ray actor pool for HPS reward inference.""" - - def __init__(self, args, placement_group=None, slots=None) -> None: - super().__init__( - actor_cls=HPSRewardActor, - actor_kwargs={ - "hps_version": args.hps_version, - "checkpoint_path": args.hps_checkpoint_path, - }, - num_workers=args.hps_num_workers, - batch_size=args.hps_batch_size, - num_gpus_per_worker=args.hps_num_gpus_per_worker, - colocate=args.hps_reward_colocate, - name="hps", - placement_group=placement_group, - slots=slots, - ) - - -async def hps_rm(args, samples: Sequence[Sample]) -> list[float]: - pool = AsyncHPSPool(args) - scores, max_queue_depth = await pool.score([s.generated_output for s in samples], [s.prompt for s in samples]) - record_reward_queue_depth(samples, "hps", max_queue_depth) - return scores diff --git a/miles/rollout/rm_hub/ocr.py b/miles/rollout/rm_hub/ocr.py index 35403befc..ad29d1965 100644 --- a/miles/rollout/rm_hub/ocr.py +++ b/miles/rollout/rm_hub/ocr.py @@ -1,23 +1,22 @@ import logging -from collections.abc import Sequence import numpy as np +import ray import torch +from Levenshtein import distance +from paddleocr import PaddleOCR from PIL import Image from miles.utils.misc import SingletonMeta -from miles.utils.processing_utils import generated_output_to_rgb_hwc_uint8_frames +from miles.utils.processing_utils import cfhw_to_fhwc, image_or_video_to_uint8 from miles.utils.types import Sample -from .core import AsyncRewardActorPool, record_reward_queue_depth +from .core import AsyncRewardActorPool logger = logging.getLogger(__name__) -def _init_paddleocr(use_gpu: bool): - # actor-only dependency: the manager imports this module just to dispatch - from paddleocr import PaddleOCR - +def _init_paddleocr(use_gpu: bool) -> PaddleOCR: return PaddleOCR( use_angle_cls=False, lang="en", @@ -42,8 +41,6 @@ def __call__(self, images: list[Image.Image] | list[np.ndarray], prompts: list[s :param prompts: Corresponding target text list :return: Reward tensor (CPU) """ - from Levenshtein import distance - prompts = [prompt.split('"')[1] for prompt in prompts] rewards = [] # Ensure input lengths are consistent @@ -83,15 +80,14 @@ def __call__(self, images: list[Image.Image] | list[np.ndarray], prompts: list[s return rewards +@ray.remote class OcrRewardActor: def __init__(self, use_gpu: bool = False): self.scorer = OcrScorer(use_gpu=use_gpu) - def score_batch(self, outputs: list[torch.Tensor], prompts: list[str]) -> list[float]: - assert len(outputs) == 1, f"OCR scores one image per call, got {len(outputs)}" - # flow_grpo feeds PaddleOCR rounded RGB (not BGR) uint8; matching it keeps the reward bit-identical - (image,) = generated_output_to_rgb_hwc_uint8_frames(outputs[0], None, round_normalized=True) - return self.scorer([image], prompts) + def score_batch(self, images: list, prompts: list[str]) -> list[float]: + assert len(images) == 1, f"OCR scores one image per call, got {len(images)}" + return self.scorer(images, prompts) class AsyncOcrPool(AsyncRewardActorPool, metaclass=SingletonMeta): @@ -105,12 +101,39 @@ def __init__(self, args) -> None: batch_size=1, num_gpus_per_worker=0, colocate=False, - name="ocr", + name="OCR", ) -async def ocr_rm(args, samples: Sequence[Sample]) -> list[float]: +def _rgb_hwc_from_generated(sample: Sample) -> np.ndarray: + """``generated_output``: ``[C, F, H, W]`` or ``[C, H, W]``; use frame index 0. + + Accepts both the local-rollout format ``[C, F, H, W]`` (video frames) and + the sglang-diffusion SD3 format ``[C, H, W]`` (static image, no frame dim). + + Feeds PaddleOCR the exact same ``(RGB, uint8 HWC)`` array that flow_grpo's + ``ocr_score`` wrapper does β€” `(images * 255).round().clamp(0,255).to(uint8)` + then ``transpose(0, 2, 3, 1)``, no channel swap. PaddleOCR's OpenCV stack + would prefer BGR, but flow_grpo trains against the (slightly-off) RGB + convention, so we match that to keep the reward signal bit-identical. + """ + t = sample.generated_output + if t is None: + raise ValueError("generated_output is None") + t = t.detach().cpu().float() + if t.ndim == 3: + t = t.unsqueeze(1) + if t.ndim != 4: + raise ValueError(f"generated_output must be 3D [C, H, W] or 4D [C, F, H, W], got {tuple(t.shape)}") + fhwc = cfhw_to_fhwc(t) + if fhwc.shape[0] != 1: + raise ValueError(f"generated_output frame dim F must be 1 for image models, got F={fhwc.shape[0]}") + return image_or_video_to_uint8(fhwc[0], round_normalized=True).numpy() + + +async def ocr_rm(args, sample: Sample): pool = AsyncOcrPool(args) - scores, max_queue_depth = await pool.score([s.generated_output for s in samples], [s.prompt for s in samples]) - record_reward_queue_depth(samples, "ocr", max_queue_depth) - return scores + image = _rgb_hwc_from_generated(sample) + scores, max_queue_depth = await pool.score([image], [sample.prompt]) + sample.reward_max_queue_depth = float(max_queue_depth) + return scores[0] diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index b58d99f2a..0fca315e9 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -2,15 +2,27 @@ from collections.abc import Sequence +import numpy as np import ray import torch from PIL import Image from miles.utils.misc import SingletonMeta -from miles.utils.processing_utils import generated_output_to_rgb_hwc_uint8_frames, sample_frame_indices +from miles.utils.processing_utils import cfhw_to_fhwc, image_or_video_to_uint8 from miles.utils.types import Sample -from .core import AsyncRewardActorPool, record_reward_queue_depth +from .core import AsyncRewardActorPool + + +def sample_frame_indices(num_total_frames: int, num_frames: int | None) -> list[int]: + if num_total_frames <= 0: + raise ValueError(f"video has no frames: {num_total_frames}") + if num_frames is None or num_total_frames <= num_frames: + return list(range(num_total_frames)) + if num_frames == 1: + return [num_total_frames // 2] + step = (num_total_frames - 1) / (num_frames - 1) + return [int(round(i * step)) for i in range(num_frames)] def _feature_tensor(features): @@ -22,6 +34,18 @@ def _feature_tensor(features): raise TypeError(f"Cannot extract embedding tensor from {type(features)!r}") +def _sample_to_rgb_hwc_uint8_frames(sample: Sample, num_frames: int | None) -> list[np.ndarray]: + cfhw = sample.generated_output + if cfhw is None: + raise ValueError("generated_output is None") + + # Convert only the frames that survive: a 107-frame clip yields 8 here, and + # converting the whole clip first cost several full-size float32 copies of it + indices = sample_frame_indices(cfhw.shape[1], num_frames) + fhwc = cfhw_to_fhwc(image_or_video_to_uint8(cfhw[:, indices].detach().cpu())) + return [np.ascontiguousarray(fhwc[i].numpy()) for i in range(len(indices))] + + class PickScoreScorer(torch.nn.Module): """CLIP PickScore for (prompt, image) pairs; raw logits scaled to ~0-1.""" @@ -67,15 +91,14 @@ def forward(self, prompts: Sequence[str], images: Sequence[Image.Image]) -> list return [float(score) for score in scores.detach().cpu()] +@ray.remote class PickScoreRewardActor: def __init__( self, *, processor_path: str, model_path: str, - frames_per_forward: int, ) -> None: - self.frames_per_forward = frames_per_forward gpu_ids = ray.get_gpu_ids() use_cuda = bool(gpu_ids) and torch.cuda.is_available() if use_cuda: @@ -87,54 +110,46 @@ def __init__( model_path=model_path, ) - def score_batch(self, outputs: list[torch.Tensor], prompts: list[str]) -> list[float]: - # a video sample is scored on every frame it arrives with and gets their mean - images, frame_counts = [], [] - for output in outputs: - frames = generated_output_to_rgb_hwc_uint8_frames(output, None) - images.extend(Image.fromarray(frame) for frame in frames) - frame_counts.append(len(frames)) - flat_prompts = [p for p, n in zip(prompts, frame_counts, strict=True) for _ in range(n)] - # forwards see --pickscore-batch-size frames, the chunking the e2e standards were recorded with - flat_scores = [] - for start in range(0, len(images), self.frames_per_forward): - end = start + self.frames_per_forward - flat_scores.extend(self.scorer(flat_prompts[start:end], images[start:end])) - scores, offset = [], 0 - for count in frame_counts: - scores.append(float(sum(flat_scores[offset : offset + count]) / count)) - offset += count - return scores + def score_batch(self, images: list, prompts: list[str]) -> list[float]: + pil_images = [Image.fromarray(image) if isinstance(image, np.ndarray) else image for image in images] + return self.scorer(prompts, pil_images) class AsyncPickScorePool(AsyncRewardActorPool, metaclass=SingletonMeta): """Ray actor pool for GPU PickScore reward inference.""" - def __init__(self, args, placement_group=None, slots=None) -> None: + def __init__(self, args) -> None: super().__init__( actor_cls=PickScoreRewardActor, actor_kwargs={ "processor_path": args.pickscore_processor_path, "model_path": args.pickscore_model_path, - "frames_per_forward": args.pickscore_batch_size, }, num_workers=args.pickscore_num_workers, batch_size=args.pickscore_batch_size, num_gpus_per_worker=args.pickscore_num_gpus_per_worker, - colocate=args.pickscore_reward_colocate, - name="pickscore", - placement_group=placement_group, - slots=slots, + colocate=args.colocate_reward, + name="PickScore", ) async def pickscore_rm(args, samples: Sequence[Sample]) -> list[float]: pool = AsyncPickScorePool(args) - # pick --pickscore-num-frames here so only the scored frames cross the object store - outputs = [ - s.generated_output[:, sample_frame_indices(s.generated_output.shape[1], args.pickscore_num_frames)] - for s in samples - ] - scores, max_queue_depth = await pool.score(outputs, [s.prompt for s in samples]) - record_reward_queue_depth(samples, "pickscore", max_queue_depth) + images: list[np.ndarray] = [] + prompts: list[str] = [] + frame_counts: list[int] = [] + for sample in samples: + frames = _sample_to_rgb_hwc_uint8_frames(sample, args.pickscore_num_frames) + images.extend(frames) + prompts.extend([sample.prompt] * len(frames)) + frame_counts.append(len(frames)) + + flat_scores, max_queue_depth = await pool.score(images, prompts) + for sample in samples: + sample.reward_max_queue_depth = float(max_queue_depth) + scores: list[float] = [] + offset = 0 + for count in frame_counts: + scores.append(float(sum(flat_scores[offset : offset + count]) / count)) + offset += count return scores diff --git a/miles/rollout/rm_hub/weighted_mixture_rm.py b/miles/rollout/rm_hub/weighted_mixture_rm.py deleted file mode 100644 index 5420c6f06..000000000 --- a/miles/rollout/rm_hub/weighted_mixture_rm.py +++ /dev/null @@ -1,51 +0,0 @@ -"""``--custom-rm-path`` example: a weighted sum of built-in rewards, weighted by ``--custom-rm-args``. - - --custom-rm-path miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm \\ - --custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weighted - -Each sample's reward is a dict holding every component plus ``"weighted"``, so each reward -gets its own ``rollout/reward/_mean`` panel while ``--reward-key`` picks what GRPO trains -on. Each named reward scores the whole batch once and keeps its own placement flags -(``---reward-colocate``, ``---num-gpus-per-worker``). Weights apply to raw scores, -whose scales differ: HPSv2.1 ~0.3, PickScore/26 ~0.85, OCR in [0, 1]. -""" - -import asyncio -from collections.abc import Sequence - -from miles.utils.types import Sample - -from .hps import hps_rm -from .ocr import ocr_rm -from .pickscore import pickscore_rm - -_REWARDS = {"hps": hps_rm, "pickscore": pickscore_rm, "ocr": ocr_rm} - - -def parse_weights(custom_rm_args: str) -> list[tuple[str, float]]: - weights = [] - # launch scripts hand the arg string to `sh`, where ";" would end the command; "," is inert - for term in custom_rm_args.split(","): - name, _, weight = term.strip().partition("=") - if name not in _REWARDS: - raise ValueError( - f"--custom-rm-args: unknown reward {name!r} in {custom_rm_args!r}; choose from {tuple(_REWARDS)}" - ) - weights.append((name, float(weight))) - return weights - - -async def weighted_mixture_rm(args, samples: Sequence[Sample], **kwargs) -> list[dict[str, float]]: - weights = parse_weights(args.custom_rm_args) - if args.reward_key not in {name for name, _ in weights} | {"weighted"}: - raise ValueError( - f"weighted_mixture_rm returns a dict per sample; pass --reward-key weighted (or one of " - f"{[name for name, _ in weights]}), got {args.reward_key!r}" - ) - per_reward = await asyncio.gather(*(_REWARDS[name](args, samples) for name, _ in weights)) - rewards = [] - for i in range(len(samples)): - components = {name: scores[i] for (name, _), scores in zip(weights, per_reward, strict=True)} - components["weighted"] = sum(weight * components[name] for name, weight in weights) - rewards.append(components) - return rewards diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py index dd5b09aab..717e10681 100644 --- a/miles/rollout/sft_rollout.py +++ b/miles/rollout/sft_rollout.py @@ -19,9 +19,10 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from miles.rollout.base_types import RolloutFnTrainOutput +from miles.rollout.rm_hub.core import get_manager_placement_group from miles.utils import tracking_utils from miles.utils.metric_utils import compute_rollout_step -from miles.utils.misc import SingletonMeta, load_function +from miles.utils.misc import load_function from miles.utils.types import Sample logger = logging.getLogger(__name__) @@ -186,26 +187,16 @@ def encode(self, items: list[dict], cache_dir: str) -> int: return len(items) +_encode_actors: list | None = None _scheduler_grid: tuple[torch.Tensor, torch.Tensor] | None = None -class SftEncodePool(metaclass=SingletonMeta): - """One encode actor per rollout placement-group bundle: encode is SFT's rollout, seated by RolloutManager.""" - - def __init__(self, args, placement_group=None) -> None: - if placement_group is None: - raise RuntimeError("SftEncodePool is seated by RolloutManager; --loss-type sft_loss was not set.") - self._args = args - self._placement_group = placement_group - self._actors: list | None = None - - @property - def actors(self) -> list: - # a fully cached dataset never encodes, so the encoders (tens of GB each) load on the first miss only - if self._actors is not None: - return self._actors - pg, bundle_indices, _ = self._placement_group - self._actors = [ +def _encode_pool(args) -> list: + global _encode_actors + if _encode_actors is None: + # Encode is SFT's rollout: the pool takes the rollout placement seats sglang engines use in RL. + pg, bundle_indices, _ = get_manager_placement_group() + _encode_actors = [ SftEncodeActor.options( num_cpus=ENCODE_GPU_FRACTION, num_gpus=ENCODE_GPU_FRACTION, @@ -213,11 +204,11 @@ def actors(self) -> list: placement_group=pg, placement_group_bundle_index=i, ), - ).remote(self._args) + ).remote(args) for i in bundle_indices ] - logger.info("SFT encode pool: %d workers at %.2f GPU each", len(self._actors), ENCODE_GPU_FRACTION) - return self._actors + logger.info("SFT encode pool: %d workers at %.2f GPU each", len(_encode_actors), ENCODE_GPU_FRACTION) + return _encode_actors def _get_scheduler_grid(args) -> tuple[torch.Tensor, torch.Tensor]: @@ -264,7 +255,7 @@ def generate_rollout(args, rollout_id, data_source, evaluation: bool = False) -> if missing: cache_dir.mkdir(parents=True, exist_ok=True) start = time.time() - actors = SftEncodePool(args).actors + actors = _encode_pool(args) miss_items = list(missing.values()) shards = [miss_items[i :: len(actors)] for i in range(len(actors))] ray.get( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c2c0e88d9..d96eb6599 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1221,7 +1221,7 @@ def add_reward_model_arguments(parser): "--rm-type", type=str, default=None, - help="Built-in reward model (pickscore / hps / ocr). Ignored when --custom-rm-path is set.", + help="Type of the reward model", ) parser.add_argument( "--reward-key", @@ -1290,10 +1290,11 @@ def add_reward_model_arguments(parser): "Use a fractional value below 1.0 for lightweight single-GPU reward models.", ) parser.add_argument( - "--pickscore-reward-colocate", + "--colocate-reward", action="store_true", default=False, - help="Seat PickScore actors on the rollout GPUs, one per placement-group bundle. Requires --colocate.", + help="Colocate reward actors onto rollout GPUs (train 0.7 + rollout 0.25 + reward 0.05). " + "Requires --colocate.", ) parser.add_argument( "--pickscore-batch-size", @@ -1319,43 +1320,6 @@ def add_reward_model_arguments(parser): default=None, help="Hugging Face model path for PickScore. Required when --rm-type pickscore.", ) - parser.add_argument( - "--hps-num-workers", - type=int, - default=1, - help="Number of Ray HPS actors used when --rm-type hps.", - ) - parser.add_argument( - "--hps-num-gpus-per-worker", - type=float, - default=1.0, - help="GPU resources per HPS actor when reward is not colocated.", - ) - parser.add_argument( - "--hps-reward-colocate", - action="store_true", - default=False, - help="Seat HPS actors on the rollout GPUs, one per placement-group bundle. Requires --colocate.", - ) - parser.add_argument( - "--hps-batch-size", - type=int, - default=8, - help="Batch size per HPS actor call.", - ) - parser.add_argument( - "--hps-version", - type=str, - default="v2.1", - choices=["v2.0", "v2.1"], - help="HPS checkpoint version used when --rm-type hps.", - ) - parser.add_argument( - "--hps-checkpoint-path", - type=str, - default=None, - help="Optional local HPS checkpoint path; otherwise download it from Hugging Face.", - ) parser.add_argument( "--rm-url", @@ -1378,18 +1342,7 @@ def add_reward_model_arguments(parser): "Signature: `async def custom_rm(args, samples: list[Sample], **kwargs) -> list[float]`. " "Wired in batched_async_rm only β€” per-sample async_rm dispatch was deliberately " "removed to avoid the (args, sample) vs (args, list) signature ambiguity. " - "If you want per-sample routing, do it inside your batched function. " - "Shipped: miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm, a weighted sum of " - 'built-in rewards configured by --custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weighted.' - ), - ) - parser.add_argument( - "--custom-rm-args", - type=str, - default=None, - help=( - "Opaque config string handed to the --custom-rm-path function as `args.custom_rm_args`; " - 'e.g. "hps=0.7,pickscore=0.3" for miles.rollout.rm_hub.weighted_mixture_rm.' + "If you want per-sample routing, do it inside your batched function." ), ) parser.add_argument( @@ -1782,38 +1735,12 @@ def miles_validate_args(args): if args.offload_rollout is None: args.offload_rollout = False - if args.hps_num_workers <= 0: - raise ValueError(f"--hps-num-workers must be positive, got {args.hps_num_workers}") - if args.hps_batch_size <= 0: - raise ValueError(f"--hps-batch-size must be positive, got {args.hps_batch_size}") - if args.hps_num_gpus_per_worker < 0: - raise ValueError(f"--hps-num-gpus-per-worker must be non-negative, got {args.hps_num_gpus_per_worker}") - - colocated_reward_workers = { - name: num_workers - for name, colocate, num_workers in ( - ("pickscore", args.pickscore_reward_colocate, args.pickscore_num_workers), - ("hps", args.hps_reward_colocate, args.hps_num_workers), - ) - if colocate - } - if colocated_reward_workers: - if not args.colocate: - raise ValueError( - f"{', '.join(f'--{name}-reward-colocate' for name in colocated_reward_workers)} requires --colocate." - ) - if sum(colocated_reward_workers.values()) > args.rollout_num_gpus: - raise ValueError( - f"colocated reward workers ({sum(colocated_reward_workers.values())}) exceed rollout_num_gpus " - f"({args.rollout_num_gpus}): the placement group has one reward slot per GPU." - ) - if args.pickscore_reward_colocate and not (args.pickscore_model_path and args.pickscore_processor_path): - raise ValueError( - "--pickscore-reward-colocate builds the PickScore actors at startup; set --pickscore-model-path " - "and --pickscore-processor-path." + if args.colocate_reward: + assert args.colocate, "--colocate-reward requires --colocate." + assert args.pickscore_num_workers <= args.rollout_num_gpus, ( + f"--colocate-reward requires pickscore_num_workers ({args.pickscore_num_workers}) " + f"<= rollout_num_gpus ({args.rollout_num_gpus}): the placement group has one bundle per GPU." ) - if args.custom_rm_args is not None and args.custom_rm_path is None: - raise ValueError("--custom-rm-args requires --custom-rm-path.") if args.eval_function_path is None: args.eval_function_path = args.rollout_function_path diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py index 513065f33..a83921bc3 100644 --- a/miles/utils/processing_utils.py +++ b/miles/utils/processing_utils.py @@ -1,8 +1,7 @@ -"""Image and video tensor processing used by diffusion output consumers.""" +"""Tensor layout conversions used by diffusion output consumers.""" from __future__ import annotations -import numpy as np import torch @@ -28,26 +27,3 @@ def image_or_video_to_uint8(tensor: torch.Tensor, *, round_normalized: bool = Fa if round_normalized: output = output.round() return output.clamp(0, 255).to(torch.uint8) - - -def sample_frame_indices(num_total_frames: int, num_frames: int | None) -> list[int]: - if num_total_frames <= 0: - raise ValueError(f"video has no frames: {num_total_frames}") - if num_frames is None or num_total_frames <= num_frames: - return list(range(num_total_frames)) - if num_frames == 1: - return [num_total_frames // 2] - step = (num_total_frames - 1) / (num_frames - 1) - return [int(round(i * step)) for i in range(num_frames)] - - -def generated_output_to_rgb_hwc_uint8_frames( - cfhw: torch.Tensor, - num_frames: int | None, - *, - round_normalized: bool = False, -) -> list[np.ndarray]: - indices = sample_frame_indices(cfhw.shape[1], num_frames) - selected = image_or_video_to_uint8(cfhw[:, indices].detach().cpu(), round_normalized=round_normalized) - fhwc = cfhw_to_fhwc(selected) - return [np.ascontiguousarray(fhwc[i].numpy()) for i in range(len(indices))] diff --git a/miles/utils/types.py b/miles/utils/types.py index 02b2efdd5..88bdd39e8 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -87,8 +87,7 @@ class Sample: inference_time_s: float | None = None # dispatch-time backlog observed when this sample's parse / reward was submitted parser_max_queue_depth: float | None = None - # per reward pool, keyed by pool name - reward_max_queue_depth: dict[str, float] | None = None + reward_max_queue_depth: float | None = None peak_memory_mb: float | None = None # Scalar from single RM (e.g. pickscore) or dict when combining multiple RMs diff --git a/requirements.txt b/requirements.txt index d084e840d..4b080313e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,7 +13,6 @@ mcp[cli]==1.25.0 memray==1.19.1 nvidia-ml-py==13.595.45 omegaconf==2.3.0 -open_clip_torch==2.24.0 openai==2.6.1 opencv-python-headless==4.10.0.84 paddleocr==2.9.1 diff --git a/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py index d9dc006be..b61d42b3f 100644 --- a/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py +++ b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py @@ -5,7 +5,7 @@ guidance 4.0, Flow-SDE noise_level=0.7, no KL, per-prompt mean + global std. Layout: train, rollout and PickScore reward all share the same 4 GPUs -(--colocate --pickscore-reward-colocate, one PickScore worker per rollout engine). +(--colocate --colocate-reward, one PickScore worker per rollout engine). SDE schedule: epoch_global_random_choice draws 2 steps per epoch from candidates 8-11. The Cosmos3 checkpoint ships a Karras flow-sigma grid whose @@ -101,7 +101,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: reward_args = ( "--rm-type pickscore " - "--pickscore-reward-colocate " + "--colocate-reward " "--pickscore-num-workers 4 " "--pickscore-batch-size 8 " "--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K " diff --git a/scripts/run_diffusion_grpo_h3_t2va_2gpu.py b/scripts/run_diffusion_grpo_h3_t2va_2gpu.py index ba310c107..54f7e91bb 100644 --- a/scripts/run_diffusion_grpo_h3_t2va_2gpu.py +++ b/scripts/run_diffusion_grpo_h3_t2va_2gpu.py @@ -184,7 +184,7 @@ def execute(args: ScriptArgs, prompt_dir: str) -> None: "--rollout-num-gpus-per-engine 2 " "--num-gpus-per-node 2 " "--colocate " - "--pickscore-reward-colocate " + "--colocate-reward " "--deterministic-mode " "--rollout-health-check-interval 60 " "--miles-router-health-check-failure-threshold 30 " diff --git a/scripts/run_diffusion_grpo_sd3_hps_sglang.py b/scripts/run_diffusion_grpo_sd3_hps_sglang.py deleted file mode 100644 index e5da1c9d0..000000000 --- a/scripts/run_diffusion_grpo_sd3_hps_sglang.py +++ /dev/null @@ -1,130 +0,0 @@ -"""SD3.5-medium HPS GRPO through the sglang-diffusion /rollout/generate path. - -2-GPU colocate: FSDP DP=2, two rollout engines, and one HPS worker share the same GPUs. - -SD3.5 is gated, so HF_TOKEN must be set even when the weights are cached β€” sglang still -fetches model_index.json from the hub at startup. - -Usage: - python3 scripts/run_diffusion_grpo_sd3_hps_sglang.py - MILES_SCRIPT_DEBUG_ALIGNMENT=1 python3 scripts/run_diffusion_grpo_sd3_hps_sglang.py -""" - -import os -from dataclasses import dataclass - -import typer - -import miles.utils.external_utils.command_utils as U - -MODEL = "stabilityai/stable-diffusion-3.5-medium" -DATASET = "rockdu/miles-diffusion-datasets" -DATASET_SUBSET = "hpdv2" -WANDB_PROJECT = "miles-diffusion-grpo" - -# master_sglang carries native SD3 /rollout/generate support; prepending it to PYTHONPATH -# shadows the editable install at /sgl-workspace/sglang. -MASTER_SGLANG_PYTHON = "/sgl-workspace/master_sglang/sglang/python" - - -@dataclass -class ScriptArgs(U.ExecuteTrainConfig): - num_rollout: int = 600 - data_dir: str = "/root/datasets" - debug_alignment: bool = False - extra_args: str = "" - - -def prepare(args: ScriptArgs) -> str: - local_dir = U.hf_download_dataset(DATASET, include=f"{DATASET_SUBSET}/**", data_dir=args.data_dir) - return f"{local_dir}/{DATASET_SUBSET}" - - -def execute(args: ScriptArgs, data_dir: str) -> None: - run_name = f"diffusion_grpo_sd3_hps_sglang_{U.create_run_id()}" - - ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt " - - rollout_args = ( - "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout " - f"--prompt-data {data_dir}/train.jsonl " - "--input-key input " - "--rollout-batch-size 8 " - "--n-samples-per-prompt 16 " - f"--num-rollout {args.num_rollout} " - "--global-batch-size 64 " - "--rollout-microgroup-size 8 " - "--train-dp-split-mode stride " - "--diffusion-num-steps 10 " - "--diffusion-guidance-scale 4.5 " - "--diffusion-negative-prompt ' ' " - "--diffusion-noise-level 0.7 " - "--diffusion-height 512 " - "--diffusion-width 512 " - "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window " - "--diffusion-num-sde-steps 10 " - "--diffusion-sde-window-range 0,10 " - ) - - eval_args = "--diffusion-eval-num-steps 40 " - - grpo_args = "--advantage-estimator grpo --diffusion-clip-range 1e-4 --diffusion-kl-beta 0.01 " - - optimizer_args = "--lr 3e-4 --adam-beta2 0.999 --weight-decay 1e-4 " - - lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 32 --lora-alpha 64 --lora-init-weights gaussian " - - reward_args = ( - "--rm-type hps " "--hps-num-workers 1 " "--hps-batch-size 8 " "--hps-version v2.1 " "--hps-reward-colocate " - ) - - wandb_args = U.get_default_wandb_args( - __file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=8, wandb_log_image_interval=10 - ) - - sglang_args = ( - "--use-miles-router " - "--sglang-server-concurrency 8 " - "--sglang-dit-precision fp16 " - "--sglang-vae-slicing " - "--update-weight-buffer-size 2147483648 " - ) - - train_backend_args = "--train-backend fsdp --diffusion-forward-dtype fp16 " - - perf_args = "--gradient-checkpointing --micro-batch-size-sample 16 --micro-batch-size-tstep 5 " - - misc_args = ( - "--actor-num-gpus-per-node 2 " - "--rollout-num-gpus 2 " - "--rollout-num-gpus-per-engine 1 " - "--num-gpus-per-node 2 " - "--colocate " - "--deterministic-mode " - ) + ("--diffusion-debug-mode --debug-skip-optimizer-step " if args.debug_alignment else "") - - U.execute_train( - train_args=( - f"{ckpt_args} {rollout_args} {eval_args} {grpo_args} {optimizer_args} " - f"{lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} {perf_args} " - f"{misc_args} {args.extra_args}" - ), - num_gpus_per_node=2, - config=args, - extra_env_vars={ - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "PYTHONPATH": MASTER_SGLANG_PYTHON, - "HF_TOKEN": os.environ.get("HF_TOKEN", ""), - **({"MILES_VERIFY_WEIGHT_SYNC": "1"} if args.debug_alignment else {}), - }, - ) - - -@U.dataclass_cli -def main(args: ScriptArgs) -> None: - data_dir = prepare(args) - execute(args, data_dir) - - -if __name__ == "__main__": - typer.run(main) diff --git a/scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py b/scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py index f8c9ecbc4..27e411fa7 100644 --- a/scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py +++ b/scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py @@ -112,7 +112,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: optimizer_args = "--lr 1e-5 --adam-beta2 0.999 --weight-decay 1e-4 " reward_placement = ( - "--pickscore-reward-colocate --pickscore-num-workers 4 " + "--colocate-reward --pickscore-num-workers 4 " if args.four_gpu_ci else "--pickscore-num-workers 4 --pickscore-num-gpus-per-worker 0.25 " ) diff --git a/tests/fast-gpu/test_hpsv2_alignment.py b/tests/fast-gpu/test_hpsv2_alignment.py deleted file mode 100644 index c232df8a2..000000000 --- a/tests/fast-gpu/test_hpsv2_alignment.py +++ /dev/null @@ -1,80 +0,0 @@ -from tests.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=180, - suite="stage-b-3-gpu-h200", - labels=[], -) - -import numpy as np -import torch -from hpsv2.src.open_clip import create_model_and_transforms, get_tokenizer -from huggingface_hub import hf_hub_download -from PIL import Image - -from miles.rollout.rm_hub.hps import HPSScorer - - -def _make_image(height: int, width: int, offset: int) -> Image.Image: - y, x = np.indices((height, width)) - rgb = np.empty((height, width, 3), dtype=np.uint8) - rgb[..., 0] = (3 * x + y + offset) % 256 - rgb[..., 1] = (x + 5 * y + 2 * offset) % 256 - rgb[..., 2] = (7 * x + 11 * y + 3 * offset) % 256 - return Image.fromarray(rgb) - - -def _official_hpsv2(checkpoint_path: str, prompts: list[str], images: list[Image.Image]): - # The strict HPS checkpoint load replaces every parameter, so skip the official scorer's redundant LAION preload. - model, _, preprocess = create_model_and_transforms( - "ViT-H-14", - pretrained=None, - precision="amp", - device="cuda", - jit=False, - force_quick_gelu=False, - force_custom_text=False, - force_patch_dropout=False, - force_image_size=None, - pretrained_image=False, - image_mean=None, - image_std=None, - light_augmentation=True, - aug_cfg={}, - output_dict=True, - with_score_predictor=False, - with_region_predictor=False, - ) - checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) - model.load_state_dict(checkpoint["state_dict"]) - model.eval() - - # The public scorer loops over pairs; batching both paths identically exercises the paired-diagonal fast path. - image_batch = torch.stack([preprocess(image) for image in images]) - text_batch = get_tokenizer("ViT-H-14")(prompts) - with torch.no_grad(), torch.amp.autocast("cuda"): - outputs = model(image_batch.cuda(), text_batch.cuda()) - scores = torch.diagonal(outputs["image_features"] @ outputs["text_features"].T) - return image_batch, scores.float().cpu() - - -def test_hps_scorer_matches_official_hpsv2(): - checkpoint_path = hf_hub_download("xswu/HPSv2", "HPS_v2.1_compressed.pt") - images = [_make_image(192, 320, 17), _make_image(320, 192, 53)] - prompts = ["a colorful geometric landscape", "an abstract portrait with vivid lines"] - - scorer = HPSScorer(device="cuda", checkpoint_path=checkpoint_path) - actual_images = torch.stack([scorer.preprocess(image) for image in images]) - actual_scores = torch.tensor(scorer(prompts, images)) - expected_images, expected_scores = _official_hpsv2(checkpoint_path, prompts, images) - - assert torch.equal(actual_images, expected_images) - torch.testing.assert_close(actual_scores, expected_scores, rtol=0, atol=0) - - -if __name__ == "__main__": - import sys - - import pytest - - sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/fast/rollout/test_hps_reward.py b/tests/fast/rollout/test_hps_reward.py deleted file mode 100644 index 827923b98..000000000 --- a/tests/fast/rollout/test_hps_reward.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Rewards receive ``generated_output`` itself and quantise to uint8 inside each actor. - -Mental model (one sample scored by two rewards): - - sample.generated_output (float CFHW) - -> HPSRewardActor.score_batch : round -> uint8 -> HPSv2 - -> PickScoreRewardActor.score_batch : truncate -> uint8 -> mean over frames -> PickScore - -Covered: HPS rounds when quantising (1); PickScore truncates and returns one mean per sample (2); -hps_rm and pickscore_rm hand the raw tensor to their pool and record their own queue depth (3); -pickscore_rm keeps only --pickscore-num-frames frames before shipping (4); batched_async_rm takes -the HPS fast path (5) and keeps sample order for mixed rm_types (6). -""" - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[]) - -from argparse import Namespace -from unittest.mock import AsyncMock - -import numpy as np -import pytest -import torch - -import miles.rollout.rm_hub.hps as hps_module -import miles.rollout.rm_hub.pickscore as pickscore_module -from miles.rollout.rm_hub import batched_async_rm -from miles.rollout.rm_hub.hps import HPSRewardActor -from miles.rollout.rm_hub.pickscore import PickScoreRewardActor -from miles.utils.types import Sample - -# one pixel row [0, 0.5, 1] per channel; rounding gives 128 in the middle, truncation 127 -_ROW = torch.tensor([0.0, 0.5, 1.0]).reshape(1, 1, 1, 3).repeat(3, 1, 1, 1) - - -def _middle_pixel(images) -> list[float]: - return [float(np.asarray(image)[0, 1, 0]) for image in images] - - -def test_hps_actor_rounds_when_quantising(): - """Truncating would shift pixels one level below what the HPSv2 reference scores.""" - actor = HPSRewardActor.__new__(HPSRewardActor) - actor.scorer = lambda prompts, images: _middle_pixel(images) - - assert actor.score_batch([_ROW], ["prompt"]) == [128.0] - - -def test_pickscore_actor_truncates_and_averages_frames_per_sample(): - """PickScore keeps flow_grpo's truncation, a two-frame sample yields one score, and each forward sees - --pickscore-batch-size frames (the chunking the e2e standards were recorded with), not one sample.""" - actor = PickScoreRewardActor.__new__(PickScoreRewardActor) - actor.frames_per_forward = 1 - forward_sizes = [] - - def scorer(prompts, images): - forward_sizes.append(len(images)) - return [p + i for p, i in zip(prompts, _middle_pixel(images), strict=True)] - - actor.scorer = scorer - two_frames = torch.cat([_ROW, _ROW], dim=1) - - assert actor.score_batch([two_frames], [1000.0]) == [1127.0] - assert forward_sizes == [1, 1] - - -@pytest.mark.asyncio -async def test_rm_functions_hand_the_raw_tensor_to_their_pool(monkeypatch): - """Decoding belongs to the actor, and each pool records its own backlog instead of overwriting a shared one.""" - pool = AsyncMock() - pool.score.side_effect = [([1.0], 2), ([1.0], 1)] - monkeypatch.setattr(hps_module, "AsyncHPSPool", lambda args: pool) - monkeypatch.setattr(pickscore_module, "AsyncPickScorePool", lambda args: pool) - sample = Sample(prompt="prompt", generated_output=_ROW) - - assert await hps_module.hps_rm(Namespace(), [sample]) == [1.0] - assert await pickscore_module.pickscore_rm(Namespace(pickscore_num_frames=None), [sample]) == [1.0] - - calls = pool.score.await_args_list - assert [call.args[1] for call in calls] == [["prompt"]] * 2 - assert all(torch.equal(call.args[0][0], _ROW) for call in calls) - assert sample.reward_max_queue_depth == {"hps": 2.0, "pickscore": 1.0} - - -@pytest.mark.asyncio -async def test_pickscore_rm_selects_frames_before_shipping(monkeypatch): - """Only the --pickscore-num-frames frames should cross the object store, not the whole video.""" - pool = AsyncMock() - pool.score.return_value = ([1.0], 0) - monkeypatch.setattr(pickscore_module, "AsyncPickScorePool", lambda args: pool) - video = torch.cat([_ROW * 0.0, _ROW, _ROW * 0.0], dim=1) - - await pickscore_module.pickscore_rm( - Namespace(pickscore_num_frames=1), [Sample(prompt="p", generated_output=video)] - ) - - (shipped,), _ = pool.score.await_args.args - assert shipped.shape[1] == 1 and torch.equal(shipped[:, 0], _ROW[:, 0]) - - -@pytest.mark.asyncio -async def test_batched_async_rm_uses_hps_batch_fast_path(monkeypatch): - hps_rm = AsyncMock(return_value=[3.0, 1.0]) - monkeypatch.setattr(hps_module, "hps_rm", hps_rm) - args = Namespace(custom_rm_path=None, rm_type="hps") - samples = [Sample(index=3), Sample(index=1)] - - assert await batched_async_rm(args, samples) == [3.0, 1.0] - hps_rm.assert_awaited_once_with(args, samples) - - -@pytest.mark.asyncio -async def test_batched_async_rm_preserves_order_for_mixed_reward_types(monkeypatch): - monkeypatch.setattr(hps_module, "hps_rm", AsyncMock(return_value=[10.0])) - monkeypatch.setattr(pickscore_module, "pickscore_rm", AsyncMock(return_value=[20.0])) - args = Namespace(custom_rm_path=None, rm_type=None) - samples = [ - Sample(metadata={"rm_type": "hps"}), - Sample(metadata={"rm_type": "pickscore"}), - ] - - assert await batched_async_rm(args, samples) == [10.0, 20.0] diff --git a/tests/fast/rollout/test_reward_pool.py b/tests/fast/rollout/test_reward_pool.py deleted file mode 100644 index 9fa0fb13c..000000000 --- a/tests/fast/rollout/test_reward_pool.py +++ /dev/null @@ -1,124 +0,0 @@ -"""The colocated reward slot ledger: which bundle each reward actor gets, and that pools never overlap. - -Mental model (8 bundles = 2 nodes x 4 GPUs, 4-GPU engines; `x` = an engine's own claim): - - node 0 bundle 0[x] 1 2 3 bundle_deal_order deals non-`x` bundles first, alternating - node 1 bundle 4[x] 5 6 7 nodes, and the `x` bundles last -> [1, 5, 2, 6, 3, 7, 0, 4] - ColocatedRewardSlots(order).allocate("hps", 3) -> [1, 5, 2]; allocate("pickscore", 3) -> [6, 3, 7] - -Covered: the deal order for single-GPU (SD3) and 4-GPU (Wan) engines (1-2); pools get disjoint -contiguous slices (3); every bundle survives partial or information-poor placements (4-6); an -over-subscribing pool is rejected with the per-bundle owner map and nothing is consumed (7); a -pool cannot allocate twice (8). -""" - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) - -import pytest - -from miles.rollout.rm_hub.core import ColocatedRewardSlots, bundle_deal_order - - -def test_sd3_bundle_deal_order(): - assert bundle_deal_order( - list(range(8)), - [0, 1, 2, 3, 0, 1, 2, 3], - num_gpus_per_node=4, - num_gpus_per_engine=1, - ) == [0, 4, 1, 5, 2, 6, 3, 7] - - -def test_wan_bundle_deal_order(): - assert bundle_deal_order( - list(range(16)), - list(range(8)) * 2, - num_gpus_per_node=8, - num_gpus_per_engine=4, - ) == [1, 9, 5, 13, 2, 10, 6, 14, 3, 11, 7, 15, 0, 8, 4, 12] - - -def test_pools_receive_disjoint_contiguous_slices(): - slots = ColocatedRewardSlots([0, 4, 1, 5, 2, 6, 3, 7]) - - hps = slots.allocate("HPS", 3) - pickscore = slots.allocate("PickScore", 3) - - assert hps == [0, 4, 1] - assert pickscore == [5, 2, 6] - assert set(hps).isdisjoint(pickscore) - - -@pytest.mark.parametrize( - ("bundle_indices", "gpu_ids", "num_gpus_per_node"), - [ - ([0, 1, 2, 3], [0, 1, 2, 3], 5), - ([7, 9], [0, 2], 3), - ], -) -def test_partial_placement_group_retains_every_bundle(bundle_indices, gpu_ids, num_gpus_per_node): - order = bundle_deal_order( - bundle_indices, - gpu_ids, - num_gpus_per_node=num_gpus_per_node, - num_gpus_per_engine=2, - ) - - assert len(order) == len(bundle_indices) - assert set(order) == set(bundle_indices) - - -def test_incomplete_engine_span_is_preferred_to_base_bundles(): - order = bundle_deal_order( - list(range(5)), - list(range(5)), - num_gpus_per_node=5, - num_gpus_per_engine=2, - ) - - assert order == [4, 1, 3, 0, 2] - - -@pytest.mark.parametrize( - "gpu_ids", - [ - [0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3], - [0, 1, 2, 3, 4, 5, 6, 7], - ], -) -def test_information_poor_placements_retain_every_bundle(gpu_ids): - bundle_indices = list(range(len(gpu_ids))) - - order = bundle_deal_order( - bundle_indices, - gpu_ids, - num_gpus_per_node=8 if len(gpu_ids) > 8 else 4, - num_gpus_per_engine=4, - ) - - assert len(order) == len(bundle_indices) - assert set(order) == set(bundle_indices) - - -def test_exhaustion_is_atomic_and_reports_ownership(): - slots = ColocatedRewardSlots([0, 1, 2, 3]) - slots.allocate("HPS", 2) - slots.allocate("PickScore", 1) - - with pytest.raises( - RuntimeError, match=r"only 1/4 remain \(bundle 0: HPS, bundle 1: HPS, bundle 2: PickScore, bundle 3: None\)" - ): - slots.allocate("Other", 2) - - assert slots.allocate("Last", 1) == [3] - - -def test_duplicate_pool_name_does_not_consume_a_slice(): - slots = ColocatedRewardSlots([0, 1, 2]) - slots.allocate("HPS", 1) - - with pytest.raises(RuntimeError, match="HPS already owns reward slots"): - slots.allocate("HPS", 1) - - assert slots.allocate("PickScore", 1) == [1] diff --git a/tests/fast/rollout/test_reward_pool_placement.py b/tests/fast/rollout/test_reward_pool_placement.py deleted file mode 100644 index 52151492a..000000000 --- a/tests/fast/rollout/test_reward_pool_placement.py +++ /dev/null @@ -1,74 +0,0 @@ -"""A colocated reward pool takes one slot per rollout bundle and only RolloutManager may seat it. - -Mental model (2 rollout bundles, --hps-reward-colocate --hps-num-workers 2): - - bundle 0 [ train | engine | reward slot ] <- worker 0, PlacementGroupSchedulingStrategy - bundle 1 [ train | engine | reward slot ] <- worker 1 - -Covered: each worker lands on its own bundle at the colocated GPU share (1); a colocated -pool built outside RolloutManager, i.e. without seats, is rejected before any actor is -created (2). Over-subscription is covered by test_reward_pool.py. -""" - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) - -import pytest - -import miles.rollout.rm_hub.core as core_module -from miles.ray.utils import COLOCATED_REWARD_GPU -from miles.rollout.rm_hub.core import AsyncRewardActorPool, ColocatedRewardSlots - - -class _FakeActorCls: - """Records the options each worker was created with instead of starting Ray actors.""" - - def __init__(self): - self.created = [] - - def options(self, **options): - self.created.append(options) - return self - - def remote(self, **kwargs): - return object() - - -@pytest.fixture(autouse=True) -def _no_ray(monkeypatch): - monkeypatch.setattr(core_module.ray, "remote", lambda actor_cls: actor_cls) - - -def _pool(actor_cls, *, num_workers, colocate, placement_group=None, slots=None): - return AsyncRewardActorPool( - actor_cls=actor_cls, - actor_kwargs={}, - num_workers=num_workers, - batch_size=8, - num_gpus_per_worker=1.0, - colocate=colocate, - name="hps", - placement_group=placement_group, - slots=slots, - ) - - -def test_colocated_workers_take_one_slot_each_at_the_colocated_share(): - """Two workers on one bundle, or a full-GPU claim on a bundle, would pend in Ray forever.""" - actor_cls = _FakeActorCls() - slots = ColocatedRewardSlots([0, 1]) - - _pool(actor_cls, num_workers=2, colocate=True, placement_group=("pg", [0, 1], [0, 1]), slots=slots) - - assert [o["scheduling_strategy"].placement_group_bundle_index for o in actor_cls.created] == [0, 1] - assert [o["num_gpus"] for o in actor_cls.created] == [COLOCATED_REWARD_GPU] * 2 - assert slots.remaining == 0 - - -def test_colocated_pool_built_without_seats_is_rejected(): - """Only RolloutManager seats colocated pools; a lazy build means the flag was set for an unseated type.""" - actor_cls = _FakeActorCls() - with pytest.raises(RuntimeError, match="not seated by RolloutManager"): - _pool(actor_cls, num_workers=1, colocate=True) - assert actor_cls.created == [] diff --git a/tests/fast/rollout/test_weighted_mixture_rm.py b/tests/fast/rollout/test_weighted_mixture_rm.py deleted file mode 100644 index 8e8b5de04..000000000 --- a/tests/fast/rollout/test_weighted_mixture_rm.py +++ /dev/null @@ -1,63 +0,0 @@ -"""weighted_mixture_rm combines built-in rewards with weights parsed from --custom-rm-args. - -Mental model (--custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weighted, one batch of 2 samples): - - weighted_mixture_rm ─┬─ hps_rm(args, [s0, s1]) -> [0.3, 0.2] x 0.7 - └─ pickscore_rm(args, [s0, s1]) -> [0.8, 0.9] x 0.3 - = [{hps: 0.3, pickscore: 0.8, weighted: 0.45}, {hps: 0.2, pickscore: 0.9, weighted: 0.41}] - -Covered: each reward scores the whole batch once and every sample gets its components plus the -weighted sum (1); an unknown reward name in --custom-rm-args is rejected (2); a --reward-key -that names neither a component nor "weighted" is rejected before any reward runs (3). -""" - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) - -from argparse import Namespace - -import pytest - -import miles.rollout.rm_hub.weighted_mixture_rm as weighted_mixture_rm_module -from miles.rollout.rm_hub.weighted_mixture_rm import parse_weights, weighted_mixture_rm - - -def _fake_rewards(calls): - def fake(scores): - async def rm(args, samples): - calls.append(len(samples)) - return scores - - return rm - - return {"hps": fake([0.3, 0.2]), "pickscore": fake([0.8, 0.9])} - - -@pytest.mark.asyncio -async def test_each_sample_gets_its_components_and_the_weighted_sum(monkeypatch): - """Fanning the batch out per sample, dropping a weight, or collapsing to a scalar would all show here.""" - calls = [] - monkeypatch.setattr(weighted_mixture_rm_module, "_REWARDS", _fake_rewards(calls)) - args = Namespace(custom_rm_args="hps=0.7,pickscore=0.3", reward_key="weighted") - - rewards = await weighted_mixture_rm(args, [object(), object()]) - - assert [r["weighted"] for r in rewards] == pytest.approx([0.45, 0.41]) - assert [(r["hps"], r["pickscore"]) for r in rewards] == [(0.3, 0.8), (0.2, 0.9)] - assert calls == [2, 2] - - -def test_unknown_reward_name_is_rejected(): - with pytest.raises(ValueError, match="unknown reward 'clip'"): - parse_weights("hps=0.7,clip=0.3") - - -@pytest.mark.asyncio -async def test_missing_reward_key_is_rejected_before_scoring(monkeypatch): - calls = [] - monkeypatch.setattr(weighted_mixture_rm_module, "_REWARDS", _fake_rewards(calls)) - - with pytest.raises(ValueError, match="--reward-key weighted"): - await weighted_mixture_rm(Namespace(custom_rm_args="hps=0.7,pickscore=0.3", reward_key=None), [object()]) - assert calls == [] diff --git a/tests/fast/utils/test_processing_utils.py b/tests/fast/utils/test_processing_utils.py index 41d541710..d695cdf4f 100644 --- a/tests/fast/utils/test_processing_utils.py +++ b/tests/fast/utils/test_processing_utils.py @@ -2,16 +2,10 @@ register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) -import numpy as np import pytest import torch -from miles.utils.processing_utils import ( - cfhw_to_fhwc, - fhwc_to_cfhw, - generated_output_to_rgb_hwc_uint8_frames, - image_or_video_to_uint8, -) +from miles.utils.processing_utils import cfhw_to_fhwc, fhwc_to_cfhw, image_or_video_to_uint8 def test_cfhw_fhwc_round_trip(): @@ -46,20 +40,3 @@ def test_image_or_video_to_uint8_does_not_round_pixel_values(): actual = image_or_video_to_uint8(torch.tensor([0.0, 127.9, 255.0]), round_normalized=True) torch.testing.assert_close(actual, torch.tensor([0, 127, 255], dtype=torch.uint8)) - - -def test_generated_output_to_rgb_hwc_uint8_frames_selects_before_conversion(): - frames = torch.tensor( - [ - [[0.0, 0.0, 0.0]], - [[0.0, 0.5, 1.0]], - [[1.0, 1.0, 1.0]], - ] - ) - - actual = generated_output_to_rgb_hwc_uint8_frames(frames.unsqueeze(0).repeat(3, 1, 1, 1), 1) - - expected = np.array([[[0, 0, 0], [127, 127, 127], [255, 255, 255]]], dtype=np.uint8) - assert len(actual) == 1 - np.testing.assert_array_equal(actual[0], expected) - assert actual[0].flags.c_contiguous