Pipelines Interface

Last updated: Oct 13, 2026 (API docstrings are auto-generated).

A pipeline in VeRL-Omni packages everything needed to plug a particular model architecture into the training loop. Two adapter families are available:

• autoregressive omni models use a training-side OmniModelBase and an optional rollout-side OmniRolloutPipelineBase;

• diffusion models use a training-side adapter subclassing DiffusionModelBase that handles scheduler setup, model-input construction, and the per-step forward / reverse-sampling logic used by RL algorithms (e.g. FlowGRPO);

• their optional rollout-side adapter is registered via VllmOmniPipelineBase that hooks into vLLM-Omni’s diffusion serving stack to expose log-probabilities.

Autoregressive training adapters are selected by (architecture, model_stage); their rollout adapters are selected by the vLLM-Omni pipeline_name. Diffusion adapters are selected by matching (DiffusionModelConfig.architecture, DiffusionModelConfig.algorithm) against a registered (architecture, algorithm) key. Diffusion architecture is read from model_index.json and the algorithm from actor_rollout_ref.model.algorithm.

verl_omni.pipelines.model_base.OmniModelBase	Abstract base class for omni model training adapters.
verl_omni.pipelines.model_base.OmniRolloutPipelineBase	Registry for omni model vLLM-Omni pipeline topologies.
verl_omni.pipelines.model_base.DiffusionModelBase	Abstract base class for diffusion model training helpers.
verl_omni.pipelines.model_base.VllmOmniPipelineBase	Registry base for vllm-omni custom diffusion pipeline classes.
verl_omni.pipelines.qwen_image_flow_grpo.QwenImage	Training adapter for the Qwen-Image diffusion model.
verl_omni.pipelines.schedulers.flow_match_sde.FlowMatchSDEDiscreteScheduler	SDE version of the FlowMatchEulerDiscreteScheduler.

Model Base

class verl_omni.pipelines.model_base.OmniModelBase

Abstract base class for omni model training adapters.

Different omni models (Qwen3-Omni, future models) have multi-stage architectures with thinker, talker, and codec components. Subclass this ABC and implement the abstract methods to plug your model into the verl RL training loop.

Unlike diffusion models, omni models are AR language models — the adapter is algorithm-agnostic. RL algorithm selection (GSPO, GRPO, RLOO, etc.) is handled by verl’s existing config fields actor.policy_loss.loss_mode and algorithm.adv_estimator.

To register, decorate your subclass with:

@OmniModelBase.register("Qwen3OmniMoeForConditionalGeneration", stage="thinker")
class Qwen3OmniThinkerAdapter(OmniModelBase):
    ...

The registry key is (architecture, stage) where architecture matches the HF config architectures[0] and stage is thinker, talker, or all.

classmethod configure_model(module, model_config)

Configure the model after loading and before FSDP wrapping.

Default implementation strips the submodules returned by get_strip_modules. Override to also:

• Redirect forward() and embedding accessors to the trainable sub-component.

• Force tie_word_embeddings=False for FSDP compatibility.

• Unfuse MoE experts for PEFT / LoRA.

Parameters:

• module – The loaded model (before FSDP wrapping).

• model_config – The OmniModelConfig.

Returns:

The configured module.

abstractmethod classmethod configure_processor(model_path: str, model_config) → Any

Load and configure the multimodal processor.

Returns a processor object that handles text, image, audio, and video inputs. Must provide apply_chat_template, __call__, tokenizer, and chat_template.

Called by the omni trainer at init time instead of verl’s default hf_processor helper.

Parameters:

• model_path – Local path to the model checkpoint.

• model_config – The OmniModelConfig.

Returns:

The configured processor (model-specific type).

abstractmethod classmethod configure_tokenizer(model_path: str, model_config) → Any

Load and configure the tokenizer.

Handles model-specific setup such as loading chat_template from a separate JSON file.

Called by the omni trainer at init time instead of verl’s default hf_tokenizer helper.

Parameters:

• model_path – Local path to the model checkpoint.

• model_config – The OmniModelConfig.

Returns:

The configured tokenizer (model-specific type).

classmethod get_class(model_config) → type[OmniModelBase]

Return the registered subclass for (architecture, model_stage).

Parameters:

model_config – An OmniModelConfig instance (or any config object with architecture and model_stage attributes).

Returns:

The registered adapter class.

Return type:

type[OmniModelBase]

Raises:

NotImplementedError – If no adapter is registered for the given (architecture, stage) key.

classmethod get_class_by_name(architecture: str, stage: str, external_lib: str | None = None) → type[OmniModelBase]

Return the registered subclass for (architecture, stage).

Parameters:

• architecture – HF config architectures[0] value.

• stage – thinker, talker, or all.

• external_lib – Optional external library to import before lookup.

Returns:

The registered adapter class.

Return type:

type[OmniModelBase]

Raises:

NotImplementedError – If no adapter is registered for the given key.

classmethod get_fsdp_ignored_module_names(model_config) → list[str]

Submodule name components to leave unsharded under FSDP2 (default: none).

FSDP2 only — under strategy=fsdp the engine raises when a non-empty list is declared. Declare a subtree whose forward is skipped for some micro-batches: unsharded forwards emit no collectives, so the skip cannot desync ranks. Ignored parameters must stay frozen — FSDP2 does not synchronize their gradients.

classmethod get_megatron_forward() → Callable

Return the pipeline’s model-forward callable for the BSHD LM engine.

The callable consumes the prepared token and multimodal inputs plus verl’s logits-processing arguments, and returns BSHD-postprocessed output. Import optional Megatron dependencies inside the override.

abstractmethod classmethod get_strip_modules(model_config) → list[str]

Return submodule prefixes to strip before FSDP init.

Multi-stage omni models contain components that are not trained in every run (e.g. the talker and codec are dead weight during thinker-only training). Stripping them before FSDP wrapping saves memory and avoids sharding unused parameters.

Parameters:

model_config – The OmniModelConfig.

Returns:

Submodule attribute names to delete. For thinker-only training this is typically ["talker", "code2wav", "code_predictor"]; for all-stage training an empty list.

Return type:

list[str]

classmethod prepare_megatron_config(model_config, engine_config)

Validate Megatron support and return the backend’s model-config view.

Called before the Megatron engine initializes. Implementations must reject unsupported configurations and preserve the original config used by the processor and rollout. FSDP adapters need not implement it.

classmethod prepare_model_inputs(model_inputs: dict[str, Any], micro_batch, model_config) → dict[str, Any]

Add model-native rollout data to an actor replay forward call.

model_inputs contains the standard language-model inputs prepared by verl. A trainable Talker stage may also need trajectory or conditioning data retained under a model-defined key in the per-sample rollout extra_fields. AgentLoopWorker batches each such key into the top level of micro_batch. A model adapter can validate its own namespaced payload and return the exact inputs required to replay the sampled policy sequence.

The default keeps the standard autoregressive path unchanged. An adapter that overrides this hook must fail closed when required fields or shapes are missing instead of reconstructing a different trajectory.

classmethod register(architecture: str, stage: str = 'thinker')

Class decorator that registers a subclass for (architecture, stage).

classmethod register_auto_classes() → None

Register optional model-package classes with Transformers auto APIs.

class verl_omni.pipelines.model_base.OmniRolloutPipelineBase

Registry for omni model vLLM-Omni pipeline topologies.

Each registered entry provides model-specific topology defaults for running the model as a multi-stage pipeline in vLLM-Omni.

To register, decorate your subclass with:

@OmniRolloutPipelineBase.register("qwen3_omni_moe")
class Qwen3OmniRolloutAdapter(OmniRolloutPipelineBase):
    ...

Registration uses a model_type key matching vLLM-Omni’s pipeline registry names (e.g. qwen3_omni_moe).

abstractmethod classmethod build_stage_configs(pipeline_mode: str = 'thinker_only') → list

Return per-stage pipeline topology for vLLM-Omni.

Each adapter defines its own pipeline_mode vocabulary (e.g. thinker_only / full for omni models, ar_only / dit_only for diffusion hybrids).

Parameters:

pipeline_mode – Model-specific mode selector.

Returns:

One frozen topology object per pipeline stage.

Return type:

list

classmethod combine_engine_outputs(outputs: list, prompt: dict) → tuple[Any, dict[str, Any]]

Select the policy output and collect architecture-specific fields.

Overriding this hook opts an adapter into retaining and assembling multiple final stage outputs. The default preserves single-output AR behavior.

classmethod ensure_pipeline_registered(pipeline_mode: str = 'thinker_only') → None

Ensure the pipeline for pipeline_mode is in vLLM-Omni’s registry.

Called before the deploy-config YAML is consumed so the engine can look up the correct PipelineConfig by its model_type.

classmethod get_class(model_type: str) → type[OmniRolloutPipelineBase] | None

Return the registered rollout adapter for model_type, or None.

Returns None when the model type is not registered — unlike OmniModelBase.get_class, which raises on miss. Rollout adapters are optional: a model may use default pipeline topology or external runner configuration.

classmethod get_engine_hf_overrides(pipeline_mode: str = 'thinker_only') → dict

Return HF config overrides per pipeline_mode.

Parameters:

pipeline_mode – The mode used to build the stages.

Returns:

HF config key-value overrides.

Return type:

dict

classmethod get_pipeline_id(pipeline_mode: str = 'thinker_only') → str

Return the vLLM-Omni pipeline model_type for pipeline_mode.

The returned string is used as the pipeline field in the generated deploy-config YAML so that the engine resolves the same PipelineConfig that the stages returned by build_stage_configs() belong to.

Parameters:

pipeline_mode – The mode used to build the stages.

Returns:

vLLM-Omni pipeline model_type.

Return type:

str

classmethod get_stage_engine_extras(stage_id: int, pipeline_mode: str = 'thinker_only') → dict

Return per-stage engine_extras to inject into the deploy config.

Parameters:

• stage_id – The pipeline stage ID.

• pipeline_mode – The mode used to build the stages.

Returns:

Extra key-value pairs merged into the stage’s engine args.

Return type:

dict

classmethod policy_stage_id(pipeline_mode='thinker_only') → int

Return the stage whose sampling parameters and logprobs define the policy.

classmethod postprocess_agent_loop_output(output, *, tokenizer, response_length: int)

Map model-native rollout data to the policy sequence used by RL.

Adapters for an omni model’s autoregressive Talker stage should put the sampled policy tokens in response_ids and align response_mask and optional response_logprobs one-to-one. Architecture-specific trajectory and conditioning data stays under a model-defined, namespaced key in extra_fields. AgentLoopWorker later exposes that key at the top level of the actor micro-batch for OmniModelBase.prepare_model_inputs().

The default preserves the standard text-token output unchanged.

classmethod prepare_engine_prompt(prompt_ids: list[int], model_config, multi_modal_data: dict, mm_processor_kwargs: dict | None = None) → dict | None

Build an architecture-specific rollout prompt when required.

classmethod register(model_type: str)

Class decorator that registers a rollout adapter for model_type.

classmethod rollout_flags(pipeline_mode='thinker_only') → dict[int, dict]

Return per-stage rollout flags for pipeline_mode.

Returns a dict[int, dict] mapping stage IDs to flags the rollout engine should apply (e.g. return_hidden_states, final_output). Default returns an empty dict — models that don’t need rollout-specific flags get this for free.

Subclasses override to add model-specific flags like return_hidden_states on intermediate AR stages in omni pipelines.

Parameters:

pipeline_mode – The mode used to build the stages.

Returns:

Per-stage flags (empty dict by default).

Return type:

dict[int, dict]

classmethod weight_sync_stage_ids(pipeline_mode='thinker_only') → list[int] | None

Return stages that receive actor weights, or all stages by default.

class verl_omni.pipelines.model_base.DiffusionModelBase

Abstract base class for diffusion model training helpers.

Different diffusion models have very different forward / sampling logic. Subclass this ABC and implement the three abstract methods to plug your model into the verl training loop.

To register, decorate your subclass with @DiffusionModelBase.register("name", algorithm="..."). The name must match the _class_name value in the pipeline’s model_index.json (which is auto-detected into DiffusionModelConfig.architecture). The algorithm must match DiffusionModelConfig.algorithm.

Example:

@DiffusionModelBase.register("QwenImagePipeline", algorithm="flow_grpo")
class QwenImage(DiffusionModelBase):
    ...

abstractmethod classmethod build_scheduler(model_config: DiffusionModelConfig) → SchedulerMixin

Build and configure the diffusion scheduler for this model. The returned scheduler should have timesteps and sigmas already set.

Parameters:

model_config (DiffusionModelConfig) – the configuration of the diffusion model.

abstractmethod classmethod forward_and_sample_previous_step(module: ModelMixin, scheduler: SchedulerMixin, model_config: DiffusionModelConfig, model_inputs: dict[str, Tensor], negative_model_inputs: dict[str, Tensor] | None, scheduler_inputs: TensorDict | dict[str, Tensor] | None, step: int)

Forward the model and sample the previous step. Used for RL-algorithms based on reversed-sampling (FlowGRPO, DanceGRPO, etc.).

Parameters:

• module (ModelMixin) – the diffusion model to be forwarded.

• scheduler (SchedulerMixin) – the scheduler used for the diffusion process.

• model_config (DiffusionModelConfig) – the configuration of the diffusion model.

• model_inputs (dict[str, torch.Tensor]) – the inputs to the diffusion model.

• negative_model_inputs (Optional[dict[str, torch.Tensor]]) – the negative inputs for guidance.

• scheduler_inputs (Optional[TensorDict | dict[str, torch.Tensor]]) – the extra inputs for the scheduler, which may contain the latents and timesteps.

• step (int) – the current step in the diffusion process.

Returns:

(log_prob, prev_sample_mean, std_dev_t, sqrt_dt)

Return type:

tuple

classmethod get_class(model_config: DiffusionModelConfig) → type[DiffusionModelBase]

Return the registered subclass for (architecture, algorithm).

classmethod get_fsdp_ignored_module_names(model_config) → list[str]

Submodule name components to leave unsharded under FSDP2 (default: none).

FSDP2 only — under strategy=fsdp the engine raises when a non-empty list is declared. Declare a subtree whose forward is skipped for some micro-batches: unsharded forwards emit no collectives, so the skip cannot desync ranks. Ignored parameters must stay frozen — FSDP2 does not synchronize their gradients.

abstractmethod classmethod prepare_model_inputs(module: ModelMixin, model_config: DiffusionModelConfig, latents: Tensor, timesteps: Tensor, prompt_embeds: Tensor, prompt_embeds_mask: Tensor, negative_prompt_embeds: Tensor, negative_prompt_embeds_mask: Tensor, micro_batch: TensorDict, step: int) → tuple[dict, dict | None]

Build architecture-specific inputs for a model forward. For reverse-trajectory algorithms, latents and timesteps usually contain the full rollout trajectory and step selects the current slice. For forward-process objectives, callers may pass an already selected/noised latent and timestep directly. The caller is responsible for universal pre-processing (common tensor extraction and nested-embed unpadding) before invoking this method.

Parameters:

• module (ModelMixin) – the diffusion transformer module.

• model_config (DiffusionModelConfig) – the configuration of the diffusion model.

• latents (torch.Tensor) – latent tensor from the micro-batch; either a full trajectory of shape (B, T, …) or a selected/noised latent of shape (B, …).

• timesteps (torch.Tensor) – timestep tensor from the micro-batch; either a full trajectory of shape (B, T) or a selected timestep of shape (B,).

• prompt_embeds (torch.Tensor) – dense positive prompt embeddings, shape (B, L, D).

• prompt_embeds_mask (torch.Tensor) – attention mask for prompt_embeds, shape (B, L).

• negative_prompt_embeds (torch.Tensor) – dense negative prompt embeddings, shape (B, L, D).

• negative_prompt_embeds_mask (torch.Tensor) – attention mask for negative_prompt_embeds.

• micro_batch (TensorDict) – the full micro-batch, available for architecture-specific metadata (e.g. height, width, vae_scale_factor).

• step (int) – the current denoising step index.

classmethod register(architecture: str, algorithm: str)

Class decorator that registers a subclass for (architecture, algorithm).

abstractmethod classmethod set_timesteps(scheduler: SchedulerMixin, model_config: DiffusionModelConfig, device: str)

Set timesteps and sigmas on the scheduler and move them to device.

Parameters:

• scheduler (SchedulerMixin) – the scheduler used for the diffusion process.

• model_config (DiffusionModelConfig) – the configuration of the diffusion model.

• device (str) – the device to move the timesteps and sigmas to.

classmethod validate_lora_config(model_config: DiffusionModelConfig) → None

Validate LoRA settings; default no-op. Override for rollout-sync-constrained models.

class verl_omni.pipelines.model_base.VllmOmniPipelineBase

Registry base for vllm-omni custom diffusion pipeline classes.

To register, decorate your custom pipeline class with @VllmOmniPipelineBase.register("name", algorithm="..."). The name must match the _class_name value in the pipeline’s model_index.json (which is auto-detected into DiffusionModelConfig.architecture). The algorithm must match DiffusionModelConfig.algorithm.

Example:

@VllmOmniPipelineBase.register("QwenImagePipeline", algorithm="flow_grpo")
class QwenImagePipelineWithLogProb(QwenImagePipeline):
    ...

classmethod get_class(architecture: str, algorithm: str) → type | None

Return the registered pipeline class for (architecture, algorithm), or None.

classmethod get_pipeline_path(architecture: str, algorithm: str) → str | None

Return the fully-qualified dotted import path for (architecture, algorithm), or None.

classmethod register(architecture: str, algorithm: str)

Class decorator that registers a pipeline for (architecture, algorithm).

Pipeline Helpers

Convenience wrappers that dispatch to the registered subclass for the current architecture. The Diffusers FSDP engine and the agent loop call into these helpers rather than touching the registry directly.

verl_omni.pipelines.utils.build_scheduler(model_config: DiffusionModelConfig) → SchedulerMixin

Build and configure the scheduler for the diffusion model. The returned scheduler has timesteps and sigmas already set.

Parameters:

model_config (DiffusionModelConfig) – the configuration of the diffusion model.

verl_omni.pipelines.utils.forward_and_sample_previous_step(module: ModelMixin, scheduler: SchedulerMixin, model_config: DiffusionModelConfig, model_inputs: dict, negative_model_inputs: dict | None, scheduler_inputs: TensorDict | dict[str, Tensor] | None, step: int)

Forward the model and sample previous step. This method is usually used for RL-algorithms based on reversed-sampling process. Such as FlowGRPO, DanceGRPO, etc.

Parameters:

• module (ModelMixin) – the diffusion model to be forwarded.

• scheduler (SchedulerMixin) – the scheduler used for the diffusion process.

• model_config (DiffusionModelConfig) – the configuration of the diffusion model.

• model_inputs (dict[str, torch.Tensor]) – the inputs to the diffusion model.

• negative_model_inputs (Optional[dict[str, torch.Tensor]]) – the negative inputs for guidance.

• scheduler_inputs (Optional[TensorDict | dict[str, torch.Tensor]]) – the extra inputs for the scheduler, which may contain the latents and timesteps.

• step (int) – the current step in the diffusion process.

verl_omni.pipelines.utils.prepare_model_inputs(module: ModelMixin, model_config: DiffusionModelConfig, latents: Tensor, timesteps: Tensor, prompt_embeds: Tensor, prompt_embeds_mask: Tensor | None, negative_prompt_embeds: Tensor | None, negative_prompt_embeds_mask: Tensor | None, micro_batch: TensorDict, step: int) → tuple[dict, dict | None]

Build architecture-specific model inputs for the forward pass. Dispatches to the registered DiffusionModelBase subclass for the current architecture.

Parameters:

• module (ModelMixin) – the diffusion transformer module.

• model_config (DiffusionModelConfig) – the configuration of the diffusion model.

• latents (torch.Tensor) – latent tensor from the micro-batch. This can be a full trajectory or an already selected/noised latent, depending on the algorithm.

• timesteps (torch.Tensor) – timestep tensor from the micro-batch. This can be a full trajectory or an already selected timestep, depending on the algorithm.

• prompt_embeds (torch.Tensor) – dense positive prompt embeddings, shape (B, L, D).

• prompt_embeds_mask (torch.Tensor) – attention mask for prompt_embeds, shape (B, L).

• negative_prompt_embeds (torch.Tensor) – dense negative prompt embeddings, shape (B, L, D).

• negative_prompt_embeds_mask (torch.Tensor) – attention mask for negative_prompt_embeds.

• micro_batch (TensorDict) – the full micro-batch, available for architecture-specific metadata (e.g. height, width, vae_scale_factor).

• step (int) – the current denoising step index.

Returns:

A pair of (model_inputs, negative_model_inputs) dicts ready to be unpacked into the transformer forward call. When the registered adapter is an I2I subclass, condition tensors are injected after the T2I prepare_model_inputs call.

Return type:

tuple[dict, Optional[dict]]

verl_omni.pipelines.utils.set_timesteps(scheduler: SchedulerMixin, model_config: DiffusionModelConfig)

Set correct timesteps and sigmas for diffusion model schedulers.

Parameters:

• scheduler (SchedulerMixin) – the scheduler used for the diffusion process.

• model_config (DiffusionModelConfig) – the configuration of the diffusion model.

Schedulers

class verl_omni.pipelines.schedulers.flow_match_sde.FlowMatchSDEDiscreteScheduler(num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, base_shift: float | None = 0.5, max_shift: float | None = 1.15, base_image_seq_len: int = 256, max_image_seq_len: int = 4096, invert_sigmas: bool = False, shift_terminal: float = None, use_karras_sigmas: bool = False, use_exponential_sigmas: bool = False, use_beta_sigmas: bool = False, time_shift_type: Literal['exponential', 'linear'] = 'exponential', stochastic_sampling: bool = False)

SDE version of the FlowMatchEulerDiscreteScheduler. The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470) and diffusers v0.37 branch.

sample_previous_step(sample: Tensor, model_output: Tensor, timestep: FloatTensor | None = None, generator: Generator | None = None, per_token_timesteps: Tensor | None = None, noise_level: float = 0.7, prev_sample: Tensor | None = None, sde_type: Literal['cps', 'sde', 'dance_sde'] = 'sde', return_logprobs: bool = True, return_sqrt_dt: bool = False, include_logprob_normalizer: bool = True)

Run a single SDE / CPS reverse step.

Parameters:

• sample (torch.FloatTensor) – A current instance of a sample created by the diffusion process.

• model_output (torch.FloatTensor) – The direct output from learned diffusion model.

• timestep (torch.FloatTensor, optional) – The current discrete timestep in the diffusion chain. When None, the internal step_index is used (sequential denoising loop).

• generator (torch.Generator, optional) – A random number generator.

• per_token_timesteps (torch.Tensor, optional) – The timesteps for each token in the sample. Currently not supported.

• noise_level (float, optional, defaults to 0.7) – The noise level used in the SDE.

• prev_sample (torch.FloatTensor, optional) – The sample from the previous timestep. If provided, it is used directly for log-probability computation instead of being sampled.

• sde_type (str, optional, defaults to “sde”) – The type of SDE to use. Choose between “sde”, “cps”, and “dance_sde”.

• return_logprobs (bool, optional, defaults to True) – Whether to return log probabilities of the previous sample.

• return_sqrt_dt (bool, optional, defaults to False) – Whether to additionally return sqrt(-dt) as a tensor of shape (batch_size,). Used by GRPO-Guard to compute the importance-ratio normalization (see GRPOGuardLoss).

• include_logprob_normalizer (bool, optional, defaults to True) – Whether to include Gaussian normalizer constants in log probabilities.

step(model_output: FloatTensor, timestep: float | FloatTensor, sample: FloatTensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = inf, s_noise: float = 1.0, generator: Generator | None = None, per_token_timesteps: Tensor | None = None, return_dict: bool = True, noise_level: float = 0.7, prev_sample: FloatTensor | None = None, sde_type: Literal['sde', 'cps', 'dance_sde'] = 'sde', return_logprobs: bool = True, include_logprob_normalizer: bool = True) → FlowMatchSDEDiscreteSchedulerOutput | tuple

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise).

Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py

Parameters:

• model_output (torch.FloatTensor) – The direct output from learned diffusion model.

• timestep (float) – The current discrete timestep in the diffusion chain.

• sample (torch.FloatTensor) – A current instance of a sample created by the diffusion process.

• s_churn (float)

• s_tmin (float)

• s_tmax (float)

• s_noise (float, defaults to 1.0) – Scaling factor for noise added to the sample.

• generator (torch.Generator, optional) – A random number generator.

• per_token_timesteps (torch.Tensor, optional) – The timesteps for each token in the sample.

• return_dict (bool) – Whether or not to return a [~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput] or tuple.

• noise_level (float, optional, defaults to 0.7) – The noise level used in the SDE.

• prev_sample (torch.FloatTensor, optional) – The sample from the previous timestep. If not provided, it will be sampled inside the function.

• sde_type (str, optional, defaults to “sde”) – The type of SDE to use. Choose between “sde”, “cps”, and “dance_sde”.

• return_logprobs (bool, optional, defaults to True) – Whether to return log probabilities of the previous sample.

• include_logprob_normalizer (bool, optional, defaults to True) – Whether to include Gaussian normalizer constants in log probabilities.

class verl_omni.pipelines.schedulers.flow_match_sde.FlowMatchSDEDiscreteSchedulerOutput(prev_sample: FloatTensor, log_prob: FloatTensor | None, prev_sample_mean: FloatTensor, std_dev_t: FloatTensor)

Output class for the scheduler’s step function output.

Parameters:

• prev_sample (torch.FloatTensor of shape (batch_size, sequence_length, num_channels) for images) – Computed sample (x_{t-1}) of previous timestep. prev_sample should be used as next model input in the denoising loop.

• log_prob (torch.FloatTensor of shape (batch_size,), optional) – The log probability of the previous sample.

• prev_sample_mean (torch.FloatTensor of shape (batch_size, sequence_length, num_channels) for images) – The mean of the computed sample of previous timestep.

• std_dev_t (torch.FloatTensor of shape (batch_size, 1, 1)) – The standard deviation used to compute prev_sample.

Built-in Pipelines

Qwen-Image (FlowGRPO)

class verl_omni.pipelines.qwen_image_flow_grpo.QwenImage

Training adapter for the Qwen-Image diffusion model.

Implements the DiffusionModelBase interface for the QwenImagePipeline architecture, providing scheduler configuration, model-input construction, and the forward/sampling step used during RL training (e.g. FlowGRPO).

Registered under "QwenImagePipeline" so it is automatically selected when DiffusionModelConfig.architecture matches that name.

classmethod build_scheduler(model_config: DiffusionModelConfig)

Build and configure the SDE scheduler for the Qwen-Image model.

Parameters:

model_config (DiffusionModelConfig) – Configuration for the diffusion model, used to determine the model path and timestep settings.

Returns:

Scheduler with timesteps already set

for the current device.

Return type:

FlowMatchSDEDiscreteScheduler

classmethod forward_and_sample_previous_step(module: QwenImageTransformer2DModel, scheduler: FlowMatchSDEDiscreteScheduler, model_config: DiffusionModelConfig, model_inputs: dict[str, Tensor], negative_model_inputs: dict[str, Tensor] | None, scheduler_inputs: TensorDict | dict[str, Tensor] | None, step: int)

Run the Qwen-Image transformer and sample the previous denoising step.

Used by RL algorithms (FlowGRPO) that require log-probabilities for reversed-sampling. Applies True-CFG guidance when model_config.true_cfg_scale > 1.0.

Parameters:

• module (QwenImageTransformer2DModel) – The Qwen-Image transformer module.

• scheduler (FlowMatchSDEDiscreteScheduler) – Scheduler used to sample the previous step and compute log-probabilities.

• model_config (DiffusionModelConfig) – Configuration providing true_cfg_scale, algo.noise_level, and algo.sde_type.

• model_inputs (dict[str, torch.Tensor]) – Positive-prompt inputs for the transformer forward pass.

• negative_model_inputs (Optional[dict[str, torch.Tensor]]) – Negative-prompt inputs used for True-CFG; may be None when CFG is disabled.

• scheduler_inputs (Optional[TensorDict | dict[str, torch.Tensor]]) – Must contain "all_latents" and "all_timesteps" tensors.

• step (int) – Current denoising step index.

Returns:

A 4-tuple of (log_prob, prev_sample_mean, std_dev_t, sqrt_dt).

Return type:

tuple

classmethod prepare_model_inputs(module: QwenImageTransformer2DModel, model_config: DiffusionModelConfig, latents: Tensor, timesteps: Tensor, prompt_embeds: Tensor, prompt_embeds_mask: Tensor, negative_prompt_embeds: Tensor, negative_prompt_embeds_mask: Tensor, micro_batch: TensorDict, step: int) → tuple[dict, dict]

Build Qwen-Image-specific inputs for the transformer forward pass.

Parameters:

• module (QwenImageTransformer2DModel) – The Qwen-Image transformer module.

• model_config (DiffusionModelConfig) – Configuration providing guidance scale and other model settings.

• latents (torch.Tensor) – Full latent tensor of shape (B, T, ...).

• timesteps (torch.Tensor) – Full timestep tensor of shape (B, T).

• prompt_embeds (torch.Tensor) – Positive prompt embeddings of shape (B, L, D).

• prompt_embeds_mask (torch.Tensor) – Attention mask for prompt_embeds of shape (B, L).

• negative_prompt_embeds (torch.Tensor) – Negative prompt embeddings of shape (B, L, D).

• negative_prompt_embeds_mask (torch.Tensor) – Attention mask for negative_prompt_embeds.

• micro_batch (TensorDict) – Micro-batch containing metadata such as height, width, and vae_scale_factor.

• step (int) – Current denoising step index used to slice latents and timesteps.

Returns:

A pair of (model_inputs, negative_model_inputs) dicts

ready to be unpacked into the transformer forward call.

Return type:

tuple[dict, dict]

classmethod set_timesteps(scheduler: FlowMatchSDEDiscreteScheduler, model_config: DiffusionModelConfig, device: str)

Configure timesteps and sigmas on the scheduler for Qwen-Image.

Parameters:

• scheduler (FlowMatchSDEDiscreteScheduler) – The scheduler whose timesteps and sigmas will be set.

• model_config (DiffusionModelConfig) – Configuration providing height, width, and number of inference steps.

• device (str) – The device (e.g. "cuda") to move the timesteps to.

class verl_omni.pipelines.qwen_image_flow_grpo.QwenImagePipelineWithLogProb(*args: Any, **kwargs: Any)

Rollout pipeline for Qwen-Image that captures per-step log-probabilities.

Extends QwenImagePipeline with a custom SDE-based scheduler and additional output fields required for RL training (e.g. FlowGRPO). In addition to the final generated image the pipeline returns all intermediate latents, their log-probabilities, and the corresponding timesteps.

Registered under "QwenImagePipeline" for vllm-omni rollout dispatch.

diffuse(prompt_embeds, prompt_embeds_mask, negative_prompt_embeds, negative_prompt_embeds_mask, latents, img_shapes, txt_seq_lens, negative_txt_seq_lens, timesteps, do_true_cfg, guidance, true_cfg_scale, noise_level, sde_window, sde_type, generator, logprobs)

Run the full SDE diffusion loop and collect per-step rollout data.

Iterates over all timesteps, optionally applying True-CFG guidance, and collects latents and log-probabilities within the SDE window.

Parameters:

• prompt_embeds (torch.Tensor) – Positive prompt embeddings.

• prompt_embeds_mask (torch.Tensor) – Attention mask for prompt_embeds.

• negative_prompt_embeds (torch.Tensor) – Negative prompt embeddings for CFG.

• negative_prompt_embeds_mask (torch.Tensor) – Attention mask for negative_prompt_embeds.

• latents (torch.Tensor) – Initial noisy latents.

• img_shapes (list) – Per-sample image shapes used by the transformer.

• txt_seq_lens (list[int]) – Sequence lengths for positive prompt embeddings.

• negative_txt_seq_lens (list[int]) – Sequence lengths for negative prompt embeddings.

• timesteps (torch.Tensor) – Scheduler timestep sequence.

• do_true_cfg (bool) – Whether to apply True-CFG guidance.

• guidance (torch.Tensor | None) – Guidance scale tensor, or None.

• true_cfg_scale (float) – Classifier-free guidance scale.

• noise_level (float) – SDE noise injection magnitude within the window.

• sde_window (tuple[int, int] | list[tuple[int, int]]) – Shared or per-row (start, end) window(s).

• sde_type (str) – SDE variant; one of "sde" or "cps".

• generator (torch.Generator | list[torch.Generator] | None) – Optional RNG; a list must have one entry per batch row.

• logprobs (bool) – Whether to compute and return per-step log-probabilities.

Returns:

A 4-tuple of

(latents, all_latents, all_log_probs, all_timesteps) where all_latents has shape (B, W+1, ...) (W = SDE-window length), all_log_probs has shape (B, W) or None when logprobs is False, and all_timesteps has shape (B, W).

Return type:

tuple

diffusion_io_spec = DiffusionIOSpec(primary=MediaSpec(modality='image', sample_rate=None, fps=None), auxiliary=())

Declares the primary rollout media stream so downstream consumers read the modality from the adapter instead of inferring it from tensor rank. Inherited by the dual-GRPO and mix-GRPO Qwen-Image subclasses.

encode_prompt(prompt_ids: Tensor | list[int] | list[list[int]] | None, attention_mask: Tensor | list[int] | list[bool] | list[list[int]] | list[list[bool]] | None = None, num_images_per_prompt: int = 1, prompt_embeds: Tensor | None = None, prompt_embeds_mask: Tensor | None = None, max_sequence_length: int = 1024)

Encode text prompt token IDs into dense embeddings.

Parameters:

• prompt_ids (torch.Tensor | list) – Token IDs of shape (B, L) or (L,).

• attention_mask (torch.Tensor | list, optional) – Boolean mask of shape (B, L) for prompt_ids; inferred as all-ones when None.

• num_images_per_prompt (int) – Number of images to generate per prompt; embeddings are repeated accordingly.

• prompt_embeds (torch.Tensor, optional) – Pre-computed embeddings; when provided prompt_ids is ignored.

• prompt_embeds_mask (torch.Tensor, optional) – Attention mask for pre-computed prompt_embeds.

• max_sequence_length (int) – Maximum sequence length; embeddings are truncated to this value.

Returns:

A pair of

(prompt_embeds, prompt_embeds_mask) tensors of shape (B * num_images_per_prompt, L, D) and (B * num_images_per_prompt, L) respectively.

Return type:

tuple[torch.Tensor, torch.Tensor]

forward(req: OmniDiffusionRequest | DiffusionRequestBatch, prompt_token_ids: torch.Tensor | list[int] | list[list[int]] | None = None, prompt_mask: torch.Tensor | list[int] | list[bool] | list[list[int]] | list[list[bool]] | None = None, negative_prompt_ids: torch.Tensor | list[int] | list[list[int]] | None = None, negative_prompt_mask: torch.Tensor | list[int] | list[bool] | list[list[int]] | list[list[bool]] | None = None, true_cfg_scale: float = 4.0, height: int | None = None, width: int | None = None, num_inference_steps: int = 50, sigmas: list[float] | None = None, guidance_scale: float = 1.0, num_images_per_prompt: int = 1, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds_mask: torch.Tensor | None = None, output_type: str | None = 'pil', attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: tuple[str, ...] = ('latents',), max_sequence_length: int = 512, noise_level: float = 0.7, sde_window_size: int | None = None, sde_window_range: tuple[int, int] = (0, 5), sde_type: Literal['sde', 'cps'] = 'sde', logprobs: bool = True) → DiffusionOutput | list[DiffusionOutput]

End-to-end image generation with rollout data collection.

Encodes the prompt, prepares latents, runs the SDE diffusion loop via diffuse(), and decodes the final latents through the VAE. Sampling parameters in req take precedence over the keyword arguments.

Parameters:

• req (OmniDiffusionRequest | DiffusionRequestBatch) – One rollout request or a request batch containing prompts and sampling parameters.

• prompt_token_ids (torch.Tensor | list[int] | list[list[int]], optional) – Token IDs for the positive prompt.

• prompt_mask (torch.Tensor | list, optional) – Attention mask for prompt_token_ids.

• negative_prompt_ids (torch.Tensor | list[int] | list[list[int]], optional) – Token IDs for the negative prompt used in True-CFG.

• negative_prompt_mask (torch.Tensor | list, optional) – Attention mask for negative_prompt_ids.

• true_cfg_scale (float) – Classifier-free guidance scale; CFG is disabled when <= 1.

• height (int, optional) – Output image height in pixels.

• width (int, optional) – Output image width in pixels.

• num_inference_steps (int) – Number of denoising steps.

• sigmas (list[float], optional) – Custom sigmas for the scheduler.

• guidance_scale (float) – Distilled guidance scale embedded in the transformer (guidance_embeds mode).

• num_images_per_prompt (int) – Number of images to generate per prompt.

• generator (torch.Generator | list[torch.Generator], optional) – Random generator(s) for reproducibility.

• latents (torch.Tensor, optional) – Pre-generated initial latents; sampled from a Gaussian when None.

• prompt_embeds (torch.Tensor, optional) – Pre-computed positive prompt embeddings; bypasses the text encoder.

• prompt_embeds_mask (torch.Tensor, optional) – Attention mask for pre-computed prompt_embeds.

• negative_prompt_embeds (torch.Tensor, optional) – Pre-computed negative prompt embeddings.

• negative_prompt_embeds_mask (torch.Tensor, optional) – Attention mask for negative_prompt_embeds.

• output_type (str, optional) – Format of the returned image; "latent" returns raw latents, otherwise the VAE-decoded image.

• attention_kwargs (dict, optional) – Extra keyword arguments forwarded to the attention layers.

• callback_on_step_end_tensor_inputs (tuple[str, ...]) – Names of tensors to expose in the step-end callback.

• max_sequence_length (int) – Maximum prompt embedding sequence length.

• noise_level (float) – SDE noise injection magnitude within the window.

• sde_window_size (int, optional) – Number of SDE steps; when None the full timestep range is used.

• sde_window_range (tuple[int, int]) – (start, end) range from which the SDE window start position is randomly sampled.

• sde_type (str) – SDE variant; "sde" or "cps".

• logprobs (bool) – Whether to compute per-step log-probabilities.

Returns:

Contains the decoded output

image plus native trajectory_* fields and prompt-embedding metadata used for FlowGRPO training.

Return type:

DiffusionOutput | list[DiffusionOutput]

post_decode(state: vllm_omni.diffusion.worker.utils.StepRequestState, **kwargs: Any) → vllm_omni.diffusion.data.DiffusionOutput

Decode final latents, package rollout trajectory, and move to CPU.

In step_execution mode the worker ships the returned DiffusionOutput across an inter-process MessageQueue to the vLLMOmniHttpServer actor. We must (a) move tensors to CPU so the receiving process does not initialise a stray CUDA context on GPU 0, and (b) populate native trajectory_* / metadata fields that forward() produces, so downstream consumers (vllm_omni_async_server.generate -> embeds_padding_2_no_padding) receive real tensors rather than None (which becomes a non-tensor LinkedList in the TensorDict and breaks mask.shape[0]).

prepare_encode(state: vllm_omni.diffusion.worker.utils.StepRequestState, **kwargs: Any) → vllm_omni.diffusion.worker.utils.StepRequestState

Populate state with encoded prompts, latents, timesteps, and CFG config.

Override of QwenImagePipeline.prepare_encode that accepts pre-tokenized prompt_ids (and optional prompt_mask) instead of raw text prompts, matching the input contract of QwenImagePipelineWithLogProb.

step_scheduler(state: vllm_omni.diffusion.worker.utils.StepRequestState, noise_pred: Tensor, **kwargs: Any) → None

One scheduler step that mirrors the per-iter body of diffuse().

The default QwenImagePipeline.step_scheduler calls the standard scheduler.step without SDE noise / log-prob bookkeeping, which means step_execution=True would silently drop all_latents / all_log_probs / all_timesteps (and the prompt_embeds_mask consumer downstream would then receive a None value that turns into a non-tensor LinkedList inside the training TensorDict). Override here to keep the step-mode and request-mode trajectories equivalent.

Qwen-Image (MixGRPO)

SD3 DPO