diff --git a/docs/models/h3/h3.md b/docs/models/h3/h3.md index c8e2aa027..361e0ebf5 100644 --- a/docs/models/h3/h3.md +++ b/docs/models/h3/h3.md @@ -61,6 +61,42 @@ python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \ --rollout-batch-size 4 --n-samples-per-prompt 16 --num-rollout 200 ``` +### 4.3 DMD2 4-step distillation (2 nodes x 8 GPUs) + +Recipe: `scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py` — FastVideo's FastH3 Preview v1 +Dense / Data-Free recipe on `--loss-type dmd2`: full-parameter student and critic, frozen +teacher, lr 2e-6, 4 student steps with an Euler rollout stopped at a random step, an x0-space +critic, and 5 critic steps of 64 samples per student step. Prompts are not in the repo: the +script downloads `vidprom_filtered_extended.txt` from the Hub (`gdhe17/Self-Forcing` at a pinned +revision and sha256) into `--data-dir` and splits it into fixed files with pinned hashes. + +**Status:** [○ NV — Not verified](../../user-guide/recipe-verification.md#nv) + +```bash +# after the multi-node bring-up, on the head node +MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py + +# references on the same eval prompts and seeds: the teacher at 50 steps, FastH3 at 4 steps +MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --reference teacher +MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --reference fasth3 +``` + +What to watch: + +| Signal | Where | Reading | +|---|---|---| +| Student media | W&B `eval_media/`, `//eval_media/eval_/` | Rollout 0 is the base model at 4 steps; later evals render the same 26 prompts and seeds. Compare with the two reference runs' `eval_0/` (same file names, `index.jsonl` lists prompt and seed) | +| Critic | `train/critic-loss` | Rising to a plateau is the normal DMD2 shape; judge by the media | +| Distribution matching | `train/dmd2_grad_abs_mean`, `dmd2_score_gap_abs_mean` (critic vs teacher x0), `dmd2_teacher_dist_abs_mean` (student vs teacher x0, g's normalizer), each also `_step{i}` | A step whose curves drift away from the others is where the student breaks | +| Student output | `train/dmd2_x0_mean`, `train/dmd2_x0_std` | Latent statistics drifting from rollout 0 flag saturation | +| Rollout vs train | `train/x0_mean_abs_diff` (and `_step{i}`), `x0_rel_max` | The trainer's recomputed x0 against the rollout's; should stay small | +| Memory | `train/peak_memory_allocated_gb` (this process), `train/device_memory_used_gb` (whole GPU), both with a `critic-` twin | Headroom on each group's steps | + +Training rollouts carry no media: with `--rm-type none` the engine skips their VAE decode, so the +student is judged on the eval renders. The engine returns audio only with an sglang-miles change that +adds it to the rollout response; until then the saved videos are silent. Only the video stream is +distilled, so listen for drift. + ## 5. Recipe notes SFT: the encode pool (`miles/rollout/encoder_hub/h3/`) imports the engine's own recipes diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index faaccbaa5..ab97f92d2 100644 --- a/docs/user-guide/cli-reference.md +++ b/docs/user-guide/cli-reference.md @@ -182,9 +182,9 @@ See [Dtype Control](../advanced/dtype-control.md). | `--diffusion-num-steps` | int | `10` | | | `--diffusion-flow-shift` | float | – | Generation-schedule shift, sent to the engine. | | `--rollout-microgroup-size` | int | `1` | | -| `--diffusion-fps` | float | – | Video only. | +| `--diffusion-fps` | float | – | Video only. Also the fps logged and saved media play at (8 when unset); H3 drops it from the request, so set it to 24 there. | | `--diffusion-output-num-frames` | int | `1` | | -| `--rollout-video-dtype` | enum | `keep` | `keep` / `uint8`: dtype of the decoded video in rollout responses. `uint8` quantizes engine-side by truncation (×255, then clamp), cutting the response ~4x. That matches the PickScore reward path; HPS, OCR and the OpenAI-API reward round instead, so keep `keep` for them and for any consumer that needs the float tensor. | +| `--rollout-video-dtype` | enum | `keep` | `keep` / `uint8`: dtype of the decoded video in rollout and eval responses. `uint8` quantizes engine-side by truncation (×255, then clamp), cutting the response ~4x. That matches the PickScore reward path; HPS, OCR and the OpenAI-API reward round instead, so keep `keep` for them and for any consumer that needs the float tensor. | | `--diffusion-guidance-scale` | float | `4.0` | | | `--diffusion-guidance-scale-2` | float | – | Wan2.2 low-noise expert; **required** when training it. | | `--diffusion-true-cfg-scale` | float | – | | @@ -296,7 +296,7 @@ See [Dtype Control](../advanced/dtype-control.md). | Flag | Type | Default | Notes | |---|---|---|---| -| `--rm-type` | str | – | `pickscore` / `hps` / `hpsv3` / `ocr` / `openai_api` / `custom_api`. Overridable per sample via `metadata.rm_type`. Ignored when `--custom-rm-path` is set. | +| `--rm-type` | str | – | `pickscore` / `hps` / `hpsv3` / `ocr` / `openai_api` / `custom_api`, or `none` (constant 0, for a loss that reads no reward such as `dmd2`; training rollouts then skip the engine's VAE decode, which needs an engine with `rollout_return_generated_output`, while eval still decodes). Overridable per sample via `metadata.rm_type`. Ignored when `--custom-rm-path` is set. | | `--custom-api-rm-config` | str | – | Required for `--rm-type custom_api`: a YAML file or inline `base64:` with `actor_class`, `actor_kwargs`, `max_concurrency`. A flat mapping configures the default OpenAI-compatible image actor. | | `--openai-api-rm-config` | str | – | Required for `--rm-type openai_api`; same format as `--custom-api-rm-config`. | | `--group-rm` | flag | off | Score a whole prompt group at once. | @@ -408,6 +408,7 @@ Every one takes a dotted path. | `--train-only` | flag | off | No engines, no weight sync, no eval. (`--debug-train-only` is a legacy alias.) | | `--debug-skip-optimizer-step` | flag | off | No backward/step — weights never drift. Use it to measure pure forward divergence from the engine. | | `--save-debug-rollout-data` | path template | – | Fills `{rollout_id}`. | +| `--save-eval-media-dir` | path | – | Keeps every eval sample's video (with its audio when the engine returns it) or image under `/eval_/`, with an `index.jsonl` of prompts, seeds and rewards. Independent of `--wandb-log-num-images`. | | `--load-debug-rollout-data` | path template | – | Replays a rollout with **no engines started**. | | `--load-debug-rollout-data-subsample` | float | – | | | `--save-debug-train-data` | path template | – | Fills `{rollout_id}` and `{rank}`. | diff --git a/docs/user-guide/recipe-verification.md b/docs/user-guide/recipe-verification.md index 6e571b932..2492eea7c 100644 --- a/docs/user-guide/recipe-verification.md +++ b/docs/user-guide/recipe-verification.md @@ -62,3 +62,4 @@ count as verification. Flow-GRPO + PickScore. - `run_diffusion_sft_wan22.py` — Wan2.2 4-GPU LoRA SFT. - `run_diffusion_sft_h3_ref2va.py` — MiniMax H3 8-GPU Ref2VA LoRA SFT. + - `run_diffusion_dmd2_h3_t2va_16gpu_multinode.py` — MiniMax H3 16-GPU t2va 4-step DMD2. diff --git a/docs/user-guide/rewards.md b/docs/user-guide/rewards.md index 9c96a9e5d..04a35f576 100644 --- a/docs/user-guide/rewards.md +++ b/docs/user-guide/rewards.md @@ -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`, `hpsv3`, `ocr`, `openai_api`, `custom_api`); ignored when `--custom-rm-path` is set | +| Reward type | `--rm-type` | Selects built-in scorer (`pickscore`, `hps`, `hpsv3`, `ocr`, `openai_api`, `custom_api`), or `none` for a constant 0 when the loss reads no reward (DMD2), which also skips decoding training rollouts; ignored when `--custom-rm-path` is set | | API reward config | `--openai-api-rm-config`, `--custom-api-rm-config` | YAML file or inline `base64:` for the two API rewards | | 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` | diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 6fb4916b5..6ff8e8bba 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -461,6 +461,7 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An device, self.models, sigma_buckets=self.args.log_loss_sigma_bucket, + num_rollout_steps=self.args.diffusion_num_steps, ) for pair_lo, pair_hi in microbatch_ranges: @@ -490,6 +491,12 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An self.scaler.step(self.optimizer) self.scaler.update() self.lr_scheduler.step() + # Peak since the last reset, so a rollout's first step also covers scoring and wake-up before it. + metrics.emit_max("peak_memory_allocated_gb", torch.cuda.max_memory_allocated() / 2**30) + # Device-wide, so it also counts the other processes on this GPU (sleeping groups, the engine). + free_bytes, total_bytes = torch.cuda.mem_get_info() + metrics.emit_max("device_memory_used_gb", (total_bytes - free_bytes) / 2**30) + torch.cuda.reset_peak_memory_stats() else: self.optimizer.zero_grad(set_to_none=True) self.global_step += 1 diff --git a/miles/backends/fsdp_utils/loss_hub/dmd2.py b/miles/backends/fsdp_utils/loss_hub/dmd2.py index 2746b5f48..2e3a56900 100644 --- a/miles/backends/fsdp_utils/loss_hub/dmd2.py +++ b/miles/backends/fsdp_utils/loss_hub/dmd2.py @@ -10,7 +10,7 @@ from miles.backends.fsdp_utils.input_dtype_policy import apply_input_dtype_policy from miles.backends.fsdp_utils.loss_hub.sft import prepare_sft_batch, sft_loss_formula from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch -from miles.backends.fsdp_utils.metrics import record_rollout_train_abs_diff +from miles.backends.fsdp_utils.metrics import record_rollout_train_abs_diff, step_metric_key from miles.utils.metric_buffer import MetricBuffer @@ -230,7 +230,19 @@ def dmd2_loss_formula( with torch.no_grad(): bsz = len(batch) metrics.emit_mean("loss", total=loss_sum, count=bsz) - metrics.emit_mean("dmd2_grad_abs_mean", total=grad.abs().mean(dim=reduce_dims).sum(), count=bsz) rollout_x0 = torch.stack([pair["x0"] for pair in batch]).to(device=x0.device, dtype=torch.float32) record_rollout_train_abs_diff(metrics, "x0", x0, rollout_x0) + per_pair = { + "dmd2_grad_abs_mean": grad.abs().mean(dim=reduce_dims), + "dmd2_score_gap_abs_mean": (x0_fake - x0_real).abs().mean(dim=reduce_dims), + "dmd2_teacher_dist_abs_mean": (x0 - x0_real).abs().mean(dim=reduce_dims), + "x0_mean_abs_diff": (x0 - rollout_x0).abs().mean(dim=reduce_dims), + } + for key in ("dmd2_grad_abs_mean", "dmd2_score_gap_abs_mean", "dmd2_teacher_dist_abs_mean"): + metrics.emit_mean(key, total=per_pair[key].sum(), count=bsz) + metrics.emit_mean("dmd2_x0_mean", total=x0.flatten(1).mean(dim=1).sum(), count=bsz) + metrics.emit_mean("dmd2_x0_std", total=x0.flatten(1).std(dim=1).sum(), count=bsz) + for index, pair in enumerate(batch): + for key, values in per_pair.items(): + metrics.emit_mean(step_metric_key(key, pair["step_index"]), total=values[index], count=1) return loss_sum diff --git a/miles/backends/fsdp_utils/metrics.py b/miles/backends/fsdp_utils/metrics.py index 437f11c9e..7b7317801 100644 --- a/miles/backends/fsdp_utils/metrics.py +++ b/miles/backends/fsdp_utils/metrics.py @@ -37,11 +37,29 @@ "nft_t_mean": MetricReduce.MEAN, "nft_num_timesteps": MetricReduce.MEAN, "dmd2_grad_abs_mean": MetricReduce.MEAN, + "dmd2_score_gap_abs_mean": MetricReduce.MEAN, + "dmd2_teacher_dist_abs_mean": MetricReduce.MEAN, + "dmd2_x0_mean": MetricReduce.MEAN, + "dmd2_x0_std": MetricReduce.MEAN, "x0_mean_abs_diff": MetricReduce.MEAN, "x0_max_abs_diff": MetricReduce.MAX, "x0_rel_max": MetricReduce.MAX, + "peak_memory_allocated_gb": MetricReduce.MAX, + "device_memory_used_gb": MetricReduce.MAX, } +# Also logged per rollout denoising step, under step_metric_key; a step no pair trained on is not logged. +STEP_SCHEMA = { + "dmd2_grad_abs_mean": MetricReduce.MEAN, + "dmd2_score_gap_abs_mean": MetricReduce.MEAN, + "dmd2_teacher_dist_abs_mean": MetricReduce.MEAN, + "x0_mean_abs_diff": MetricReduce.MEAN, +} + + +def step_metric_key(key: str, step_index: int) -> str: + return f"{key}_step{step_index}" + def sigma_bucket_key(stream_name: str, bucket: int, num_buckets: int) -> str: # Keep integer edges as floats (0.0, 1.0) so key names stay stable across bucket counts. @@ -53,10 +71,13 @@ def edge(value: float) -> str: def new_metric_buffer( - group, device: torch.device, components: Collection[str], sigma_buckets: int = 0 + group, device: torch.device, components: Collection[str], sigma_buckets: int = 0, num_rollout_steps: int = 0 ) -> MetricBuffer: """Buffer for one optimizer step, reducing over the DP group.""" schema = dict(SCHEMA) + for key, reduce in STEP_SCHEMA.items(): + for step_index in range(num_rollout_steps): + schema[step_metric_key(key, step_index)] = reduce # SFT loss by sigma bucket: velocity-MSE magnitude varies strongly with the # corruption level, so only bucketed curves are comparable across steps. for stream_name in ("visual", "audio"): diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index fe9b5c0af..2407e509d 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -122,6 +122,7 @@ def __init__(self, args, pg): 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() self.rollout_id = -1 + self._media_threads: list[threading.Thread] = [] self._metric_checker = MetricChecker.maybe_create(args) self._health_monitor = None if self.args.use_fault_tolerance: @@ -154,6 +155,9 @@ def _try_ci_fault_injection(self): def dispose(self): from miles.dashboard import hooks + # A media thread still encoding or uploading would lose the last eval's files. + for thread in self._media_threads: + thread.join() hooks.detach_and_flush() if self._metric_checker is not None: self._metric_checker.dispose() @@ -247,7 +251,8 @@ def eval(self, rollout_id): self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=True) metrics = _log_eval_rollout_data(rollout_id, self.args, data, result.metrics) max_images = self.args.wandb_log_num_images - if max_images > 0: + save_dir = self.args.save_eval_media_dir + if max_images > 0 or save_dir is not None: self._log_images_async( { f"eval_media/{name}_images": payload["samples"] @@ -258,6 +263,7 @@ def eval(self, rollout_id): step_key="eval/step", step_value=compute_rollout_step(self.args, rollout_id), reward_key=self.args.eval_reward_key, + save_dir=None if save_dir is None else f"{save_dir}/eval_{rollout_id}", ) if self._metric_checker is not None: self._metric_checker.on_eval(metrics) @@ -460,7 +466,9 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl def _log_images_async(self, *args, **kwargs) -> None: # mp4 encode + wandb upload take ~1min per log round; keep them off the train barrier - threading.Thread(target=self._log_images, args=args, kwargs=kwargs, daemon=True).start() + thread = threading.Thread(target=self._log_images, args=args, kwargs=kwargs, daemon=True) + thread.start() + self._media_threads.append(thread) def _log_images( self, @@ -470,6 +478,7 @@ def _log_images( step_key: str, step_value: int, reward_key: str | None, + save_dir: str | None = None, ) -> None: """Log per-key sample image grids under their own media namespace. @@ -478,39 +487,45 @@ def _log_images( media panels do not honor ``step_metric`` (they slide on the internal commit step), so panels pile up over a run — keeping them in their own namespace at least groups them in one UI section. + + With ``save_dir``, every sample's file is kept there, next to an ``index.jsonl`` of captions. """ + import json import os import tempfile - import imageio import wandb + fps = self.args.diffusion_fps or 8 with tempfile.TemporaryDirectory() as tmpdir: + media_dir = save_dir or tmpdir + os.makedirs(media_dir, exist_ok=True) log_dict: dict = {} + index_rows = [] for media_key, samples in media_key_to_samples.items(): media = [] - for sample_index, s in enumerate(samples[:max_images]): + for sample_index, s in enumerate(samples if save_dir else samples[:max_images]): t = s.generated_output if t is None or t.ndim != 4: continue - # generated_output is [C, F, H, W]; rescale to uint8 once. - arr = t.float().cpu().numpy() - if arr.max() <= 1.0 + 1e-3: - arr = arr * 255.0 - arr = np.clip(arr, 0, 255).astype(np.uint8) reward = s.reward if not reward_key else (s.reward or {}).get(reward_key) caption = f"{str(s.prompt)[:160]} | reward={reward}" - num_frames = arr.shape[1] - if num_frames > 1: - frames = [arr[:, f, :, :].transpose(1, 2, 0) for f in range(num_frames)] - path = os.path.join(tmpdir, f"{media_key.replace('/', '_')}_{sample_index}.mp4") - imageio.mimsave(path, frames, fps=8, codec="libx264", format="FFMPEG") - media.append(wandb.Video(path, caption=caption, format="mp4", fps=8)) + stem = os.path.join(media_dir, f"{media_key.replace('/', '_')}_{sample_index}") + path = write_sample_media(s, stem, fps=fps) + index_rows.append( + {"file": os.path.basename(path), "prompt": s.prompt, "seed": s.seed, "reward": reward} + ) + if sample_index >= max_images: + continue + if path.endswith(".mp4"): + media.append(wandb.Video(path, caption=caption, format="mp4", fps=fps)) else: - frame = arr[:, 0, :, :].transpose(1, 2, 0) - media.append(wandb.Image(frame, caption=caption)) + media.append(wandb.Image(path, caption=caption)) if media: log_dict[media_key] = media + if save_dir is not None: + with open(os.path.join(save_dir, "index.jsonl"), "w") as f: + f.writelines(json.dumps(row) + "\n" for row in index_rows) if not log_dict: return log_dict[step_key] = step_value @@ -520,6 +535,64 @@ def set_train_parallel_config(self, config: dict): self.train_parallel_config = config +def write_sample_media(sample: Sample, stem: str, *, fps: float) -> str: + """Write ``sample.generated_output`` as ``stem.mp4`` (with its audio track, if any) or ``stem.png``.""" + import imageio + + # generated_output is [C, F, H, W]; rescale to uint8 once. + arr = sample.generated_output.float().cpu().numpy() + if arr.max() <= 1.0 + 1e-3: + arr = arr * 255.0 + arr = np.clip(arr, 0, 255).astype(np.uint8) + if arr.shape[1] == 1: + path = f"{stem}.png" + imageio.imwrite(path, arr[:, 0, :, :].transpose(1, 2, 0)) + return path + path = f"{stem}.mp4" + frames = [arr[:, f, :, :].transpose(1, 2, 0) for f in range(arr.shape[1])] + video_path = path if sample.generated_audio is None else f"{stem}.video.mp4" + imageio.mimsave(video_path, frames, fps=fps, codec="libx264", format="FFMPEG") + if sample.generated_audio is not None: + _mux_audio(video_path, sample.generated_audio, sample.audio_sample_rate, path) + Path(video_path).unlink() + return path + + +def _mux_audio(video_path: str, audio: torch.Tensor, sample_rate: int, path: str) -> None: + import subprocess + + import imageio_ffmpeg + + # [C, L] waveform as interleaved float32 PCM on stdin. + pcm = audio.float().clamp(-1, 1).t().contiguous().numpy().tobytes() + subprocess.run( + [ + imageio_ffmpeg.get_ffmpeg_exe(), + "-y", + "-loglevel", + "error", + "-i", + video_path, + "-f", + "f32le", + "-ar", + str(sample_rate), + "-ac", + str(audio.shape[0]), + "-i", + "pipe:0", + "-c:v", + "copy", + "-c:a", + "aac", + "-shortest", + path, + ], + input=pcm, + check=True, + ) + + def init_rollout_engines(args, pg, all_rollout_engines): if args.train_only: return 0 diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py index b526dff44..a52ef7f54 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -12,7 +12,9 @@ def _resolve_rm_type(args, sample: Sample) -> str: async def async_rm(args, sample: Sample, **kwargs): rm_type = _resolve_rm_type(args, sample) - if rm_type == "ocr": + if rm_type == "none": + return 0.0 + elif rm_type == "ocr": from .ocr import ocr_rm return (await ocr_rm(args, [sample]))[0] @@ -69,6 +71,8 @@ async def batched_async_rm( if samples: rm_types = [_resolve_rm_type(args, sample) for sample in samples] + if all(rm_type == "none" for rm_type in rm_types): + return [0.0] * len(samples) if all(rm_type == "pickscore" for rm_type in rm_types): from .pickscore import pickscore_rm diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 6b23cf662..41ddc0bca 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -46,6 +46,7 @@ def build_rollout_sampling_params( "num_inference_steps": num_steps, "guidance_scale": args.diffusion_guidance_scale, "true_cfg_scale": args.diffusion_true_cfg_scale, + "rollout_video_dtype": args.rollout_video_dtype, } if neg is not None: @@ -61,13 +62,15 @@ 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, } ) if args.use_output_store: sampling_params["rollout_return_outputs_via_store"] = True + if args.rm_type == "none" and args.custom_rm_path is None: + # No reward reads the decoded media, so the engine skips the VAE decode of training rollouts. + sampling_params["rollout_return_generated_output"] = False if args.diffusion_output_num_frames is not None: sampling_params["num_frames"] = int(args.diffusion_output_num_frames) @@ -415,12 +418,13 @@ async def generate_rollout_async( def release_eval_sample_media(args: Namespace, sample: Sample) -> None: - """Drop a scored eval sample's media, which only image logging and debug dumps still read.""" - if args.save_debug_rollout_data is not None: + """Drop a scored eval sample's media, which only image logging, saved eval media and debug dumps still read.""" + if args.save_debug_rollout_data is not None or args.save_eval_media_dir is not None: return if sample.index < args.wandb_log_num_images: return sample.generated_output = None + sample.generated_audio = None # eval only diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index b77bae51f..a95856749 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -308,7 +308,7 @@ def add_rollout_arguments(parser): choices=["keep", "uint8"], default="keep", help=( - "Dtype of the decoded video in rollout responses. 'keep' returns the " + "Dtype of the decoded video in rollout and eval 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." @@ -318,7 +318,8 @@ def add_rollout_arguments(parser): "--diffusion-fps", type=float, default=None, - help="Video fps for rollout; None for image models.", + help="Video fps for rollout, also the fps logged and saved media play at (8 when unset); " + "None for image models.", ) parser.add_argument( "--diffusion-output-num-frames", @@ -1384,6 +1385,13 @@ def add_debug_arguments(parser): "The file will be saved to `save_debug_rollout_data.format(rollout_id)`." ), ) + parser.add_argument( + "--save-eval-media-dir", + type=str, + default=None, + help="Keep every eval sample's video (with its audio, if the engine returns it) or image under " + "/eval_/, next to an index.jsonl of prompts, seeds and rewards.", + ) parser.add_argument( "--load-debug-rollout-data", type=str, @@ -1464,8 +1472,9 @@ def add_reward_model_arguments(parser): "--rm-type", type=str, default=None, - help="Built-in reward (pickscore / hps / hpsv3 / ocr / openai_api), or custom_api for a configurable actor. " - "Ignored when --custom-rm-path is set.", + help="Built-in reward (pickscore / hps / hpsv3 / ocr / openai_api), custom_api for a configurable actor, " + "or none for a constant 0 when the loss reads no reward; none also has the engine skip decoding " + "training rollouts, whose media nothing reads. Ignored when --custom-rm-path is set.", ) for rm_type in ("custom_api", "openai_api"): parser.add_argument( diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 0ef09d353..1d2cd3fa1 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -224,6 +224,8 @@ def apply_rollout_image_response( deserialize_func = _store_deserialize_func(sample.rollout_store_ref, deserialize_func) sample.generated_output = _normalize_generated_output(deserialize_func(body.get("generated_output"))) + sample.generated_audio = deserialize_func(body.get("generated_audio")) + sample.audio_sample_rate = body.get("audio_sample_rate") # Eval-mode rollout (rollout=False) sends no log_probs; train-mode always does. sample.rollout_log_probs = deserialize_func(body.get("rollout_log_probs")) sample.rollout_debug_tensors = _parse_rollout_debug_tensors( diff --git a/miles/utils/types.py b/miles/utils/types.py index bfce07c05..8b84fb645 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -83,6 +83,9 @@ class Sample: seed: int | None = None # Eager tensor on CPU. Image rollout shape: ``[C, T, H, W]`` (``T==1`` typical). 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 + audio_sample_rate: int | None = None rollout_log_probs: torch.Tensor | None = None rollout_debug_tensors: RolloutDebugTensors | None = None denoising_env: DenoisingEnv | None = None diff --git a/requirements.txt b/requirements.txt index e94e2257f..56d20fe36 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,9 @@ datasets==4.4.2 diffusers==0.41.0 fastapi==0.128.8 httpx[http2]==0.28.1 +# miles/ray/rollout.py muxes eval audio through imageio_ffmpeg's bundled ffmpeg; the GPU image only had it +# transitively (cosmos-guardrail, moviepy) and the hosted CPU suite not at all. +imageio-ffmpeg==0.5.1 ltx-core @ git+https://github.com/Lightricks/LTX-2.git@780984275fd47128b02bef9b5c085404276866ee#subdirectory=packages/ltx-core mcp[cli]==1.25.0 memray==1.19.1 diff --git a/scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py b/scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py new file mode 100644 index 000000000..e13c3ed20 --- /dev/null +++ b/scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py @@ -0,0 +1,265 @@ +"""MiniMax H3 t2va 4-step DMD2 distillation, 2 nodes x 8 GPU, data-free (prompts only). + +Multi-node: this script submits into an EXISTING ray cluster of 2 x 8 colocated train/rollout GPUs -- +see docs/user-guide/launch-script.md (Multi-node training) for the bring-up, then run it on the head +node with MILES_SCRIPT_EXTERNAL_RAY=1. + +Recipe: FastVideo's FastH3 Preview v1 Dense / Data-Free (checkpoint_metadata.json of +FastVideo/FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree, released at step 1000), mapped onto +--loss-type dmd2: + * full-parameter student and critic, both from MiniMax-H3, frozen MiniMax-H3 teacher; fp32 master, + bf16 forward; AdamW (0.9, 0.999), weight decay 0.01, constant lr 2e-6 for both, grad clip 1.0; + * 4 student steps on the uniform shift-12 grid (FastH3's rungs 999/749/500/250 within 1.4e-4 in + sigma), backward simulation by stopping each rollout at a random step, plain Euler (ode) rollout; + * teacher unguided (H3 has no CFG; FastH3 real_score_guidance_scale 1.0), critic loss in x0 space; + * scoring and critic sigma: unshifted range 0.001-0.999 read through shift 5. FastH3 sets + score_timestep_shift 2.4 on its own clock, documented as FastGen's shift 5 (12 / 5 = 2.4); + * one rollout = 5 critic steps of 64 samples + 1 student step on the last 64 (FastH3's + generator_update_interval 5 at global batch 64), so 1000 FastH3 steps = 200 rollouts. +Deviations: 768p 16:9 only (FastH3 mixes shapes); only the video stream is distilled (the rollout +records no audio trajectory), so audio follows whatever the shared layers learn -- listen to the media. + +Prompts: the 248,219 unique VidProM prompts (248,221 lines) of vidprom_filtered_extended.txt (the list Self-Forcing, NVlabs/rcm +and FastVideo's Wan data-free DMD train on), downloaded at a pinned revision and split here +deterministically: train 247,219 / test 1,000 / eval_media 26 (16 test prompts + 10 FastGen-DMD +demo prompts that also describe sound). The output hashes are pinned, so every node builds the same files. + +What to watch: eval media at rollout 0 is the base model at 4 steps; every --eval-interval rollouts +the student renders the same prompts and seeds. Videos (with audio once the engine returns it) go to +W&B eval_media/ and to //eval_media/eval_/. References on the same prompts +and seeds, for side-by-side comparison: + MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --reference teacher + MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --reference fasth3 +teacher = MiniMax-H3 at its default 50 steps; fasth3 = the FastH3 Dense / Data-Free student at 4 steps. + +Usage (after the multi-node bring-up, on the head node): + MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py +""" + +import hashlib +import json +import random +from dataclasses import dataclass +from pathlib import Path + +import typer + +import miles.utils.external_utils.command_utils as U + +MODEL = "MiniMaxAI/MiniMax-H3" +FASTH3_MODEL = "FastVideo/FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree" +WANDB_PROJECT = "miles-diffusion-dmd2" + +PROMPT_REPO = "gdhe17/Self-Forcing" +PROMPT_FILE = "vidprom_filtered_extended.txt" +PROMPT_REVISION = "2f8b779212da279d212c22a509b66ad6552f350e" +PROMPT_SHA256 = "7896742f468bc8aef9e4547424d1ce0a951acdb2a82233790155401a99bf5aa5" +PROMPT_SPLIT_SEED = 1234 +NUM_TEST_PROMPTS = 1000 +NUM_EVAL_TEST_PROMPTS = 16 +# FastGen-DMD scripts/inference/prompts/flux3_t2va_demo.txt: short prompts that also describe sound. +DEMO_PROMPTS = [ + "a small red boat bobbing on gentle ocean waves at sunset, seagulls calling in the distance", + "a golden retriever running through autumn leaves in a park, leaves crunching underfoot", + "close-up of rain falling on a city street at night, neon reflections, soft rain sounds", + "a steam locomotive pulling into a rural station, whistle blowing, steam hissing", + "a campfire crackling on a beach under a starry sky, waves breaking softly", + "a busy street market in Marrakech, vendors calling out, colorful fabrics swaying", + "time-lapse of clouds rolling over snowy mountain peaks, wind howling", + "a jazz drummer playing a fast solo in a smoky club, cymbals shimmering", + "children playing soccer on a sunny field, laughter and a whistle", + "a waterfall cascading into a jungle pool, birds singing, water roaring", +] +SPLIT_SHA256 = { + "train.jsonl": "ec9ecd057ec5d3aa4b383224bc9071edfb73cf4a77b6594a5c9708068c39bc71", + "test.jsonl": "be22ea8ecdbe6d909099cd943fcd20cb4e5b0870160212b3483ad085e9f0684c", + "eval_media.jsonl": "9a9faf270a94fe98adbe28dfcb059a6acbc6d5f3d5f8da7a3e9a51c57e18fe60", +} + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + num_rollout: int = 200 + # 5 critic steps x 64 samples; the student steps on the last 64. + rollout_batch_size: int = 320 + eval_interval: int = 5 + save_interval: int = 50 + duration_seconds: float = 5.0 + data_dir: str = "/root/datasets" + # "teacher" or "fasth3": render the eval prompts with that model instead of training. + reference: str = "" + extra_args: str = "" + + +def _sha256(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def _write_jsonl(path: Path, prompts: list[str]) -> None: + with open(path, "w", encoding="utf-8") as f: + for prompt in prompts: + f.write(json.dumps({"input": prompt}, ensure_ascii=False) + "\n") + + +def prepare(args: ScriptArgs) -> str: + """Download the pinned prompt list and split it; returns the directory of the jsonl files.""" + from huggingface_hub import hf_hub_download + + root = Path(args.data_dir) / "vidprom-extended" + prompt_dir = root / "prompts" + if all((prompt_dir / name).exists() and _sha256(prompt_dir / name) == h for name, h in SPLIT_SHA256.items()): + return str(prompt_dir) + + raw = Path(hf_hub_download(PROMPT_REPO, PROMPT_FILE, revision=PROMPT_REVISION, local_dir=root / "raw")) + if _sha256(raw) != PROMPT_SHA256: + raise RuntimeError(f"{raw} does not match the pinned sha256 {PROMPT_SHA256}") + prompts = list(dict.fromkeys(line.strip() for line in raw.read_text(encoding="utf-8").splitlines())) + prompts = [prompt for prompt in prompts if prompt] + order = list(range(len(prompts))) + random.Random(PROMPT_SPLIT_SEED).shuffle(order) + shuffled = [prompts[i] for i in order] + test = shuffled[:NUM_TEST_PROMPTS] + splits = { + "train.jsonl": shuffled[NUM_TEST_PROMPTS:], + "test.jsonl": test, + "eval_media.jsonl": test[:NUM_EVAL_TEST_PROMPTS] + DEMO_PROMPTS, + } + prompt_dir.mkdir(parents=True, exist_ok=True) + for name, split in splits.items(): + _write_jsonl(prompt_dir / name, split) + if _sha256(prompt_dir / name) != SPLIT_SHA256[name]: + raise RuntimeError(f"{prompt_dir / name} does not match the pinned split sha256") + return str(prompt_dir) + + +def execute(args: ScriptArgs, prompt_dir: str) -> None: + run_name = f"diffusion_dmd2_h3_t2va_{args.reference or 'student'}_{U.create_run_id()}" + run_dir = f"{args.output_dir}/{run_name}" + + rollout_args = ( + "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout " + f"--prompt-data {prompt_dir}/train.jsonl " + "--input-key input " + f"--rollout-batch-size {args.rollout_batch_size} " + "--n-samples-per-prompt 1 " + "--num-steps-per-rollout 1 " + "--rollout-microgroup-size 1 " + "--rm-type none " + "--rollout-video-dtype uint8 " + ) + + diffusion_args = ( + "--diffusion-task t2va " + "--diffusion-num-steps 4 " + "--diffusion-guidance-scale 1.0 " + "--diffusion-h3-aspect-ratio 16:9 " + f"--diffusion-h3-duration-seconds {args.duration_seconds} " + "--diffusion-flow-shift 12.0 " + "--diffusion-audio-flow-shift 3.0 " + # H3 drops fps from the request; this sets the playback rate of the logged and saved media. + "--diffusion-fps 24 " + ) + + eval_args = ( + f"--eval-prompt-data dmd2_media {prompt_dir}/eval_media.jsonl " + "--n-samples-per-eval-prompt 1 " + f"--save-eval-media-dir {run_dir}/eval_media " + ) + + # Every eval video; training rollouts carry no media, since --rm-type none skips their decode. + wandb_args = U.get_default_wandb_args( + __file__, + run_id=run_name, + project=WANDB_PROJECT, + wandb_log_num_images=NUM_EVAL_TEST_PROMPTS + len(DEMO_PROMPTS), + ) + + sglang_args = ( + "--use-miles-router " + "--sglang-server-concurrency 2 " + "--sglang-tp-size 2 " + "--sglang-sp-degree 1 " + "--sglang-ulysses-degree 1 " + "--sglang-ring-degree 1 " + "--sglang-dit-precision bf16 " + "--update-weight-buffer-size 4294967296 " + ) + + topology_args = ( + "--actor-num-nodes 2 " + "--actor-num-gpus-per-node 8 " + "--num-gpus-per-node 8 " + "--rollout-num-gpus 16 " + "--rollout-num-gpus-per-engine 2 " + "--rollout-parser-num-workers 16 " + "--colocate " + # H3 denoising blocks the engine's event loop for tens of seconds per sample. + "--rollout-health-check-interval 60 " + "--miles-router-health-check-failure-threshold 30 " + ) + + if args.reference: + if args.reference == "teacher": + checkpoint, eval_steps = MODEL, 50 + elif args.reference == "fasth3": + checkpoint, eval_steps = FASTH3_MODEL, 4 + else: + raise ValueError(f"--reference must be teacher or fasth3, got {args.reference!r}") + mode_args = ( + f"--hf-checkpoint {checkpoint} " + # Every engine fetches the checkpoint by its Hub id; the FastH3 name matches no H3 pattern. + "--diffusion-model-family h3 " + "--debug-rollout-only " + "--num-rollout 0 " + "--eval-interval 1 " + f"--diffusion-eval-num-steps {eval_steps} " + ) + else: + mode_args = ( + f"--hf-checkpoint {MODEL} " + f"--save {run_dir}/ckpt " + f"--save-interval {args.save_interval} " + f"--num-rollout {args.num_rollout} " + f"--eval-interval {args.eval_interval} " + "--loss-type dmd2 " + f"--teacher-load {MODEL} " + "--dmd2-teacher-guidance-scale 1.0 " + "--dmd2-critic-steps-per-rollout 5 " + "--dmd2-score-sigma-range 0.001 0.999 " + "--dmd2-critic-loss-space x0 " + "--fsdp-flow-shift visual=5.0 " + "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.stop_at_random_step " + "--diffusion-sde-type ode " + "--lr 2e-6 " + "--critic-lr 2e-6 " + "--weight-decay 0.01 " + "--train-backend fsdp " + "--fsdp-master-dtype fp32 " + "--fsdp-reduce-dtype fp32 " + "--diffusion-forward-dtype bf16 " + "--gradient-checkpointing " + ) + + U.execute_train( + train_args=( + f"{mode_args} {rollout_args} {diffusion_args} {eval_args} {wandb_args} {sglang_args} " + f"{topology_args} {args.extra_args}" + ), + num_gpus_per_node=8, + config=args, + extra_env_vars={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs) -> None: + prompt_dir = prepare(args) + execute(args, prompt_dir) + + +if __name__ == "__main__": + typer.run(main) diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py index 6aed8b941..eca56850e 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py @@ -41,6 +41,7 @@ test_surrogate_gradient_matches_the_reference same gradient as the reference's batch-mean MSE test_critic_matching_the_teacher_gives_no_gradient x0_fake == x0_real ──► zero loss and gradient test_gaussian_student_moves_toward_the_teacher analytic scores: sign and normalization move mu + test_health_metrics_are_recorded_overall_and_per_step score gap, teacher distance, x0 stats, split by stop step """ from tests.ci.ci_register import register_cpu_ci @@ -358,13 +359,14 @@ def _generator_pairs(num_pairs): "sigma": 0.75 - 0.25 * n, "denoising_env": DenoisingEnv(pos_cond_kwargs=_cond(1.0 + n), neg_cond_kwargs=None), "sample_index": 100 + n, + "step_index": n, } for n in range(num_pairs) ] def _metrics(): - return new_metric_buffer(None, torch.device("cpu"), ["transformer"]) + return new_metric_buffer(None, torch.device("cpu"), ["transformer"], num_rollout_steps=4) def _loss(ctx, pairs, x0): @@ -485,3 +487,26 @@ def test_trainer_and_rollout_x0_diff_is_recorded(): reduced = metrics.reduce() assert reduced["x0_mean_abs_diff"] == pytest.approx(0.25) assert reduced["x0_max_abs_diff"] == pytest.approx(0.25) + + +def test_health_metrics_are_recorded_overall_and_per_step(): + ctx = _ctx() + pairs = _generator_pairs(2) + # Pair n stopped at step n: x0 = 0.5 + n against a teacher x0 of 0 and a critic x0 of 0.25. + x0 = torch.stack([pair["x0"] for pair in pairs]) + torch.tensor([0.5, 1.5]).view(2, 1, 1, 1) + ctx.external_data = {"critic": _scores(pairs, torch.zeros(2, *LATENT_SHAPE), torch.full((2, *LATENT_SHAPE), 0.25))} + + _, metrics = _loss(ctx, pairs, x0) + + reduced = metrics.reduce() + assert reduced["dmd2_score_gap_abs_mean"] == pytest.approx(0.25) + assert reduced["dmd2_teacher_dist_abs_mean"] == pytest.approx(1.0) + assert reduced["dmd2_x0_mean"] == pytest.approx(1.0) + assert reduced["dmd2_x0_std"] == pytest.approx(0.0, abs=1e-6) + for step, distance in enumerate([0.5, 1.5]): + assert reduced[f"dmd2_teacher_dist_abs_mean_step{step}"] == pytest.approx(distance) + assert reduced[f"dmd2_grad_abs_mean_step{step}"] == pytest.approx(0.25 / distance) + assert reduced[f"dmd2_score_gap_abs_mean_step{step}"] == pytest.approx(0.25) + assert reduced[f"x0_mean_abs_diff_step{step}"] == pytest.approx(distance) + # No pair stopped at steps 2 or 3, so they log nothing. + assert not any(key.endswith(("_step2", "_step3")) for key in reduced) diff --git a/tests/fast/backends/fsdp_utils/test_train_input_hook.py b/tests/fast/backends/fsdp_utils/test_train_input_hook.py index b6525ac0d..52b85f9b5 100644 --- a/tests/fast/backends/fsdp_utils/test_train_input_hook.py +++ b/tests/fast/backends/fsdp_utils/test_train_input_hook.py @@ -65,6 +65,7 @@ def _actor(train_input, train_output, forwarded): diffusion_recompute_old_log_prob=False, debug_skip_optimizer_step=True, log_loss_sigma_bucket=0, + diffusion_num_steps=4, ), _forward_dtype=torch.float32, parallel_state=SimpleNamespace(dp_rank=0, dp_group=None), diff --git a/tests/fast/rollout/test_eval_media.py b/tests/fast/rollout/test_eval_media.py new file mode 100644 index 000000000..a3dfb68f5 --- /dev/null +++ b/tests/fast/rollout/test_eval_media.py @@ -0,0 +1,69 @@ +"""Logged and saved media: videos play at --diffusion-fps and carry the engine's audio; eval media can be kept on disk.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=20, suite="stage-a-cpu", labels=[]) + +import json +import subprocess +from argparse import Namespace +from types import SimpleNamespace + +import imageio_ffmpeg +import pytest +import torch + +from miles.ray import rollout +from miles.utils.types import Sample + + +def _streams(path): + probe = subprocess.run([imageio_ffmpeg.get_ffmpeg_exe(), "-i", path], capture_output=True, text=True) + return {kind for kind in ("Video", "Audio") if f": {kind}:" in probe.stderr} + + +def _video_sample(**fields): + return Sample(prompt="a boat", seed=7, generated_output=torch.rand(3, 12, 16, 16), **fields) + + +def test_video_plays_at_the_given_fps(tmp_path): + path = rollout.write_sample_media(_video_sample(), str(tmp_path / "clip"), fps=24) + + frames, seconds = imageio_ffmpeg.count_frames_and_secs(path) + assert path.endswith(".mp4") and frames == 12 + assert seconds == pytest.approx(0.5, abs=0.05) + assert _streams(path) == {"Video"} + + +def test_audio_is_muxed_into_the_video(tmp_path): + sample = _video_sample(generated_audio=torch.rand(2, 24000) * 2 - 1, audio_sample_rate=48000) + + path = rollout.write_sample_media(sample, str(tmp_path / "clip"), fps=24) + + assert _streams(path) == {"Video", "Audio"} + assert sorted(p.name for p in tmp_path.iterdir()) == ["clip.mp4"] + + +def test_a_single_frame_is_saved_as_an_image(tmp_path): + sample = Sample(generated_output=torch.rand(3, 1, 16, 16)) + + assert rollout.write_sample_media(sample, str(tmp_path / "image"), fps=8).endswith(".png") + + +def test_saved_eval_media_keeps_every_sample_without_wandb(tmp_path): + manager = SimpleNamespace(args=Namespace(diffusion_fps=24)) + samples = [_video_sample(), Sample(prompt="no media")] + + rollout.RolloutManager.__ray_metadata__.modified_class._log_images( + manager, + {"eval_media/demo_images": samples}, + max_images=0, + step_key="eval/step", + step_value=0, + reward_key=None, + save_dir=str(tmp_path), + ) + + assert sorted(p.name for p in tmp_path.iterdir()) == ["eval_media_demo_images_0.mp4", "index.jsonl"] + rows = [json.loads(line) for line in (tmp_path / "index.jsonl").read_text().splitlines()] + assert rows == [{"file": "eval_media_demo_images_0.mp4", "prompt": "a boat", "seed": 7, "reward": None}] diff --git a/tests/fast/rollout/test_release_eval_sample_media.py b/tests/fast/rollout/test_release_eval_sample_media.py index 174ebf708..9bba64990 100644 --- a/tests/fast/rollout/test_release_eval_sample_media.py +++ b/tests/fast/rollout/test_release_eval_sample_media.py @@ -11,21 +11,30 @@ def make_sample(index: int) -> Sample: - return Sample(index=index, generated_output=torch.zeros(3, 1, 4, 4)) + return Sample(index=index, generated_output=torch.zeros(3, 1, 4, 4), generated_audio=torch.zeros(2, 8)) + + +def make_args(save_debug_rollout_data=None, wandb_log_num_images=2, save_eval_media_dir=None) -> Namespace: + return Namespace( + save_debug_rollout_data=save_debug_rollout_data, + wandb_log_num_images=wandb_log_num_images, + save_eval_media_dir=save_eval_media_dir, + ) def test_media_released_when_not_logged(): sample = make_sample(index=3) - release_eval_sample_media(Namespace(save_debug_rollout_data=None, wandb_log_num_images=2), sample) + release_eval_sample_media(make_args(), sample) assert sample.generated_output is None + assert sample.generated_audio is None def test_media_kept_for_logged_samples(): sample = make_sample(index=1) - release_eval_sample_media(Namespace(save_debug_rollout_data=None, wandb_log_num_images=2), sample) + release_eval_sample_media(make_args(), sample) assert sample.generated_output is not None @@ -33,6 +42,15 @@ def test_media_kept_for_logged_samples(): def test_media_kept_when_debug_dump_enabled(): sample = make_sample(index=3) - release_eval_sample_media(Namespace(save_debug_rollout_data="/tmp/dump", wandb_log_num_images=0), sample) + release_eval_sample_media(make_args(save_debug_rollout_data="/tmp/dump", wandb_log_num_images=0), sample) + + assert sample.generated_output is not None + + +def test_media_kept_when_eval_media_is_saved(): + sample = make_sample(index=3) + + release_eval_sample_media(make_args(wandb_log_num_images=0, save_eval_media_dir="/tmp/media"), sample) assert sample.generated_output is not None + assert sample.generated_audio is not None diff --git a/tests/fast/rollout/test_rm_none.py b/tests/fast/rollout/test_rm_none.py new file mode 100644 index 000000000..9ab131873 --- /dev/null +++ b/tests/fast/rollout/test_rm_none.py @@ -0,0 +1,27 @@ +"""--rm-type none: a constant 0 reward for losses that read none, with no scorer behind it.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) + +import asyncio +from argparse import Namespace + +from miles.rollout.rm_hub import async_rm, batched_async_rm +from miles.utils.types import Sample + + +def _args(rm_type): + return Namespace(rm_type=rm_type, custom_rm_path=None) + + +def test_none_rewards_every_sample_with_zero(): + samples = [Sample(index=index) for index in range(3)] + + assert asyncio.run(batched_async_rm(_args("none"), samples)) == [0.0, 0.0, 0.0] + + +def test_per_sample_none_overrides_the_global_reward(): + sample = Sample(index=0, metadata={"rm_type": "none"}) + + assert asyncio.run(async_rm(_args("pickscore"), sample)) == 0.0 diff --git a/tests/fast/rollout/test_rollout_negative_prompt.py b/tests/fast/rollout/test_rollout_negative_prompt.py index f7feb1aba..3d6f5436e 100644 --- a/tests/fast/rollout/test_rollout_negative_prompt.py +++ b/tests/fast/rollout/test_rollout_negative_prompt.py @@ -34,6 +34,8 @@ def _args(**overrides): rollout_video_dtype="keep", train_pipeline_config_path=None, use_output_store=False, + rm_type="pickscore", + custom_rm_path=None, ) values.update(overrides) return Namespace(**values) diff --git a/tests/fast/rollout/test_rollout_payload_media.py b/tests/fast/rollout/test_rollout_payload_media.py new file mode 100644 index 000000000..070a77c24 --- /dev/null +++ b/tests/fast/rollout/test_rollout_payload_media.py @@ -0,0 +1,55 @@ +"""Which media the engine decodes and ships: training rollouts skip the decode when no reward reads it.""" + +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 + +from miles.rollout.sglang_diffusion_rollout import build_rollout_sampling_params + + +def _args(**overrides): + values = dict( + diffusion_negative_prompt=None, + diffusion_generator_device=None, + diffusion_quality="exact", + diffusion_eval_num_steps=None, + diffusion_num_steps=4, + diffusion_width=832, + diffusion_height=480, + diffusion_guidance_scale=1.0, + diffusion_guidance_scale_2=None, + diffusion_true_cfg_scale=None, + diffusion_output_num_frames=None, + diffusion_fps=None, + diffusion_sde_type="ode", + diffusion_noise_level=0.7, + diffusion_log_prob_no_const=False, + diffusion_debug_mode=False, + rollout_video_dtype="uint8", + train_pipeline_config_path=None, + rm_type="none", + custom_rm_path=None, + use_output_store=False, + ) + values.update(overrides) + return Namespace(**values) + + +def test_training_rollouts_skip_the_decode_when_no_reward_reads_it(): + assert build_rollout_sampling_params(_args())["rollout_return_generated_output"] is False + # Eval renders are what gets looked at, so they always decode. + assert "rollout_return_generated_output" not in build_rollout_sampling_params(_args(), evaluation=True) + + +@pytest.mark.parametrize("overrides", [{"rm_type": "pickscore"}, {"custom_rm_path": "my.rm"}]) +def test_a_reward_keeps_the_decode(overrides): + assert "rollout_return_generated_output" not in build_rollout_sampling_params(_args(**overrides)) + + +def test_video_dtype_applies_to_eval_too(): + for evaluation in (False, True): + assert build_rollout_sampling_params(_args(), evaluation=evaluation)["rollout_video_dtype"] == "uint8" diff --git a/tests/fast/utils/test_diffusion_rollout_response.py b/tests/fast/utils/test_diffusion_rollout_response.py index 4a67185cb..ad211d4da 100644 --- a/tests/fast/utils/test_diffusion_rollout_response.py +++ b/tests/fast/utils/test_diffusion_rollout_response.py @@ -5,7 +5,8 @@ import pytest import torch -from miles.utils.diffusion_rollout_response import _normalize_generated_output +from miles.utils.diffusion_rollout_response import _normalize_generated_output, apply_rollout_image_response +from miles.utils.types import Sample @pytest.mark.parametrize( @@ -103,3 +104,15 @@ def _tensor_bytes(tensor: torch.Tensor) -> bytes: from safetensors.torch import save return save({"t": tensor}) + + +def test_response_audio_fills_the_sample(): + audio = torch.zeros(2, 480) + body = {"generated_audio": audio, "audio_sample_rate": 48000} + + sample = apply_rollout_image_response(Sample(), body, deserialize_func=lambda value: value) + + assert sample.generated_audio is audio + assert sample.audio_sample_rate == 48000 + silent = apply_rollout_image_response(Sample(), {}, deserialize_func=lambda value: value) + assert silent.generated_audio is None and silent.audio_sample_rate is None