diff --git a/tests/workers/config/test_diffusion_config_on_cpu.py b/tests/workers/config/test_diffusion_config_on_cpu.py index 5ee271683..a67bc13d8 100644 --- a/tests/workers/config/test_diffusion_config_on_cpu.py +++ b/tests/workers/config/test_diffusion_config_on_cpu.py @@ -168,6 +168,19 @@ def test_defaults(self): class TestDiffusionRolloutConfig: + def test_prompt_embed_cache_config(self): + cfg = DiffusionRolloutConfig( + name="vllm_omni", + enable_prompt_embed_cache=True, + prompt_embed_cache_size=64, + ) + assert cfg.enable_prompt_embed_cache is True + assert cfg.prompt_embed_cache_size == 64 + + def test_invalid_prompt_embed_cache_size_raises(self): + with pytest.raises(ValueError, match="prompt_embed_cache_size must be positive"): + DiffusionRolloutConfig(name="vllm_omni", prompt_embed_cache_size=0) + def test_invalid_rollout_adapter_raises(self): with pytest.raises(ValueError, match="Invalid diffusion rollout rollout_adapter"): DiffusionRolloutConfig(name="vllm_omni", rollout_adapter="bogus") diff --git a/verl_omni/experimental/qwen_image_flow_grpo_stepwise/vllm_omni_rollout_adapter.py b/verl_omni/experimental/qwen_image_flow_grpo_stepwise/vllm_omni_rollout_adapter.py index f8a8af831..37eeff344 100644 --- a/verl_omni/experimental/qwen_image_flow_grpo_stepwise/vllm_omni_rollout_adapter.py +++ b/verl_omni/experimental/qwen_image_flow_grpo_stepwise/vllm_omni_rollout_adapter.py @@ -117,6 +117,10 @@ def encode_prompt( ``(B * num_images_per_prompt, L, D)`` and ``(B * num_images_per_prompt, L)`` respectively. """ + if isinstance(prompt_ids, list): + prompt_ids = torch.tensor(prompt_ids, device=self.device) + if isinstance(attention_mask, list): + attention_mask = torch.tensor(attention_mask, device=self.device) prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids attention_mask = ( attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask @@ -173,7 +177,7 @@ def _tokenize_text_prompt(self, text: str | list[str]): truncation=True, return_tensors="pt", ).to(self.device) - return tokens.input_ids, tokens.attention_mask + return tokens.input_ids.tolist(), tokens.attention_mask.tolist() def prepare_encode( self, @@ -191,12 +195,6 @@ def prepare_encode( state.prompts or [] ) - # Normalize list inputs to tensors on device. - if isinstance(prompt_ids, list): - prompt_ids = torch.tensor(prompt_ids, device=self.device) - if isinstance(negative_prompt_ids, list): - negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) - if prompt_ids is None: raise ValueError( "QwenImagePipelineWithLogProbStepwise.prepare_encode requires either " @@ -221,10 +219,18 @@ def prepare_encode( self._current_timestep = None self._interrupt = False - if prompt_ids is not None: + prompt_embed_cache = getattr(self, "_prompt_embed_cache", None) + prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled) + if not prompt_embed_cache_enabled: + if isinstance(prompt_ids, list): + prompt_ids = torch.tensor(prompt_ids, device=self.device) + if isinstance(negative_prompt_ids, list): + negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) + + if isinstance(prompt_ids, torch.Tensor): batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1 else: - batch_size = 1 + batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1 has_neg_prompt = negative_prompt_ids is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt diff --git a/verl_omni/pipelines/qwen_image_diffusion_nft/vllm_omni_rollout_adapter.py b/verl_omni/pipelines/qwen_image_diffusion_nft/vllm_omni_rollout_adapter.py index 3f16d7187..abdfb6ef0 100644 --- a/verl_omni/pipelines/qwen_image_diffusion_nft/vllm_omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen_image_diffusion_nft/vllm_omni_rollout_adapter.py @@ -72,13 +72,19 @@ def _prepare_token_id_generation_context( self._current_timestep = None self._interrupt = False - if isinstance(prompt_ids, list): - prompt_ids = torch.tensor(prompt_ids, device=self.device) - if isinstance(negative_prompt_ids, list): - negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) + prompt_embed_cache = getattr(self, "_prompt_embed_cache", None) + prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled) + if not prompt_embed_cache_enabled: + if isinstance(prompt_ids, list): + prompt_ids = torch.tensor(prompt_ids, device=self.device) + if isinstance(negative_prompt_ids, list): + negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) if prompt_ids is not None: - batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1 + if isinstance(prompt_ids, torch.Tensor): + batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1 + else: + batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1 elif prompt_embeds is not None: batch_size = prompt_embeds.shape[0] else: diff --git a/verl_omni/pipelines/qwen_image_dpo/vllm_omni_rollout_adapter.py b/verl_omni/pipelines/qwen_image_dpo/vllm_omni_rollout_adapter.py index 0b2a0994e..fdab89400 100644 --- a/verl_omni/pipelines/qwen_image_dpo/vllm_omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen_image_dpo/vllm_omni_rollout_adapter.py @@ -82,6 +82,10 @@ def encode_prompt( prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): + if isinstance(prompt_ids, list): + prompt_ids = torch.tensor(prompt_ids, device=self.device) + if isinstance(attention_mask, list): + attention_mask = torch.tensor(attention_mask, device=self.device) prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids attention_mask = ( attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask @@ -150,16 +154,21 @@ def forward( self._current_timestep = None self._interrupt = False + prompt_embed_cache = getattr(self, "_prompt_embed_cache", None) + prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled) if prompt_ids is not None: - if isinstance(prompt_ids, list): + if not prompt_embed_cache_enabled and isinstance(prompt_ids, list): prompt_ids = torch.tensor(prompt_ids, device=self.device) - batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1 + if isinstance(prompt_ids, torch.Tensor): + batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1 + else: + batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1 elif prompt_embeds is not None: batch_size = prompt_embeds.shape[0] else: return DiffusionOutput(output=None, custom_output={}) - if isinstance(negative_prompt_ids, list): + if not prompt_embed_cache_enabled and isinstance(negative_prompt_ids, list): negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) has_neg_prompt = negative_prompt_ids is not None or ( diff --git a/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py b/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py index 2ac1eda1a..f380e7985 100644 --- a/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py @@ -307,10 +307,15 @@ def forward( self._current_timestep = None self._interrupt = False + prompt_embed_cache = getattr(self, "_prompt_embed_cache", None) + prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled) if prompt_token_ids is not None: - if isinstance(prompt_token_ids, list): + if not prompt_embed_cache_enabled and isinstance(prompt_token_ids, list): prompt_token_ids = torch.tensor(prompt_token_ids, device=self.device) - batch_size = prompt_token_ids.shape[0] if prompt_token_ids.ndim == 2 else 1 + if isinstance(prompt_token_ids, torch.Tensor): + batch_size = prompt_token_ids.shape[0] if prompt_token_ids.ndim == 2 else 1 + else: + batch_size = len(prompt_token_ids) if prompt_token_ids and isinstance(prompt_token_ids[0], list) else 1 elif prompt_embeds is not None: batch_size = prompt_embeds.shape[0] else: @@ -318,7 +323,7 @@ def forward( # Return a minimal dummy output to avoid crashing. return DiffusionOutput(output=None, custom_output={}) - if isinstance(negative_prompt_ids, list): + if not prompt_embed_cache_enabled and isinstance(negative_prompt_ids, list): negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) has_neg_prompt = negative_prompt_ids is not None or ( diff --git a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml index 69e001b30..5a4434305 100644 --- a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml @@ -207,6 +207,8 @@ actor_rollout_ref: enable: false engine_kwargs: vllm_omni: {} + enable_prompt_embed_cache: false + prompt_embed_cache_size: 32 val_kwargs: _target_: verl_omni.workers.config.diffusion.DiffusionSamplingConfig 'n': 1 diff --git a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml index eed93cf0a..e5a663f3c 100644 --- a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml @@ -248,6 +248,8 @@ actor_rollout_ref: enable: false engine_kwargs: vllm_omni: {} + enable_prompt_embed_cache: false + prompt_embed_cache_size: 32 val_kwargs: _target_: verl_omni.workers.config.diffusion.DiffusionSamplingConfig 'n': 1 diff --git a/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml b/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml index 91b436d6f..1a3aa1a88 100644 --- a/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml +++ b/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml @@ -122,6 +122,12 @@ engine_kwargs: # vllm-omni engine config vllm_omni: {} +# Reuse text-encoder outputs for repeated diffusion prompts. +enable_prompt_embed_cache: false + +# Maximum number of prompt embeddings cached per diffusion worker. +prompt_embed_cache_size: 32 + # Sampling parameters used during validation (diffusion-specific). val_kwargs: diff --git a/verl_omni/workers/config/diffusion/rollout.py b/verl_omni/workers/config/diffusion/rollout.py index 830169d1f..1f2019fac 100644 --- a/verl_omni/workers/config/diffusion/rollout.py +++ b/verl_omni/workers/config/diffusion/rollout.py @@ -120,6 +120,9 @@ class DiffusionRolloutConfig(BaseConfig): # the *_stepwise variant of the pipeline (e.g. flow_grpo_stepwise). step_execution: bool = False + enable_prompt_embed_cache: bool = False + prompt_embed_cache_size: int = 32 + # note that the logprob computation should belong to the actor log_prob_micro_batch_size_per_gpu: Optional[int] = None @@ -179,6 +182,8 @@ def __post_init__(self): raise ValueError( f"Invalid diffusion rollout rollout_adapter: {self.rollout_adapter}. Must be one of ['default', 'old']." ) + if self.prompt_embed_cache_size <= 0: + raise ValueError(f"prompt_embed_cache_size must be positive, got {self.prompt_embed_cache_size}.") if self.pipeline_model_parallel_size > 1: if self.name == "vllm_omni": diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py index 4f75ad6e3..587f990a1 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py @@ -170,6 +170,9 @@ async def run_server(self, args: argparse.Namespace): engine_args["enable_dummy_pipeline"] = True engine_args["custom_pipeline_args"] = {"pipeline_class": pipeline_path} + engine_args["enable_prompt_embed_cache"] = self.config.enable_prompt_embed_cache + engine_args["prompt_embed_cache_size"] = self.config.prompt_embed_cache_size + if getattr(self.config, "step_execution", False): engine_args["step_execution"] = True @@ -360,7 +363,9 @@ def _preprocess_input( custom_prompt: OmniCustomPrompt = {"prompt_token_ids": prompt_ids} if prompt_mask is not None: - custom_prompt["prompt_mask"] = prompt_mask + custom_prompt["prompt_mask"] = ( + prompt_mask.tolist() if isinstance(prompt_mask, torch.Tensor) else prompt_mask + ) if len(default_params_list) > 1: # Multi-stage pipelines tag the diffusion stage so the orchestrator can route inputs correctly. custom_prompt["modalities"] = ["image"]