diff --git a/.github/workflows/_run-ci.yml b/.github/workflows/_run-ci.yml index e96d43942..3e4113c24 100644 --- a/.github/workflows/_run-ci.yml +++ b/.github/workflows/_run-ci.yml @@ -62,6 +62,10 @@ jobs: # (e.g. the label name '["h200", "2gpu"]' with brackets), which no # runner carries, and the job queues forever. runs-on: ${{ fromJSON(inputs.runs_on) }} + # The CI node's RDMA devices are RoCE, and a RoCE connection needs the device's network interface, which the + # runner's bridge network does not have. `ci-roce` is an ipvlan network on bond0 that adds it. Created once on + # the node; without it the container is not created: + # docker network create -d ipvlan --internal --subnet 192.168.250.0/24 -o parent=bond0 ci-roce container: image: ${{ inputs.container_image }} options: > @@ -80,6 +84,7 @@ jobs: -v /data/miles_ci/paddleocr:/github/home/.paddleocr --privileged --ulimit nofile=65535:65535 + --network ci-roce defaults: run: working-directory: ${{ github.workspace }} @@ -90,6 +95,9 @@ jobs: # Gated models (e.g. stabilityai/stable-diffusion-3.5-medium for the SD3 # e2e) need an HF token with repo access; unset for public-model tests. HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Mooncake's whitelist of RDMA devices: only bond0's is reachable through ci-roce, and a transfer that + # tries any of the other seven fails. + MC_MS_FILTERS: mlx5_bond_0 steps: - name: Checkout repository uses: actions/checkout@v4 diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 289adbce5..59e60ce95 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -319,7 +319,7 @@ jobs: --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }} ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || needs.ci-gate.outputs.docs_only == 'true' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }} ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }} - ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} + ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }} secrets: inherit # Stage B: CPU-only slower bucket (currently empty; reserved for future @@ -337,7 +337,7 @@ jobs: --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }} ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || needs.ci-gate.outputs.docs_only == 'true' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }} ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }} - ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} + ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }} secrets: inherit # Stage B GPU: fast-gpu bucket, runs once stage-a-cpu is green (fast-fail: @@ -364,7 +364,7 @@ jobs: --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }} ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }} ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }} - ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} + ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }} secrets: inherit # Stage B GPU: 4-GPU distributed tests on the 5gpu runner. @@ -387,7 +387,7 @@ jobs: --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }} ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }} ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }} - ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} + ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }} secrets: inherit # Fan-in so each stage-c needs one job instead of every stage-b. @@ -422,7 +422,7 @@ jobs: --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }} ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }} ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }} - ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} + ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }} secrets: inherit # Stage C GPU: e2e suite on the 5gpu runner. @@ -445,5 +445,5 @@ jobs: --labels ${{ join(github.event.pull_request.labels.*.name, ' ') }} ${{ (github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || contains(github.event.pull_request.labels.*.name, 'run-ci-image') || contains(github.event.pull_request.labels.*.name, 'run-ci-all')) && '--match-all-labels' || '' }} ${{ (github.event_name == 'schedule' || inputs.nightly) && '--nightly' || '' }} - ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} + ${{ (github.event_name == 'schedule' || github.event_name == 'pull_request' || inputs.continue_on_error) && '--continue-on-error' || '' }} secrets: inherit diff --git a/docs/advanced/streaming-reward.md b/docs/advanced/streaming-reward.md index 9be270d59..c31bcffe9 100644 --- a/docs/advanced/streaming-reward.md +++ b/docs/advanced/streaming-reward.md @@ -21,9 +21,8 @@ flowchart LR The rollout pipeline decomposes into `generate` / `deserialize` / `reward` components. With `--use-output-store`, the -output is read back from the store and deleted there between `deserialize` and `reward`, still inside the generation -slot. Because every microgroup is its own `asyncio` task, the engine starts the next request while previous responses -are still being unpacked and scored. +reward actors read the output from the store, and it is deleted there once scored. Because every microgroup is its +own `asyncio` task, the engine starts the next request while previous responses are still being unpacked and scored. Nothing waits for the full batch. ## 2. Deserialization: msgpack + parser-actor pool diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md index d4eed2185..e81fb2a69 100644 --- a/docs/user-guide/customization.md +++ b/docs/user-guide/customization.md @@ -57,6 +57,10 @@ def generate_rollout( Legacy list/dict returns are still wrapped into the corresponding output dataclass by `miles.rollout.base_types.call_rollout_fn`. +With `--use-output-store`, a rollout function that does not go through `generate_rollout_async` sets +`GenerateState(args).media_outputs_left` to the number of outputs to keep for its media logging at the start of each +rollout, and calls `release_generated_outputs` after a reward it runs itself. + ### `--eval-function-path` Same shape as the rollout function. Defaults to `--rollout-function-path` when @@ -92,10 +96,8 @@ async def custom_generate( `evaluation=` is optional — if your signature omits it, the caller skips that kwarg. -With `--use-output-store`, `generate_microgroup` leaves `sample.generated_output` as a `StoredTensor`. -`generate_and_rm_microgroup` reads it into the tensor after your function returns and before the reward. Code that -reads the output earlier, such as a custom generate function that refines its own first pass, calls -`await load_generated_outputs(GenerateState(args), microgroup)` first. +With `--use-output-store`, `generate_microgroup` leaves `sample.generated_output` as a `StoredTensor`: the output +stays in the store for the reward actors, and a custom generate function cannot read it. ### `--diffusion-step-strategy-path` @@ -134,6 +136,13 @@ Built-in scorers (`--rm-type pickscore` / `hps` / `ocr` / `custom_api` / `openai [Rewards](rewards.md). `--custom-rm-path` replaces that dispatch entirely; `--custom-reward-post-process-path` replaces the normalization that runs after scoring. +With `--use-output-store`, the rollout process does not read a training rollout's output. There, +`sample.generated_output` is a `StoredTensor`, and the actors of a reward pool read it when they score it. A +`--custom-rm-path` function therefore scores through a pool, a built-in `*_rm` or its own `AsyncRewardActorPool`, and +does not read the tensor itself. Once a sample is scored its output is `None`; the first `--wandb-log-num-images` +scored samples of a rollout that logs media keep it as a tensor, so a rollout whose first groups a filter drops logs +fewer. Every hook that runs after scoring sees that. + ### `--custom-rm-path` ```python @@ -183,6 +192,9 @@ async def api_rm(args, samples: list[Sample], **kwargs) -> list[float]: --rm-url http://localhost:8000/score ``` +With `--use-output-store` this function cannot read the output. Send it through `--custom-api-rm-config`, whose +`ApiRewardActor` reads it in the actor, or through your own `AsyncRewardActorPool(args=args, ...)`. + ### `--custom-reward-post-process-path` Replace GRPO advantage normalization in diff --git a/miles/dashboard/events.py b/miles/dashboard/events.py index a125637b3..f5559f619 100644 --- a/miles/dashboard/events.py +++ b/miles/dashboard/events.py @@ -33,8 +33,6 @@ class TrajectoryEventKind(StrEnum): DESER_END = "deser_end" REWARD_START = "reward_start" REWARD_END = "reward_end" - STORE_READ_START = "store_read_start" - STORE_READ_END = "store_read_end" # stage name (StageTimer) -> (start kind, end kind) @@ -42,7 +40,6 @@ class TrajectoryEventKind(StrEnum): "generate": (TrajectoryEventKind.GEN_START, TrajectoryEventKind.GEN_END), "deserialize": (TrajectoryEventKind.DESER_START, TrajectoryEventKind.DESER_END), "reward": (TrajectoryEventKind.REWARD_START, TrajectoryEventKind.REWARD_END), - "store_read": (TrajectoryEventKind.STORE_READ_START, TrajectoryEventKind.STORE_READ_END), } # kind -> (base span, is_start), for the reader to fold start/end back into spans @@ -53,8 +50,6 @@ class TrajectoryEventKind(StrEnum): "deser_end": ("deser", False), "reward_start": ("reward", True), "reward_end": ("reward", False), - "store_read_start": ("store_read", True), - "store_read_end": ("store_read", False), } diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 2407e509d..f2fec6717 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -449,12 +449,15 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl reward_stats["rollout/step"] = compute_rollout_step(self.args, self.rollout_id) tracking_utils.log(self.args, reward_stats, step_key="rollout/step") - max_images = self.args.wandb_log_num_images - interval = self.args.wandb_log_image_interval - if max_images > 0 and self.rollout_id % interval == 0: + if tracking_utils.logs_rollout_media(self.args, self.rollout_id): + # with the output store on, only the samples the rollout kept for media hold a tensor self._log_images_async( - {"rollout_media/sample_images": samples}, - max_images=max_images, + { + "rollout_media/sample_images": [ + sample for sample in samples if isinstance(sample.generated_output, torch.Tensor) + ] + }, + max_images=self.args.wandb_log_num_images, step_key="rollout/step", step_value=compute_rollout_step(self.args, self.rollout_id), reward_key=self.args.reward_key, diff --git a/miles/rollout/rm_hub/api.py b/miles/rollout/rm_hub/api.py index 6a346d09d..fae659c2f 100644 --- a/miles/rollout/rm_hub/api.py +++ b/miles/rollout/rm_hub/api.py @@ -66,6 +66,7 @@ def __init__(self, args) -> None: if not isinstance(actor_cls, type) or not issubclass(actor_cls, self.actor_base_cls): raise TypeError(f"API reward actor_class must be an {self.actor_base_cls.__name__} subclass") super().__init__( + args=args, actor_cls=actor_cls, actor_kwargs=config.actor_kwargs, num_workers=1, diff --git a/miles/rollout/rm_hub/core.py b/miles/rollout/rm_hub/core.py index 066129447..0ee754f5a 100644 --- a/miles/rollout/rm_hub/core.py +++ b/miles/rollout/rm_hub/core.py @@ -9,6 +9,8 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from miles.ray.utils import COLOCATED_REWARD_GPU +from miles.utils.diffusion_rollout_response import read_generated_outputs +from miles.utils.output_store import RegisteredTensorReader, StoredTensor, connect_output_store logger = logging.getLogger(__name__) @@ -79,10 +81,27 @@ def record_reward_queue_depth(samples, name: str, max_queue_depth: int) -> None: sample.reward_max_queue_depth = depths +class _StoredOutputRewardActor: + """A reward actor that reads the stored outputs it is given from the output store, then scores them.""" + + def __init__(self, *, scorer_cls, scorer_kwargs: dict, args) -> None: + self._scorer = scorer_cls(**scorer_kwargs) + self._reader = RegisteredTensorReader(connect_output_store(args)) + + def score_batch(self, outputs: list, prompts: list[str]) -> list[float]: + # eval outputs arrive inline + stored_outputs = [output for output in outputs if isinstance(output, StoredTensor)] + if stored_outputs: + read_outputs = iter(read_generated_outputs(self._reader, stored_outputs)) + outputs = [next(read_outputs) if isinstance(output, StoredTensor) else output for output in outputs] + return self._scorer.score_batch(outputs, prompts) + + class AsyncRewardActorPool: """Round-robin pool for Ray reward actors exposing ``score_batch(outputs, prompts)``. - Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves. + Actors receive the float ``generated_output`` tensors and quantise to uint8 themselves. With the output store + on they are given stored outputs and read them first. 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. @@ -91,6 +110,7 @@ class AsyncRewardActorPool: def __init__( self, *, + args, actor_cls, actor_kwargs: dict, num_workers: int, @@ -113,6 +133,9 @@ def __init__( num_gpus_per_worker = COLOCATED_REWARD_GPU else: strategies = ["DEFAULT"] * num_workers + if args.use_output_store: + actor_kwargs = {"scorer_cls": actor_cls, "scorer_kwargs": actor_kwargs, "args": args} + actor_cls = _StoredOutputRewardActor self._actors = [ ray.remote(actor_cls) diff --git a/miles/rollout/rm_hub/hps.py b/miles/rollout/rm_hub/hps.py index 0311f102f..0dcbccf36 100644 --- a/miles/rollout/rm_hub/hps.py +++ b/miles/rollout/rm_hub/hps.py @@ -128,6 +128,7 @@ class AsyncHPSPool(AsyncRewardActorPool, metaclass=SingletonMeta): def __init__(self, args, placement_group=None, slots=None) -> None: super().__init__( + args=args, actor_cls=HPSRewardActor, actor_kwargs={ "hps_version": args.hps_version, diff --git a/miles/rollout/rm_hub/hpsv3.py b/miles/rollout/rm_hub/hpsv3.py index 3c089319d..1e457d645 100644 --- a/miles/rollout/rm_hub/hpsv3.py +++ b/miles/rollout/rm_hub/hpsv3.py @@ -233,6 +233,7 @@ class AsyncHPSv3Pool(AsyncRewardActorPool, metaclass=SingletonMeta): def __init__(self, args, placement_group=None, slots=None) -> None: super().__init__( + args=args, actor_cls=HPSv3RewardActor, actor_kwargs={"checkpoint_path": args.hpsv3_checkpoint_path}, num_workers=args.hpsv3_num_workers, diff --git a/miles/rollout/rm_hub/ocr.py b/miles/rollout/rm_hub/ocr.py index 35403befc..46e8b8dc9 100644 --- a/miles/rollout/rm_hub/ocr.py +++ b/miles/rollout/rm_hub/ocr.py @@ -99,6 +99,7 @@ class AsyncOcrPool(AsyncRewardActorPool, metaclass=SingletonMeta): def __init__(self, args) -> None: super().__init__( + args=args, actor_cls=OcrRewardActor, actor_kwargs={"use_gpu": False}, num_workers=args.ocr_num_workers, diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index b58d99f2a..a4e971c03 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -7,6 +7,7 @@ from PIL import Image from miles.utils.misc import SingletonMeta +from miles.utils.output_store import StoredTensor from miles.utils.processing_utils import generated_output_to_rgb_hwc_uint8_frames, sample_frame_indices from miles.utils.types import Sample @@ -74,8 +75,10 @@ def __init__( processor_path: str, model_path: str, frames_per_forward: int, + num_frames: int | None, ) -> None: self.frames_per_forward = frames_per_forward + self.num_frames = num_frames gpu_ids = ray.get_gpu_ids() use_cuda = bool(gpu_ids) and torch.cuda.is_available() if use_cuda: @@ -88,10 +91,10 @@ def __init__( ) 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 + # a video is scored on --pickscore-num-frames of its frames, averaged; already picked frames pass unchanged images, frame_counts = [], [] for output in outputs: - frames = generated_output_to_rgb_hwc_uint8_frames(output, None) + frames = generated_output_to_rgb_hwc_uint8_frames(output, self.num_frames) 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)] @@ -112,11 +115,13 @@ class AsyncPickScorePool(AsyncRewardActorPool, metaclass=SingletonMeta): def __init__(self, args, placement_group=None, slots=None) -> None: super().__init__( + args=args, 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_frames": args.pickscore_num_frames, }, num_workers=args.pickscore_num_workers, batch_size=args.pickscore_batch_size, @@ -130,10 +135,15 @@ def __init__(self, args, placement_group=None, slots=None) -> None: 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 + # pick --pickscore-num-frames here so only the scored frames cross the object store; a stored output is not + # here to pick from, so its actor picks them outputs = [ - s.generated_output[:, sample_frame_indices(s.generated_output.shape[1], args.pickscore_num_frames)] - for s in samples + ( + output + if isinstance(output, StoredTensor) + else output[:, sample_frame_indices(output.shape[1], args.pickscore_num_frames)] + ) + for output in (s.generated_output 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) diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index ee78f844a..0f380cf01 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -22,7 +22,8 @@ from miles.utils.eval_config import EvalDatasetConfig from miles.utils.http_utils import post from miles.utils.misc import SingletonMeta, load_function -from miles.utils.output_store import StoredTensor, connect_output_store, record_dropped_samples +from miles.utils.output_store import RegisteredTensorReader, StoredTensor, connect_output_store, record_dropped_samples +from miles.utils.tracking_utils import logs_rollout_media from miles.utils.types import Sample from .rm_hub import batched_async_rm @@ -132,6 +133,9 @@ def __init__(self, args: Namespace) -> None: self._parser_rr = 0 self.parser_inflight = [0] * args.rollout_parser_num_workers self.output_store = connect_output_store(args) if args.use_output_store else None + self.output_reader = RegisteredTensorReader(self.output_store) if args.use_output_store else None + # outputs of the running rollout still to keep for its media logging + self.media_outputs_left = 0 self.reset() @@ -249,33 +253,44 @@ async def generate_microgroup( return microgroup -async def load_generated_outputs(state: GenerateState, samples: list[Sample]) -> None: - """Read the samples' stored outputs into tensors and delete them from the store. +async def release_generated_outputs(state: GenerateState, samples: list[Sample]) -> None: + """Delete the stored outputs of scored samples from the store, first reading the few media logging keeps. - The store then holds an output only from the engine's put to this read. The rollout's release deletes whatever a - failed delete leaves. + The samples' outputs become ``None``, except the kept ones, which become tensors. The rollout's release deletes + whatever a failed delete leaves. """ - stored = [sample for sample in samples if isinstance(sample.generated_output, StoredTensor)] - if not stored: + stored_samples = [sample for sample in samples if isinstance(sample.generated_output, StoredTensor)] + if not stored_samples: return - outputs = [sample.generated_output for sample in stored] - refs = [ - msgspec.structs.replace( - output.ref, fields=tuple(spec for spec in output.ref.fields if spec.name == output.field) - ) - for output in outputs - ] - st = hooks.StageTimer() - with st.stage("store_read"): - tensors = await asyncio.to_thread(read_generated_outputs, state.output_store, outputs) + # taken before the read is awaited, so concurrent microgroups keep at most --wandb-log-num-images in all + media_samples = stored_samples[: state.media_outputs_left] + state.media_outputs_left -= len(media_samples) + media_tensors = [] + if media_samples: + media_outputs = [sample.generated_output for sample in media_samples] try: - await asyncio.to_thread(state.output_store.delete_many, refs) + media_tensors = await asyncio.to_thread(read_generated_outputs, state.output_reader, media_outputs) except Exception: logger.warning( - "Failed to delete %d read outputs; the rollout's release deletes them", len(refs), exc_info=True + "Failed to read %d outputs kept for media; they are not logged", len(media_samples), exc_info=True ) - st.attach(stored) - for sample, tensor in zip(stored, tensors, strict=True): + media_samples = [] + output_refs = [ + msgspec.structs.replace( + stored_output.ref, + fields=tuple(spec for spec in stored_output.ref.fields if spec.name == stored_output.field), + ) + for stored_output in (sample.generated_output for sample in stored_samples) + ] + try: + await asyncio.to_thread(state.output_store.delete_many, output_refs) + except Exception: + logger.warning( + "Failed to delete %d scored outputs; the rollout's release deletes them", len(output_refs), exc_info=True + ) + for sample in stored_samples: + sample.generated_output = None + for sample, tensor in zip(media_samples, media_tensors, strict=True): sample.generated_output = tensor @@ -299,8 +314,6 @@ async def generate_and_rm_microgroup( microgroup = await custom_generate_func(args, microgroup, sampling_params) else: microgroup = await generate_microgroup(args, microgroup, sampling_params, evaluation=evaluation) - # inside the semaphore, so the store holds at most one microgroup's outputs per generation slot - await load_generated_outputs(state, microgroup) # for the rm that need the whole group, we will not do the rm here if args.group_rm: @@ -314,13 +327,14 @@ async def generate_and_rm_microgroup( for sample, reward in zip(microgroup, rewards, strict=True): sample.reward = reward hooks.record_trajectory(sample) + await release_generated_outputs(state, microgroup) return microgroup async def generate_and_rm_group( args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False ) -> list[Sample]: - GenerateState(args) + state = GenerateState(args) # N-spaced base so sgl-d's seed→[seed+0..seed+N-1] expansion stays disjoint # per (rollout, prompt-group); group_index is monotonic across the run. @@ -347,6 +361,7 @@ async def generate_and_rm_group( rewards = await batched_async_rm(args, group) for sample, reward in zip(group, rewards, strict=False): sample.reward = reward + await release_generated_outputs(state, group) return group @@ -369,6 +384,7 @@ async def generate_rollout_async( assert args.rollout_global_dataset state = GenerateState(args) + state.media_outputs_left = args.wandb_log_num_images if logs_rollout_media(args, rollout_id) else 0 # instantiate data filters dynamic_filter = ( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index d943d311c..86bf607b4 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -591,8 +591,11 @@ def add_rollout_arguments(parser): default=False, help=( "Rollout requests put their latents, debug and cond tensors and their generated output in " - "the run's Mooncake (sglang-d --output-store-backend) and return a ref; the rollout reads the " - "output back right after generation and the trainer reads the tensors." + "the run's Mooncake (sglang-d --output-store-backend) and return a ref; the reward actors read " + "the output and the trainer reads the tensors. In the rollout process a sample's " + "generated_output is a StoredTensor until its reward and None after it, except the first " + "--wandb-log-num-images scored samples of a rollout that logs media, which become tensors. " + "Outputs wait in the store until scored; see --output-store-global-segment-size." ), ) parser.add_argument( @@ -616,18 +619,24 @@ def add_rollout_arguments(parser): parser.add_argument( "--output-store-global-segment-size", type=str, - default="16gb", + default="72gb", help=( "Memory each training node lends the store, mounted by its first trainer actor; the run's " - "stored objects live in these segments, so their sum bounds what a rollout can store." + "stored objects live in these segments, so their sum bounds what a rollout can store. Outputs " + "wait there until scored: the segments should hold one rollout's outputs plus two rollouts of " + "latents, or a reward slower than generation fills the store and requests fail. The default " + "holds an LTX-2.3 rollout of 64 float videos." ), ) parser.add_argument( "--output-store-protocol", type=str, choices=("rdma", "tcp"), - default="tcp", - help="Mooncake transfer protocol of every output store client; rdma needs an RDMA device on every node.", + default="rdma", + help=( + "Mooncake transfer protocol of every output store client. rdma needs an RDMA device on every " + "node; tcp works without one and is much slower, above all for puts." + ), ) parser.add_argument( "--output-store-device-name", diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index ff65a9c23..6155870b4 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -4,19 +4,16 @@ import time from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import Any import msgpack import ray import torch -from miles.utils.output_store import StoredTensor, materialize_stored_tensors +from miles.utils.output_store import RegisteredTensorReader, StoredTensor from miles.utils.processing_utils import fhwc_to_cfhw from miles.utils.types import CondKwargs, DenoisingEnv, DiTTrajectory, RolloutDebugTensors, Sample, decode_tensor -if TYPE_CHECKING: - from radix_transfer_plugins import OutputStore - __all__ = [ "apply_rollout_image_response", "RolloutImageResponseParserActor", @@ -80,11 +77,9 @@ def _normalize_generated_output(tensor: torch.Tensor | None) -> torch.Tensor | N return canonical.contiguous() -def read_generated_outputs(store: OutputStore, outputs: list[StoredTensor]) -> list[torch.Tensor]: +def read_generated_outputs(reader: RegisteredTensorReader, stored_outputs: list[StoredTensor]) -> list[torch.Tensor]: """Read stored outputs in one store read, in the ``[C, F, H, W]`` layout an inline output is parsed to.""" - tensors = list(outputs) - materialize_stored_tensors(tensors, store=store) - return [_normalize_generated_output(tensor) for tensor in tensors] + return [_normalize_generated_output(tensor) for tensor in reader.read(stored_outputs)] def _parse_tensor_or_list( diff --git a/miles/utils/output_store.py b/miles/utils/output_store.py index b999c2f0e..c3edc5f53 100644 --- a/miles/utils/output_store.py +++ b/miles/utils/output_store.py @@ -3,12 +3,16 @@ from __future__ import annotations import dataclasses +import itertools import json +import threading from argparse import Namespace from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass from typing import TYPE_CHECKING, Any +import torch + from miles.utils.types import Sample if TYPE_CHECKING: @@ -100,6 +104,56 @@ class StoredTensor: field: str +class RegisteredTensorReader: + """Reads stored tensors into one buffer that stays registered with the store. + + ``OutputStore.read_many`` registers fresh memory on every call, and under RDMA a registration pins that memory. + The buffer is registered when first needed and again only to grow, and stays for the life of the process. Only + fields that are not empty and stored uncompressed can be read; every tensor the engines store is. + """ + + def __init__(self, store: OutputStore) -> None: + self._store = store + self._buffer = None + # the store allows one read into a buffer at a time, and API reward actors read from several threads + self._lock = threading.Lock() + + def read(self, stored_tensors: list[StoredTensor]) -> list[torch.Tensor]: + """The tensors of ``stored_tensors``, each a copy the caller owns, in one store read.""" + from radix_transfer_plugins import RangeRead + + specs = [ + next(spec for spec in stored_tensor.ref.fields if spec.name == stored_tensor.field) + for stored_tensor in stored_tensors + ] + *offsets, total_nbytes = itertools.accumulate((spec.nbytes for spec in specs), initial=0) + range_reads = [ + RangeRead( + ref=stored_tensor.ref, + field=stored_tensor.field, + src_offset=0, + nbytes=spec.nbytes, + dst_offset=offset, + ) + for stored_tensor, spec, offset in zip(stored_tensors, specs, offsets, strict=True) + ] + with self._lock: + if self._buffer is None or self._buffer.nbytes < total_nbytes: + if self._buffer is not None: + self._buffer.close() + self._buffer = self._store.read_buffer(total_nbytes) + self._store.read_ranges(range_reads, into=self._buffer) + tensors = [] + for spec, offset in zip(specs, offsets, strict=True): + dtype = getattr(torch, spec.dtype) + tensor = torch.frombuffer( + self._buffer.view, dtype=dtype, count=spec.nbytes // dtype.itemsize, offset=offset + ) + # copied out: the next read overwrites the buffer + tensors.append(tensor.reshape(spec.shape).clone()) + return tensors + + def store_refs_of(samples: Iterable[Any]) -> dict[str, StoreRef]: """Every output store ref the samples hold, by key; groups may nest samples in lists.""" refs = {} diff --git a/miles/utils/tracking_utils.py b/miles/utils/tracking_utils.py index 5929ac05d..f11f85203 100644 --- a/miles/utils/tracking_utils.py +++ b/miles/utils/tracking_utils.py @@ -16,6 +16,11 @@ def init_tracking(args, primary: bool = True, **kwargs): wandb_utils.init_wandb_secondary(args, **kwargs) +def logs_rollout_media(args, rollout_id: int) -> bool: + """Whether train rollout ``rollout_id`` logs media to W&B.""" + return args.wandb_log_num_images > 0 and rollout_id % args.wandb_log_image_interval == 0 + + def log(args, metrics, step_key, is_media=False): # E2E metrics tee (tests/ci/e2e_metrics_registry.py): mirror every metric # dict to MILES_METRICS_JSONL; O_APPEND keeps concurrent writers line-atomic. diff --git a/miles/utils/types.py b/miles/utils/types.py index 40df0cfd8..93f5cbc25 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -82,7 +82,7 @@ class Sample: # reproducibility seed: int | None = None # Eager tensor on CPU. Image rollout shape: ``[C, T, H, W]`` (``T==1`` typical). - # With --use-output-store, a StoredTensor between generation and load_generated_outputs. + # With --use-output-store, a StoredTensor from generation until its reward, then None (or a tensor kept for media). generated_output: torch.Tensor | None = None # [C, L] waveform at audio_sample_rate, for engines that also return audio (e.g. H3 t2va). generated_audio: torch.Tensor | None = None diff --git a/scripts/run_diffusion_grpo_krea2_async.py b/scripts/run_diffusion_grpo_krea2_async.py index 4a7b225db..209575e23 100644 --- a/scripts/run_diffusion_grpo_krea2_async.py +++ b/scripts/run_diffusion_grpo_krea2_async.py @@ -99,6 +99,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: __file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=8, wandb_log_image_interval=10 ) + # the launcher starts the Mooncake master on the --output-store-mooncake-master port sglang_args = ( "--use-miles-router " "--sglang-server-concurrency 8 " @@ -107,9 +108,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--sglang-dit-precision bf16 " "--sglang-vae-slicing " "--update-weight-buffer-size 2147483648 " - # the launcher starts the Mooncake master on this port; debug mode stores every step's debug tensors, - # ~12 GB per 240-sample rollout with two rollouts in flight, beyond the 16gb default segment - "--use-output-store --output-store-mooncake-master 127.0.0.1:50070 --output-store-global-segment-size 64gb " + "--use-output-store --output-store-mooncake-master 127.0.0.1:50070 " ) train_backend_args = "--train-backend fsdp --diffusion-forward-dtype bf16 --micro-batch-size 2 " diff --git a/tests/e2e/short/test_krea2_nft_ocr_async_8xGPU_4xGPU_proxy.py b/tests/e2e/short/test_krea2_nft_ocr_async_8xGPU_4xGPU_proxy.py index b7a23b2de..0712ee70f 100644 --- a/tests/e2e/short/test_krea2_nft_ocr_async_8xGPU_4xGPU_proxy.py +++ b/tests/e2e/short/test_krea2_nft_ocr_async_8xGPU_4xGPU_proxy.py @@ -17,8 +17,13 @@ est_time=2400, suite="stage-c-5-gpu-h200", script="scripts/run_diffusion_nft_krea2_async.py", - # Inline on purpose: the legacy (no output store) gate for rollouts fetched by the rollout manager. - args=["--num-rollout", "4", "--four-gpu-ci"], + args=[ + "--num-rollout", + "4", + "--four-gpu-ci", + "--extra-args", + "--use-output-store --output-store-mooncake-master 127.0.0.1:50063", + ], labels=["e2e"], metrics=[ "rollout/reward/raw_num_samples", diff --git a/tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py b/tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py index 16bd14570..cf4ed869e 100644 --- a/tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py +++ b/tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py @@ -10,8 +10,7 @@ est_time=1200, suite="stage-c-3-gpu-h200", script="scripts/run_diffusion_grpo_sd3_ocr_sglang.py", - # Inline on purpose: the legacy (no output store) gate for rollouts fetched in the parser actors. - args=["--num-rollout", "2"], + args=["--num-rollout", "2", "--extra-args", "--use-output-store --output-store-mooncake-master 127.0.0.1:50067"], metrics=[ "rollout/reward/raw_num_samples", "rollout/reward/raw_mean", diff --git a/tests/fast/rollout/test_api_reward_pool.py b/tests/fast/rollout/test_api_reward_pool.py index d285b3dcd..8a92a98be 100644 --- a/tests/fast/rollout/test_api_reward_pool.py +++ b/tests/fast/rollout/test_api_reward_pool.py @@ -59,7 +59,7 @@ def factory(**kwargs): def test_pool_reuses_one_concurrent_zero_gpu_worker(ray_worker, pool_cls): remote, actor_cls = ray_worker config = ApiRewardConfig(actor_kwargs={"model": "judge", "api_key_env": "TEST_RM_KEY"}, max_concurrency=2) - args = Namespace(**{f"_{pool_cls.name}_rm_config": config}) + args = Namespace(use_output_store=False, **{f"_{pool_cls.name}_rm_config": config}) pool = pool_cls(args) assert pool_cls(args) is pool @@ -120,7 +120,7 @@ def handler(request): return httpx.Response(200, json={"reward": float(payload["prompt"])}) api_transport(handler) - pool = AsyncApiRewardPool(Namespace(_custom_api_rm_config=config)) + pool = AsyncApiRewardPool(Namespace(use_output_store=False, _custom_api_rm_config=config)) remote.assert_called_once_with(CustomApiRewardActor) actor_cls.remote.assert_called_once_with(**kwargs) actor_cls.options.assert_called_once_with(num_cpus=0, num_gpus=0, scheduling_strategy="DEFAULT", max_concurrency=3) @@ -155,7 +155,7 @@ def test_custom_and_openai_pools_use_independent_configs_and_workers(ray_worker) max_concurrency=3, ) openai_config = ApiRewardConfig(actor_kwargs={"model": "judge", "api_key_env": "TEST_RM_KEY"}, max_concurrency=7) - args = Namespace(_custom_api_rm_config=custom_config, _openai_api_rm_config=openai_config) + args = Namespace(use_output_store=False, _custom_api_rm_config=custom_config, _openai_api_rm_config=openai_config) custom_pool = AsyncApiRewardPool(args) openai_pool = AsyncOpenAIPool(args) assert custom_pool is not openai_pool diff --git a/tests/fast/rollout/test_generated_output_store.py b/tests/fast/rollout/test_generated_output_store.py index 5d1038b58..910bd3cee 100644 --- a/tests/fast/rollout/test_generated_output_store.py +++ b/tests/fast/rollout/test_generated_output_store.py @@ -1,28 +1,36 @@ -"""A rollout's generated output through the output store, from the engine's marker to the reward. +"""A rollout's generated output through the output store, from the engine's marker to its delete. -Mental model (one microgroup, store on): +Mental model (one microgroup, store on; the rollout process never reads an output it does not log): engine response {"__store__": "generated_output."} - | apply_rollout_image_response (1) a StoredTensor, layout left as stored + | apply_rollout_image_response (1) a StoredTensor, layout left as stored v sample.generated_output = StoredTensor(ref, "generated_output.") - | load_generated_outputs, inside the semaphore - | read_generated_outputs (2) == the inline-parsed tensor (CHW, HWC); one read per ref - | delete the output fields (3) latents stay for the trainer; inline (eval) outputs untouched; - | read + delete is the dashboard's store_read stage - | (4) a failing delete is left to the release; a failing read fails + | slot released; the reward is given the refs (2) at every call site: microgroup, custom generate, + v group reward + reward pool actor = _StoredOutputRewardActor (3) with the store on, not with it off + | read_generated_outputs (4) == the inline parse; the buffer is registered once, + | and again only to grow + | (5) reads from several threads get their own tensors + | scorer.score_batch(tensors) (6) refs read, inline tensors passed through + | PickScore: pickscore_rm hands a stored video over unsliced (7) + | the actor picks its frames of the whole video (8) == the frames picked inline v - sample.generated_output = tensor (5) every call site (microgroup, group_rm, custom generate) - | rewards tensors, read while holding the semaphore + release_generated_outputs after the reward (9) only output fields deleted, outputs None, inline + | untouched, nothing read when no media is logged + | keep the first N for media as tensors (10) at most N across concurrent microgroups + | (11) a failing read or delete does not fail the rollout v - reward, media logging, trainer: as with the store off + generate_rollout_async sets the media counter (12) only on rollouts that log media + rollout manager logs the samples holding a tensor (13) """ from tests.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[]) +register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) import asyncio +import concurrent.futures import contextlib from argparse import Namespace from types import SimpleNamespace @@ -33,28 +41,51 @@ from radix_transfer_plugins.stashes.tensor_stash import TensorStash from tests.fast.utils.test_output_store_materialize import MemoryBackend +import miles.ray.rollout as rollout_module +import miles.rollout.rm_hub.core as core_module +import miles.rollout.rm_hub.pickscore as pickscore_module import miles.rollout.sglang_diffusion_rollout as rollout_fn_module -from miles.dashboard.events import STAGE_KINDS +from miles.ray.rollout import RolloutManager +from miles.rollout.rm_hub.core import _StoredOutputRewardActor from miles.rollout.sglang_diffusion_rollout import ( generate_and_rm_group, generate_and_rm_microgroup, - load_generated_outputs, + generate_rollout_async, + release_generated_outputs, ) from miles.utils.diffusion_rollout_response import ( _normalize_generated_output, apply_rollout_image_response, read_generated_outputs, ) -from miles.utils.output_store import StoredTensor +from miles.utils.output_store import RegisteredTensorReader, StoredTensor +from miles.utils.processing_utils import sample_frame_indices from miles.utils.types import Sample _CHW = torch.arange(3 * 5 * 6, dtype=torch.float32).reshape(3, 5, 6) / 90 _HWC = _CHW.permute(1, 2, 0).contiguous() _LATENTS = "dit_trajectory.latents.0.0" +_RolloutManager = RolloutManager.__ray_metadata__.modified_class class _PartsBackend(MemoryBackend): - """Deletes the named parts only and ignores missing ones, as Mooncake does.""" + """Deletes the named parts only and reads byte ranges, as Mooncake does; counts registrations and range reads.""" + + def __init__(self): + super().__init__() + self.registrations = 0 + self.range_reads = 0 + + def register(self, memory): + self.registrations += 1 + + def get_ranges(self, reads, *, into): + self.range_reads += 1 + for read in reads: + data = self.objects[read.ref.key][read.field] + into[read.dst_offset : read.dst_offset + read.nbytes] = data[ + read.src_offset : read.src_offset + read.nbytes + ] def delete_many(self, keys, *, parts): for key, names in zip(keys, parts, strict=True): @@ -69,6 +100,20 @@ def _stored(tensors: dict[str, torch.Tensor]): return store, backend, ref +def _stored_samples(ref, indices) -> list[Sample]: + return [Sample(index=b, generated_output=StoredTensor(ref, f"generated_output.{b}")) for b in indices] + + +def _state(store, *, media_outputs_left=0): + return SimpleNamespace( + semaphore=asyncio.Semaphore(1), + dp_rank_context=contextlib.nullcontext, + output_store=store, + output_reader=RegisteredTensorReader(store) if store is not None else None, + media_outputs_left=media_outputs_left, + ) + + def test_parse_keeps_a_stored_output_as_a_store_ref(): """Canonicalizing a StoredTensor at parse time would fail; its layout is fixed where it is read.""" store, _, ref = _stored({"generated_output.0": _HWC}) @@ -79,106 +124,307 @@ def test_parse_keeps_a_stored_output_as_a_store_ref(): assert sample.generated_output == StoredTensor(sample.rollout_store_ref, "generated_output.0") -def test_read_matches_the_inline_parse_with_one_read_per_ref(): - """A reward must score the same tensor as on the inline path, whatever layout the engine stored.""" - store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _HWC}) +@pytest.mark.asyncio +@pytest.mark.parametrize("call_site", ["microgroup", "custom_generate", "group_rm"]) +async def test_every_call_site_hands_the_reward_refs_and_releases_them(monkeypatch, call_site): + """The rollout process must not read an output on any path, generation must not wait for a reward, and a + path that skipped the release would leave every output in the store until the train step.""" + store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _CHW, _LATENTS: _CHW}) + state = _state(store) + rewarded = [] + + async def generate(args, microgroup, sampling_params, evaluation=False): + for sample in microgroup: + sample.generated_output = StoredTensor(ref, f"generated_output.{sample.index}") + return microgroup + + async def reward(args, samples): + rewarded.extend((sample.generated_output, state.semaphore.locked()) for sample in samples) + return [1.0] * len(samples) - outputs = read_generated_outputs( - store, [StoredTensor(ref, "generated_output.0"), StoredTensor(ref, "generated_output.1")] + monkeypatch.setattr(rollout_fn_module, "GenerateState", lambda args: state) + monkeypatch.setattr(rollout_fn_module, "generate_microgroup", generate) + monkeypatch.setattr(rollout_fn_module, "load_function", lambda path: generate) + monkeypatch.setattr(rollout_fn_module, "batched_async_rm", reward) + monkeypatch.setattr(rollout_fn_module.hooks, "record_trajectory", lambda sample: None) + args = Namespace( + custom_generate_function_path="custom.generate" if call_site == "custom_generate" else None, + group_rm=call_site == "group_rm", + n_samples_per_prompt=2, + rollout_seed=0, + rollout_microgroup_size=2, ) + group = [Sample(index=b, group_index=0) for b in (0, 1)] + + if call_site == "group_rm": + await generate_and_rm_group(args, group, {}) + else: + await generate_and_rm_microgroup(args, group, {}) + + assert len(rewarded) == 2 + assert all(isinstance(output, StoredTensor) and not slot_held for output, slot_held in rewarded) + assert backend.range_reads == 0 and set(backend.objects[ref.key]) == {_LATENTS} + assert [sample.generated_output for sample in group] == [None, None] + + +@pytest.mark.parametrize("use_output_store", [True, False], ids=["store_on", "store_off"]) +def test_pool_actors_read_the_store_only_with_it_on(monkeypatch, use_output_store): + """Unwrapped actors would get refs they cannot score; with the store off there is no store to connect to. + A PickScore actor without its frame count would average every frame of a stored video.""" + created = [] + + class FakeRemote: + def __init__(self, actor_cls): + self.actor_cls = actor_cls + + def options(self, **options): + return self + + def remote(self, **kwargs): + created.append((self.actor_cls, kwargs)) + + monkeypatch.setattr(core_module.ray, "remote", FakeRemote) + monkeypatch.setattr(pickscore_module.SingletonMeta, "_instances", {}) + args = Namespace( + use_output_store=use_output_store, + pickscore_processor_path="processor", + pickscore_model_path="model", + pickscore_batch_size=8, + pickscore_num_workers=1, + pickscore_num_gpus_per_worker=0, + pickscore_reward_colocate=False, + pickscore_num_frames=3, + ) + + pickscore_module.AsyncPickScorePool(args) + + ((actor_cls, kwargs),) = created + scorer_kwargs = {"processor_path": "processor", "model_path": "model", "frames_per_forward": 8, "num_frames": 3} + if use_output_store: + assert actor_cls is _StoredOutputRewardActor + assert kwargs == { + "scorer_cls": pickscore_module.PickScoreRewardActor, + "scorer_kwargs": scorer_kwargs, + "args": args, + } + else: + assert actor_cls is pickscore_module.PickScoreRewardActor and kwargs == scorer_kwargs + + +def test_read_matches_the_inline_parse_and_registers_its_buffer_once(): + """A reward must score the same tensor as on the inline path, whatever layout the engine stored. Under RDMA a + registration pins the memory, so a reader that registered on every read would be slower than the transfer.""" + store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _HWC, "generated_output.2": _CHW}) + reader = RegisteredTensorReader(store) + one_output = [StoredTensor(ref, "generated_output.0")] + three_outputs = [StoredTensor(ref, f"generated_output.{b}") for b in range(3)] + + first = read_generated_outputs(reader, one_output) + again = read_generated_outputs(reader, one_output) + grown = read_generated_outputs(reader, three_outputs) + + assert torch.equal(first[0], _normalize_generated_output(_CHW)) and torch.equal(again[0], first[0]) + assert torch.equal(grown[1], _normalize_generated_output(_HWC)) + assert torch.equal(grown[2], _normalize_generated_output(_CHW)) + # the first read registers, the second reuses, the larger one registers again + assert backend.registrations == 2 and backend.range_reads == 3 + + +def test_concurrent_reads_return_their_own_tensors(): + """API reward actors read from several threads; a read that shared the buffer would hand one thread another's + bytes.""" + tensors = {f"generated_output.{b}": torch.full((3, 5, 6), float(b)) for b in range(8)} + store, _, ref = _stored(tensors) + reader = RegisteredTensorReader(store) + + def reads_its_own_tensor(read_index): + field = f"generated_output.{read_index % 8}" + (tensor,) = reader.read([StoredTensor(ref, field)]) + return torch.equal(tensor, tensors[field]) - assert torch.equal(outputs[0], _normalize_generated_output(_CHW)) - assert torch.equal(outputs[1], _normalize_generated_output(_HWC)) - assert backend.reads == [[(ref.key, ["generated_output.0", "generated_output.1"])]] + with concurrent.futures.ThreadPoolExecutor(8) as pool: + assert all(pool.map(reads_its_own_tensor, range(200))) + + +def test_reward_actor_reads_refs_and_passes_tensors_through(monkeypatch): + """The scorer must get tensors in the order it was given outputs; eval outputs arrive inline.""" + store, _, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _HWC}) + monkeypatch.setattr(core_module, "connect_output_store", lambda args: store) + received = [] + + class FakeScorer: + def score_batch(self, outputs, prompts): + received.extend(outputs) + return [1.0] * len(outputs) + + actor = _StoredOutputRewardActor(scorer_cls=FakeScorer, scorer_kwargs={}, args=Namespace()) + inline = torch.zeros(3, 1, 4, 5) + outputs = [StoredTensor(ref, "generated_output.0"), inline, StoredTensor(ref, "generated_output.1")] + + assert actor.score_batch(outputs, ["a", "b", "c"]) == [1.0, 1.0, 1.0] + assert torch.equal(received[0], _normalize_generated_output(_CHW)) + assert received[1] is inline + assert torch.equal(received[2], _normalize_generated_output(_HWC)) @pytest.mark.asyncio -async def test_load_reads_the_outputs_and_deletes_only_their_fields(): - """Latents stay for the trainer; the outputs leave the store once read; an inline (eval) output is untouched. - The dashboard drops a stage name it does not know, so the read time must use a known one.""" - store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _HWC, _LATENTS: _CHW}) +async def test_pickscore_hands_a_stored_video_over_unsliced(monkeypatch): + """Slicing a StoredTensor in the rollout process would crash; its frames are picked where it is read.""" + _, _, ref = _stored({"generated_output.0": torch.rand(3, 7, 4, 5)}) + scored = [] + + class FakePool: + async def score(self, outputs, prompts): + scored.extend(outputs) + return [1.0] * len(outputs), 0 + + monkeypatch.setattr(pickscore_module, "AsyncPickScorePool", lambda args: FakePool()) + stored_video = StoredTensor(ref, "generated_output.0") + samples = [ + Sample(prompt="a", generated_output=stored_video), + Sample(prompt="b", generated_output=torch.rand(3, 7, 4, 5)), + ] + + await pickscore_module.pickscore_rm(Namespace(pickscore_num_frames=3), samples) + + assert scored[0] is stored_video and scored[1].shape == (3, 3, 4, 5) + + +@pytest.mark.parametrize( + ("total_frames", "num_frames"), + [(7, 3), (7, 1), (7, None), (2, 3)], + ids=["3_of_7", "1_of_7", "all_of_7", "3_of_2"], +) +def test_pickscore_picks_the_same_frames_from_a_whole_stored_video(total_frames, num_frames): + """A stored video reaches the actor whole; it must score as the frames the rollout process picks inline. + The unpicked frames are on the 0-255 scale, so quantising before picking would rescale the picked ones.""" + actor = pickscore_module.PickScoreRewardActor.__new__(pickscore_module.PickScoreRewardActor) + actor.frames_per_forward = 8 + actor.num_frames = num_frames + actor.scorer = lambda prompts, images: [float(sum(image.getdata(0))) for image in images] + picked_frames = sample_frame_indices(total_frames, num_frames) + video = torch.rand(3, total_frames, 4, 6) + for frame in set(range(total_frames)) - set(picked_frames): + video[:, frame] = 200.0 + + whole = actor.score_batch([video], ["p"]) + picked = actor.score_batch([video[:, picked_frames]], ["p"]) + + assert whole == picked + + +@pytest.mark.asyncio +async def test_release_deletes_only_the_scored_outputs(): + """Latents stay for the trainer; the outputs go once scored, the samples drop their dangling refs, and an + inline (eval) output is untouched. A rollout that logs no media reads nothing in the rollout process.""" + store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _CHW, _LATENTS: _CHW}) inline = torch.zeros(3, 1, 2, 2) - samples = [Sample(index=b, generated_output=StoredTensor(ref, f"generated_output.{b}")) for b in (0, 1)] - samples.append(Sample(index=2, generated_output=inline)) + samples = [*_stored_samples(ref, (0, 1)), Sample(index=2, generated_output=inline)] - await load_generated_outputs(SimpleNamespace(output_store=store), samples) + await release_generated_outputs(_state(store), samples) - assert torch.equal(samples[0].generated_output, _normalize_generated_output(_CHW)) - assert torch.equal(samples[1].generated_output, _normalize_generated_output(_HWC)) + assert set(backend.objects[ref.key]) == {_LATENTS} and backend.range_reads == 0 + assert [sample.generated_output for sample in samples[:2]] == [None, None] assert samples[2].generated_output is inline - assert set(backend.objects[ref.key]) == {_LATENTS} - stages = [ - [segment["stage"] for segment in (sample.metadata or {}).get("lifecycle_stages", [])] for sample in samples - ] - assert stages == [["store_read"], ["store_read"], []] and "store_read" in STAGE_KINDS + + +@pytest.mark.asyncio +async def test_logged_rollout_keeps_at_most_the_logged_count_as_tensors(): + """Media logging needs N outputs after the fields are gone; concurrent microgroups must not keep more than N.""" + store, backend, ref = _stored({f"generated_output.{b}": _CHW * b for b in range(4)}) + groups = [_stored_samples(ref, pair) for pair in ([0, 1], [2, 3])] + state = _state(store, media_outputs_left=1) + + await asyncio.gather(*(release_generated_outputs(state, group) for group in groups)) + + outputs = [sample.generated_output for group in groups for sample in group] + assert torch.equal(outputs[0], _normalize_generated_output(_CHW * 0)) + assert outputs[1:] == [None, None, None] + assert backend.objects[ref.key] == {} @pytest.mark.asyncio @pytest.mark.parametrize("failing", ["read", "delete"]) -async def test_a_failing_delete_is_left_to_the_release_but_a_failing_read_fails(failing): - """The reward cannot run without the output; a field left behind is deleted with the object by the release.""" - store, _, ref = _stored({"generated_output.0": _CHW}) +async def test_failing_release_does_not_fail_the_rollout(failing): + """A failed media read or delete must not lose a scored rollout; the rollout's release deletes the object.""" + store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _CHW}) def fail(*args, **kwargs): raise RuntimeError(f"Mooncake {failing} failed") - setattr(store, "read_many" if failing == "read" else "delete_many", fail) - sample = Sample(index=0, generated_output=StoredTensor(ref, "generated_output.0")) + setattr(store, "read_ranges" if failing == "read" else "delete_many", fail) + samples = _stored_samples(ref, (0, 1)) + + await release_generated_outputs(_state(store, media_outputs_left=1), samples) if failing == "read": - with pytest.raises(RuntimeError, match="read failed"): - await load_generated_outputs(SimpleNamespace(output_store=store), [sample]) + assert samples[0].generated_output is None and backend.objects[ref.key] == {} else: - await load_generated_outputs(SimpleNamespace(output_store=store), [sample]) - assert torch.equal(sample.generated_output, _normalize_generated_output(_CHW)) + assert torch.equal(samples[0].generated_output, _normalize_generated_output(_CHW)) + assert samples[1].generated_output is None @pytest.mark.asyncio -@pytest.mark.parametrize("call_site", ["microgroup", "group_rm", "custom_generate"]) -async def test_every_call_site_rewards_tensors_read_inside_the_semaphore(monkeypatch, call_site): - """A call site without the load would hand rewards store refs; a load outside the semaphore would let the - store fill with outputs of more microgroups than there are generation slots.""" - store, backend, ref = _stored({"generated_output.0": _CHW, "generated_output.1": _CHW, _LATENTS: _CHW}) - state = SimpleNamespace(semaphore=asyncio.Semaphore(1), dp_rank_context=contextlib.nullcontext, output_store=store) - read_inside_semaphore = [] - read_many = store.read_many - - def recording_read_many(*args, **kwargs): - read_inside_semaphore.append(state.semaphore.locked()) - return read_many(*args, **kwargs) +@pytest.mark.parametrize(("rollout_id", "media_outputs_kept"), [(10, 4), (11, 0)], ids=["logged", "not_logged"]) +async def test_only_a_rollout_that_logs_media_keeps_outputs(monkeypatch, rollout_id, media_outputs_kept): + """A rollout must keep exactly the outputs the rollout manager will log for it, and none otherwise.""" - store.read_many = recording_read_many - rewarded = [] + class _Stop(Exception): + pass - async def generate(args, microgroup, sampling_params, evaluation=False): - for sample in microgroup: - sample.generated_output = StoredTensor(ref, f"generated_output.{sample.index}") - return microgroup + class FakeState: + media_outputs_left = None - async def reward(args, samples): - rewarded.extend(sample.generated_output for sample in samples) - return [1.0] * len(samples) + @property + def remaining_batch_size(self): + raise _Stop + state = FakeState() monkeypatch.setattr(rollout_fn_module, "GenerateState", lambda args: state) - monkeypatch.setattr(rollout_fn_module, "generate_microgroup", generate) - monkeypatch.setattr(rollout_fn_module, "load_function", lambda path: generate) - monkeypatch.setattr(rollout_fn_module, "batched_async_rm", reward) - monkeypatch.setattr(rollout_fn_module.hooks, "record_trajectory", lambda sample: None) args = Namespace( - custom_generate_function_path="custom.generate" if call_site == "custom_generate" else None, - group_rm=call_site == "group_rm", - n_samples_per_prompt=2, - rollout_seed=0, - rollout_microgroup_size=2, + rollout_global_dataset=True, + wandb_log_num_images=4, + wandb_log_image_interval=10, + dynamic_sampling_filter_path=None, + rollout_batch_size=1, + over_sampling_batch_size=1, + n_samples_per_prompt=1, ) - group = [Sample(index=b, group_index=0) for b in (0, 1)] - if call_site == "group_rm": - await generate_and_rm_group(args, group, {}) - else: - await generate_and_rm_microgroup(args, group, {}) + with pytest.raises(_Stop): + await generate_rollout_async(args, rollout_id, data_source=None) - assert len(rewarded) == 2 - assert all(torch.equal(output, _normalize_generated_output(_CHW)) for output in rewarded) - assert read_inside_semaphore == [True] - assert set(backend.objects[ref.key]) == {_LATENTS} + assert state.media_outputs_left == media_outputs_kept + + +def test_media_logging_takes_the_samples_that_kept_an_output(monkeypatch): + """With the first samples' outputs dropped, logging the first N samples would log nothing; a store ref a + rollout function never released would crash the logging thread.""" + _, _, ref = _stored({"generated_output.0": _CHW}) + media_sample = Sample(index=2, reward=1.0, generated_output=torch.zeros(3, 1, 2, 2)) + samples = [ + Sample(index=0, reward=1.0, generated_output=None), + Sample(index=1, reward=1.0, generated_output=StoredTensor(ref, "generated_output.0")), + media_sample, + ] + logged = [] + manager = SimpleNamespace( + args=Namespace( + wandb_log_num_images=1, + wandb_log_image_interval=1, + advantage_estimator="grpo", + n_samples_per_prompt=1, + reward_key=None, + ), + rollout_id=0, + custom_convert_samples_to_train_data_func=None, + custom_expand_samples_to_train_pairs_func=None, + _post_process_rewards=lambda samples: ([1.0] * len(samples), [0.0] * len(samples)), + _log_images_async=lambda media, **kwargs: logged.append(media["rollout_media/sample_images"]), + ) + monkeypatch.setattr(rollout_module.tracking_utils, "log", lambda *args, **kwargs: None) + monkeypatch.setattr(rollout_module, "flow_grpo_expand_samples_to_train_pairs", lambda *args: None) + + _RolloutManager._convert_samples_to_train_data(manager, samples) + + assert logged == [[media_sample]] diff --git a/tests/fast/rollout/test_hps_reward.py b/tests/fast/rollout/test_hps_reward.py index 827923b98..b0bed75fe 100644 --- a/tests/fast/rollout/test_hps_reward.py +++ b/tests/fast/rollout/test_hps_reward.py @@ -51,6 +51,7 @@ def test_pickscore_actor_truncates_and_averages_frames_per_sample(): --pickscore-batch-size frames (the chunking the e2e standards were recorded with), not one sample.""" actor = PickScoreRewardActor.__new__(PickScoreRewardActor) actor.frames_per_forward = 1 + actor.num_frames = None forward_sizes = [] def scorer(prompts, images): diff --git a/tests/fast/rollout/test_reward_pool_placement.py b/tests/fast/rollout/test_reward_pool_placement.py index 4f808cdd4..7ea78cf3f 100644 --- a/tests/fast/rollout/test_reward_pool_placement.py +++ b/tests/fast/rollout/test_reward_pool_placement.py @@ -14,6 +14,8 @@ register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) +from argparse import Namespace + import pytest import miles.rollout.rm_hub.core as core_module @@ -42,6 +44,7 @@ def _no_ray(monkeypatch): def _pool(actor_cls, *, num_workers, colocate, placement_group=None, slots=None): return AsyncRewardActorPool( + args=Namespace(use_output_store=False), actor_cls=actor_cls, actor_kwargs={}, num_workers=num_workers, diff --git a/tests/fast/rollout/test_rollout_store_refs.py b/tests/fast/rollout/test_rollout_store_refs.py index 221fc464b..7752326aa 100644 --- a/tests/fast/rollout/test_rollout_store_refs.py +++ b/tests/fast/rollout/test_rollout_store_refs.py @@ -61,6 +61,8 @@ def _args() -> Namespace: n_samples_per_prompt=1, dynamic_sampling_filter_path=f"{__name__}.drop_zero_reward", rollout_sample_filter_path=None, + wandb_log_num_images=0, + wandb_log_image_interval=1, ) diff --git a/tests/fast/utils/test_output_store_args.py b/tests/fast/utils/test_output_store_args.py index 978a0b561..e5ffbc21d 100644 --- a/tests/fast/utils/test_output_store_args.py +++ b/tests/fast/utils/test_output_store_args.py @@ -48,7 +48,7 @@ def test_flags_parse(monkeypatch: pytest.MonkeyPatch) -> None: assert args.use_output_store assert args.output_store_mooncake_master == "10.0.0.1:50051" assert args.output_store_local_buffer_size == "2gb" - assert args.output_store_global_segment_size == "16gb" + assert args.output_store_global_segment_size == "72gb" assert args.output_store_protocol == "tcp" assert args.output_store_device_name == ""