diff --git a/miles/backends/fsdp_utils/checkpoint.py b/miles/backends/fsdp_utils/checkpoint.py index 1d0dbcd5..ff1a3377 100644 --- a/miles/backends/fsdp_utils/checkpoint.py +++ b/miles/backends/fsdp_utils/checkpoint.py @@ -85,13 +85,25 @@ def state_dict(self): def load_state_dict(self, state_dict): options = StateDictOptions(strict=False) if self.lora_only else None - set_state_dict( + incompatible = set_state_dict( self.model, optimizers=[], model_state_dict=state_dict["model"], optim_state_dict=None, options=options, ) + # strict=False swallows mismatches, so a checkpoint whose keys match + # nothing "loads" as a silent no-op. In lora_only mode missing base keys + # are expected; unclaimed checkpoint keys or unfilled LoRA params are not. + unclaimed = list(incompatible.unexpected_keys) + unfilled = [k for k in incompatible.missing_keys if not self.lora_only or _is_lora_param_name(k)] + if unclaimed or unfilled: + logger.error( + f"[FSDP] Checkpoint mismatch: {len(unclaimed)} unclaimed checkpoint keys, " + f"{len(unfilled)} unfilled model params (e.g. {(unclaimed + unfilled)[:3]})" + ) + else: + logger.info(f"[FSDP] Applied {len(state_dict['model'])} params from the checkpoint") class OptimizerState(Stateful): diff --git a/miles/rollout/rm_hub/core.py b/miles/rollout/rm_hub/core.py index 5dcdc066..0ad3d879 100644 --- a/miles/rollout/rm_hub/core.py +++ b/miles/rollout/rm_hub/core.py @@ -82,11 +82,16 @@ def _next_actor(self): return actor async def score(self, images: list, prompts: list[str]) -> list[float]: - refs = [] - for start in range(0, len(images), self._batch_size): - end = start + self._batch_size - refs.append(self._next_actor().score_batch.remote(images[start:end], prompts[start:end])) + chunks = [ + (self._next_actor(), images[start:start + self._batch_size], prompts[start:start + self._batch_size]) + for start in range(0, len(images), self._batch_size) + ] + + def submit_and_collect(): + # .remote() copies each chunk into plasma in the calling thread, so + # submitting on the loop stalls every concurrent request for that copy. + return ray.get([actor.score_batch.remote(img, prm) for actor, img, prm in chunks]) loop = asyncio.get_running_loop() - chunked_scores = await loop.run_in_executor(None, ray.get, refs) + chunked_scores = await loop.run_in_executor(None, submit_and_collect) return [float(score) for chunk in chunked_scores for score in chunk] diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index 8b7d7d26..67a9c059 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -8,7 +8,7 @@ from PIL import Image from miles.utils.misc import SingletonMeta -from miles.utils.processing_utils import cfhw_to_fhwc, image_or_video_to_uint8 +from miles.utils.processing_utils import cfhw_to_fhwc from miles.utils.types import Sample from .core import AsyncRewardActorPool @@ -39,9 +39,17 @@ def _sample_to_rgb_hwc_uint8_frames(sample: Sample, num_frames: int | None) -> l if cfhw is None: raise ValueError("generated_output is None") - fhwc = image_or_video_to_uint8(cfhw_to_fhwc(cfhw.detach().cpu())) - indices = sample_frame_indices(fhwc.shape[0], num_frames) - return [np.ascontiguousarray(fhwc[i].numpy()) for i in indices] + # Convert only the frames that survive: a 107-frame clip yields 8 here, and + # converting the whole clip first costs several full-size float32 copies of + # it. The unit-scale decision still comes from the whole clip, so selecting + # first cannot change how a frame is quantised. + indices = sample_frame_indices(cfhw.shape[1], num_frames) + unit_scale = float(cfhw.max()) <= 1.0 + 1e-3 + selected = cfhw[:, indices].detach().cpu().float() + if unit_scale: + selected = selected * 255.0 + fhwc = cfhw_to_fhwc(selected.clamp(0, 255).to(torch.uint8)) + return [np.ascontiguousarray(fhwc[i].numpy()) for i in range(len(indices))] class PickScoreScorer(torch.nn.Module): diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index cde28da0..981c321e 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -2,6 +2,7 @@ import copy import inspect import logging +import random from argparse import Namespace from collections.abc import Callable from contextlib import contextmanager @@ -60,6 +61,7 @@ def build_rollout_sampling_params( "rollout_noise_level": args.diffusion_noise_level, "rollout_log_prob_no_const": args.diffusion_log_prob_no_const, "rollout_debug_mode": args.diffusion_debug_mode, + "rollout_video_dtype": args.rollout_video_dtype, "rollout_return_denoising_env": True, "rollout_return_dit_trajectory": True, } @@ -225,6 +227,8 @@ async def generate_and_rm_microgroup( state = GenerateState(args) # generate + if args.rollout_request_stagger > 0: + await asyncio.sleep(random.uniform(0, args.rollout_request_stagger)) async with state.semaphore: with state.dp_rank_context() as _: if args.custom_generate_function_path is not None: diff --git a/miles/router/router.py b/miles/router/router.py index 935ae060..ab28baa2 100644 --- a/miles/router/router.py +++ b/miles/router/router.py @@ -64,6 +64,7 @@ def _setup_routes(self): """Setup all the HTTP routes""" # sglang-router api self.app.post("/add_worker")(self.add_worker) + self.app.post("/remove_worker")(self.remove_worker) self.app.get("/list_workers")(self.list_workers) # Catch-all route for proxying to SGLang - must be registered LAST self.app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])(self.proxy) @@ -177,6 +178,30 @@ async def add_worker(self, request: Request): return {"status": "success", "worker_urls": self.worker_request_counts} + async def remove_worker(self, request: Request): + """Drop a worker from the routing pool, same URL forms as add_worker. + + Engines call this on shutdown. Without the route it fell through to the + catch-all and was proxied to a worker as if it were a generate request. + """ + worker_url = request.query_params.get("url") or request.query_params.get("worker_url") + if not worker_url: + body = await request.body() + payload = json.loads(body) if body else {} + worker_url = payload.get("url") or payload.get("worker_url") + + if not worker_url: + return JSONResponse( + status_code=400, content={"error": "worker_url is required (use query ?url=... or JSON body)"} + ) + + self.worker_request_counts.pop(worker_url, None) + self.worker_failure_counts.pop(worker_url, None) + self.dead_workers.discard(worker_url) + if self.verbose: + print(f"[miles-router] Removed worker: {worker_url}") + return {"status": "success", "worker_urls": self.worker_request_counts} + async def list_workers(self, request: Request): """List all registered workers""" return {"urls": list(self.worker_request_counts.keys())} diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index a5b749c9..3a7f400f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -252,6 +252,30 @@ def add_rollout_arguments(parser): default=1, help="Samples per prompt sent in one rollout request (sub-batch of the group).", ) + parser.add_argument( + "--rollout-video-dtype", + type=str, + choices=["keep", "uint8"], + default="keep", + help=( + "Dtype of the decoded video in rollout responses. 'keep' returns the " + "engine's raw float tensor; 'uint8' quantises engine-side with the same " + "formula the reward path applies, cutting the response body ~4x. Use " + "'keep' for any consumer that needs the unquantised tensor." + ), + ) + parser.add_argument( + "--rollout-request-stagger", + type=float, + default=0.0, + help=( + "Max uniform random delay (s) before a request first tries to acquire a " + "concurrency slot. Requests admitted together stay in phase for the whole " + "rollout because generation time is near-constant, so their response bodies " + "land on the transfer path simultaneously; a one-shot jitter breaks that. " + "Applied before the acquire, so it never idles a held slot. 0 disables." + ), + ) parser.add_argument( "--diffusion-fps", type=float, @@ -1482,6 +1506,9 @@ def miles_validate_args(args): if args.wandb_log_image_interval < 1: raise ValueError(f"wandb_log_image_interval must be >= 1, got {args.wandb_log_image_interval}") + if args.rollout_request_stagger < 0: + raise ValueError(f"--rollout-request-stagger must be non-negative, got {args.rollout_request_stagger}") + args.rollout_patch_groups = [name for name in (args.rollout_patch_group or "").split(",") if name] if not args.hf_checkpoint: diff --git a/tests/fast/rollout/test_rollout_negative_prompt.py b/tests/fast/rollout/test_rollout_negative_prompt.py index 265e52b3..27168ca2 100644 --- a/tests/fast/rollout/test_rollout_negative_prompt.py +++ b/tests/fast/rollout/test_rollout_negative_prompt.py @@ -30,6 +30,7 @@ def _args(**overrides): diffusion_noise_level=0.7, diffusion_log_prob_no_const=False, diffusion_debug_mode=False, + rollout_video_dtype="keep", train_pipeline_config_path=None, ) values.update(overrides)