diff --git a/docs/algo/performance.md b/docs/algo/performance.md index df48f9cad..7d3fd7378 100644 --- a/docs/algo/performance.md +++ b/docs/algo/performance.md @@ -1,10 +1,61 @@ (performance)= # Performance Reference -Last updated: 09/09/2026 +Last updated: 09/10/2026 Below are reference benchmark results for VeRL-Omni training runs. +## Diffusion actor output retention + +The FSDP/FSDP2 PPO and NFT engines discard per-timestep model outputs after +backward during ordinary training. Losses and metrics are still aggregated, but +unused prediction/latent tensors are not retained and stacked across all +timesteps and micro-batches. Forward-only calls keep their complete outputs, +including reference-policy inference. + +Direct engine callers that need outputs during training can opt in using the +existing verl batch metadata: `tu.assign_non_tensor(batch, return_model_output=True)`. +The engine result retains the `model_output` key; it is an empty dictionary when +outputs are discarded. The training worker continues to return metrics only. + +This changes output lifetime, not the loss, gradient accumulation, input-trajectory +placement, or DPO's one-shot update. It does not make total device memory independent +of trajectory length; capacity and throughput gains require workload-specific measurement. + +### Opt-in timestep input staging + +For Qwen-Image FlowGRPO or DiffusionNFT training with FSDP/FSDP2 on GPU and +`ulysses_sequence_parallel_size=1`, enable: + +```bash +actor_rollout_ref.actor.enable_timestep_staging=true +``` + +This default-false option keeps the caller's trajectory on CPU, copies shared +prompt conditions once per micro-batch, and transfers only the current step's +inputs before forward/backward. PPO transfers the current/next latent pair and +matching loss fields; NFT keeps the clean latent and any shared noise on device +and transfers per-step noise when supplied. Dtypes, step order and gradient +accumulation are unchanged. Unused driver tensors are not copied to the device. +Consumed tensor inputs must be on CPU and must not require gradients. + +The validated scope is `QwenImagePipeline` with `flow_grpo` or `diffusion_nft` +on GPU, FSDP/FSDP2 and SP=1; see the +[Qwen-Image README](../../examples/flowgrpo_trainer/qwen_image/README.md#optional-timestep-input-staging). +The shared engine does not enforce a model/device allowlist; trainer config +validation rejects sequence parallelism. Inference keeps its existing input/output behavior. Training output +opt-in still works, but retaining those outputs reintroduces trajectory-length +dependent output memory. + +Direct engine callers can set `tu.assign_non_tensor(batch, enable_timestep_staging=True)` +on a training batch. The public actor worker supplies this metadata from its actor +configuration; the reference engine uses the unchanged upstream `FSDPEngineConfig`. + +Transfers are synchronous: there is no prefetch, pinned-memory pool, or overlap +guarantee. Full CPU trajectory storage is unchanged. Measure both peak allocated/ +reserved device memory and update time on the intended workload before choosing +this memory-for-transfer-cost tradeoff; staging is not a throughput guarantee. + ## DAPO Phase 1: LoRA Training on Qwen3-Omni Thinker AVQA This reference uses the {doc}`Thinker DAPO Phase-1 recipe <../examples/dapo_trainer>`: diff --git a/docs/examples/qwen_image/flowgrpo_trainer_qwen_image.md b/docs/examples/qwen_image/flowgrpo_trainer_qwen_image.md new file mode 120000 index 000000000..f0c0d819e --- /dev/null +++ b/docs/examples/qwen_image/flowgrpo_trainer_qwen_image.md @@ -0,0 +1 @@ +../../../examples/flowgrpo_trainer/qwen_image/README.md \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index f38c4f908..ec3430a28 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,6 @@ # Welcome to VeRL-Omni's documentation! -Last updated: 09/09/2026 +Last updated: 09/10/2026 [VeRL-Omni](https://github.com/verl-project/verl-omni) is a general RL training framework focused on multimodal generative models, built on top of [verl](https://github.com/verl-project/verl). It originated from the multi-modal generation RL effort in `verl`, and now has a dedicated home so it can evolve in a more focused way. @@ -86,6 +86,7 @@ examples/diffusionopd_trainer.md examples/flowgrpo_trainer_sd35_drm.md examples/bagel/flowgrpo_trainer_bagel.md examples/qwen3_tts/grpo_trainer_qwen3_tts.md +examples/qwen_image/flowgrpo_trainer_qwen_image.md examples/qwen_image_edit/flowgrpo_trainer_qwen_image_edit.md examples/ltx2/flowgrpo_trainer_ltx2.md examples/minimax_h3/diffusionnft_trainer_minimax_h3.md diff --git a/examples/diffusionnft_trainer/README.md b/examples/diffusionnft_trainer/README.md index 9afe08073..f77c68618 100644 --- a/examples/diffusionnft_trainer/README.md +++ b/examples/diffusionnft_trainer/README.md @@ -1,6 +1,6 @@ # DiffusionNFT Trainer -Last updated: 09/04/2026 +Last updated: 09/10/2026 This example shows how to post-train `Qwen-Image` with DiffusionNFT on an OCR-style image generation task using `vllm-omni` rollout and a visual generative reward model (`Qwen3-VL-8B-Instruct` in this example). @@ -21,6 +21,12 @@ For the full installation guide, see [Installation](../../docs/start/install.md) ## Installation +For optional Qwen-Image timestep input staging, use +`actor_rollout_ref.actor.enable_timestep_staging=true` and follow the +[shared staging contract](../flowgrpo_trainer/qwen_image/README.md#optional-timestep-input-staging). +This validation scope is Qwen-Image with FSDP/FSDP2 on GPU, SP=1; it does not +extend to the MiniMax H3 recipes above. + Follow the [installation guide](../../docs/start/install.md) to set up the base environment, then install the OCR reward dependency: ```bash diff --git a/examples/flowgrpo_trainer/qwen_image/README.md b/examples/flowgrpo_trainer/qwen_image/README.md new file mode 100644 index 000000000..8354c8f48 --- /dev/null +++ b/examples/flowgrpo_trainer/qwen_image/README.md @@ -0,0 +1,37 @@ +# Qwen-Image FlowGRPO + +Last updated: 09/10/2026 + +See the [FlowGRPO trainer guide](../../../docs/examples/flowgrpo_trainer.md) for installation, OCR data and +reward-model setup. + +## Optional timestep input staging + +For Qwen-Image (`QwenImagePipeline`) FlowGRPO or DiffusionNFT with FSDP/FSDP2 on +GPU and `ulysses_sequence_parallel_size=1`, the default-off actor option can +reduce the number of trajectory inputs resident on the device: + +```bash +bash examples/flowgrpo_trainer/qwen_image/run_qwen_image_ocr_lora.sh \ + actor_rollout_ref.actor.enable_timestep_staging=true +``` + +The same override works with the +[Qwen-Image DiffusionNFT recipe](../../../docs/examples/diffusionnft_trainer.md). +It keeps the caller's trajectory on CPU, copies shared prompt conditions once +per micro-batch, then transfers the current timestep's inputs. FlowGRPO stages +the current/next latent pair and matching loss fields; DiffusionNFT keeps the +clean latent and shared noise on device and stages per-step noise when present. +Consumed tensor inputs must be on CPU without gradients. + +Qwen-Image FlowGRPO/DiffusionNFT on GPU with FSDP/FSDP2 and SP=1 is the validated +scope. Other models, algorithms and backends are not covered by this feature's +validation; there is no model-name or device allowlist in the shared engine. +The trainer's config validation rejects staging with sequence parallelism. +Inference does not stage inputs. Training outputs remain discarded after +backward unless explicitly requested. + +Transfers are synchronous, without prefetch or an overlap guarantee. Full CPU +trajectory storage is unchanged. Measure both peak memory and update time: +this trades device memory for transfers, not a guaranteed throughput gain. +See [output lifetime and staging details](../../../docs/algo/performance.md#diffusion-actor-output-retention). diff --git a/tests/utils/test_config_on_cpu.py b/tests/utils/test_config_on_cpu.py index 5240e1eb5..0afad40e3 100644 --- a/tests/utils/test_config_on_cpu.py +++ b/tests/utils/test_config_on_cpu.py @@ -31,3 +31,20 @@ def test_validate_config_rejects_unknown_resume_mode(): def test_validate_config_requires_resume_path(): with pytest.raises(ValueError, match="resume_from_path"): validate_config(_config(resume_mode="resume_path")) + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("sp_size", [1, 2]) +@pytest.mark.parametrize("as_dict", [False, True]) +def test_validate_config_timestep_staging(enabled, sp_size, as_dict): + config = _config() + config.actor_rollout_ref = { + "actor": {"enable_timestep_staging": enabled, "fsdp_config": {"ulysses_sequence_parallel_size": sp_size}} + } + if as_dict: + config = OmegaConf.to_container(config) + if enabled and sp_size != 1: + with pytest.raises(ValueError, match="sequence_parallel_size=1"): + validate_config(config) + else: + validate_config(config) diff --git a/tests/workers/test_diffusers_output_lifetime_on_cpu.py b/tests/workers/test_diffusers_output_lifetime_on_cpu.py new file mode 100644 index 000000000..5553c428e --- /dev/null +++ b/tests/workers/test_diffusers_output_lifetime_on_cpu.py @@ -0,0 +1,188 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Exercise PPO/NFT output lifetime through the real engine batch entry points.""" + +import weakref +from types import SimpleNamespace + +import pytest +import torch +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu + +from verl_omni.workers.engine.fsdp import diffusers_impl + + +@pytest.fixture(params=["ppo", "nft"]) +def engine_case(request, monkeypatch): + engine_cls, timesteps_key = { + "ppo": (diffusers_impl.PPODiffusersFSDPEngine, "all_timesteps"), + "nft": (diffusers_impl.NFTDiffusersFSDPEngine, "train_timesteps"), + }[request.param] + + def make_engine(): + # FSDP initialization requires devices; the batch loop and optimizer are real. + engine = object.__new__(engine_cls) + engine.ulysses_sequence_parallel_size = 1 + engine.ulysses_device_mesh = None + engine.module = torch.nn.Linear(2, 2, bias=False, dtype=torch.float64) + with torch.no_grad(): + engine.module.weight.copy_(torch.tensor([[0.1, 0.2], [0.3, 0.4]], dtype=torch.float64)) + engine.optimizer = torch.optim.SGD(engine.module.parameters(), lr=0.01) + engine.optimizer_config = SimpleNamespace(clip_grad=1000.0) + engine.get_data_parallel_group = lambda: None + observed = SimpleNamespace(refs=[], live_before_step=[], calls=[], backward_calls=0) + + def forward_step(micro_batch, loss_function, forward_only, step): + observed.live_before_step.append(sum(ref() is not None for ref in observed.refs)) + observed.calls.append((micro_batch["sample_id"].tolist(), step, torch.is_grad_enabled())) + prediction = engine.module(micro_batch["features"] + step) + model_output = {"prediction": prediction, "auxiliary": prediction.detach().square()} + output_refs = tuple(weakref.ref(value) for value in model_output.values()) + observed.refs.extend(output_refs) + + if loss_function is None: + assert forward_only + loss = prediction.new_tensor(1.0) + metrics = {} + else: + loss, metrics = loss_function(model_output, micro_batch) + + if loss.requires_grad: + + def check_backward(gradient): + assert all(ref() is not None for ref in output_refs) + observed.backward_calls += 1 + return gradient + + loss.register_hook(check_backward) + + return loss, {"model_output": model_output, "loss": loss.detach().item(), "metrics": metrics} + + engine.forward_step = forward_step + return engine, observed + + def prepare_micro_batches(data, dp_group, same_micro_num_in_dp): + assert dp_group is None + assert same_micro_num_in_dp + assert tu.get_non_tensor_data(data, "use_dynamic_bsz", default=None) is False + assert tu.get_non_tensor_data(data, "sp_size", default=None) == 1 + return list(data.split(tu.get_non_tensor_data(data, "micro_batch_size_per_gpu", default=None))), None + + monkeypatch.setattr(diffusers_impl, "get_device_id", lambda: "cpu") + monkeypatch.setattr(diffusers_impl, "prepare_micro_batches", prepare_micro_batches) + return make_engine, timesteps_key + + +def _batch(timesteps_key, num_steps, micro_batch_size, return_model_output=None): + batch = TensorDict( + { + "features": torch.arange(12, dtype=torch.float64).reshape(6, 2) / 10, + "sample_id": torch.arange(6), + timesteps_key: torch.arange(num_steps).expand(6, -1), + }, + batch_size=[6], + ) + tu.assign_non_tensor(batch, micro_batch_size_per_gpu=micro_batch_size) + if return_model_output is not None: + tu.assign_non_tensor(batch, return_model_output=return_model_output) + return batch + + +def _loss(model_output, data): + loss = (model_output["prediction"] - 1).square().mean() + loss = loss / tu.get_non_tensor_data(data, "gradient_accumulation_steps", default=None) + return loss, {"objective": loss.detach().item()} + + +@pytest.mark.parametrize("num_steps", [1, 4]) +@pytest.mark.parametrize("micro_batch_size", [2, 6]) +@pytest.mark.parametrize("forward_only", [False, True]) +@pytest.mark.parametrize("return_model_output", [None, False, True]) +def test_output_retention_contract(engine_case, num_steps, micro_batch_size, forward_only, return_model_output): + make_engine, timesteps_key = engine_case + engine, observed = make_engine() + batch = _batch(timesteps_key, num_steps, micro_batch_size, return_model_output) + initial_weight = engine.module.weight.detach().clone() + + if forward_only: + output = engine.infer_batch(batch, loss_function=_loss) + else: + output = engine.train_batch(batch, loss_function=_loss) + + num_micro_batches = 6 // micro_batch_size + num_calls = num_micro_batches * num_steps + assert observed.calls == [ + (list(range(start, start + micro_batch_size)), step, not forward_only) + for start in range(0, 6, micro_batch_size) + for step in range(num_steps) + ] + assert observed.backward_calls == (0 if forward_only else num_calls) + assert len(output["loss"]) == num_micro_batches + assert all(len(losses) == num_steps for losses in output["loss"]) + assert output["metrics"]["objective"] == [loss for losses in output["loss"] for loss in losses] + + keep_outputs = forward_only or return_model_output is True + expected_live_counts = [2 * index if keep_outputs else 0 for index in range(num_calls)] + assert observed.live_before_step == expected_live_counts + # Postprocessing creates new stacked tensors; no per-step tensor should survive it. + assert all(ref() is None for ref in observed.refs) + + if keep_outputs: + expected = torch.cat( + [ + torch.stack( + [torch.nn.functional.linear(features + step, initial_weight) for step in range(num_steps)], + dim=1, + ) + for features in batch["features"].split(micro_batch_size) + ], + dim=0, + ) + torch.testing.assert_close(output["model_output"]["prediction"], expected, rtol=0, atol=0) + torch.testing.assert_close(output["model_output"]["auxiliary"], expected.square(), rtol=0, atol=0) + else: + assert output["model_output"] == {} + if forward_only: + assert engine.module.weight.grad is None + torch.testing.assert_close(engine.module.weight, initial_weight, rtol=0, atol=0) + + +def test_training_update_matches_retained_outputs_across_repeated_calls(engine_case): + make_engine, timesteps_key = engine_case + optimized, _ = make_engine() + reference, _ = make_engine() + + for _ in range(2): + actual = optimized.train_batch(_batch(timesteps_key, 4, 2), loss_function=_loss) + expected = reference.train_batch(_batch(timesteps_key, 4, 2, True), loss_function=_loss) + + assert actual["model_output"] == {} + assert expected["model_output"] + assert actual["loss"] == expected["loss"] + assert actual["metrics"] == expected["metrics"] + torch.testing.assert_close(optimized.module.weight.grad, reference.module.weight.grad, rtol=0, atol=0) + torch.testing.assert_close(optimized.module.weight, reference.module.weight, rtol=0, atol=0) + + +def test_inference_without_loss_preserves_outputs(engine_case): + make_engine, timesteps_key = engine_case + engine, observed = make_engine() + output = engine.infer_batch(_batch(timesteps_key, 3, 2)) + + assert output["model_output"]["prediction"].shape == (6, 3, 2) + assert output["loss"] == [[1.0] * 3] * 3 + assert output["metrics"] == {} + assert observed.backward_calls == 0 + assert engine.module.weight.grad is None diff --git a/tests/workers/test_diffusers_timestep_staging_on_cpu.py b/tests/workers/test_diffusers_timestep_staging_on_cpu.py new file mode 100644 index 000000000..9487f9bcb --- /dev/null +++ b/tests/workers/test_diffusers_timestep_staging_on_cpu.py @@ -0,0 +1,436 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Check timestep selection, transfer ownership and public engine/config wiring.""" + +import weakref +from contextlib import nullcontext +from inspect import unwrap +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config import FSDPEngineConfig + +import verl_omni +from verl_omni.pipelines.schedulers import FlowMatchSDEDiscreteScheduler +from verl_omni.utils.config import validate_config +from verl_omni.workers.config import FSDPDiffusionActorConfig +from verl_omni.workers.engine.fsdp import diffusers_impl +from verl_omni.workers.engine_workers import ActorRolloutRefWorker + +PPO_OPTIONAL = ( + "ref_log_prob", + "ref_prev_sample_mean", + "teacher_prev_sample_mean", + "old_prev_sample_mean", + "rollout_is_weights", +) + + +def _batch(algorithm, steps=4, noise_mode="per_step"): + generator = torch.Generator().manual_seed(100) + + def values(*shape): + return torch.randn(*shape, generator=generator, dtype=torch.float64) + + tensors = { + "prompt_embeds": values(4, 3, 2), + "prompt_embeds_mask": torch.ones(4, 3, dtype=torch.int64), + "negative_prompt_embeds": values(4, 3, 2), + "negative_prompt_embeds_mask": torch.ones(4, 3, dtype=torch.int64), + "unused_trajectory": values(4, steps * 2, 3, 2), + } + timesteps = torch.arange(steps, dtype=torch.float64).flip(0)[None].expand(4, -1) + if algorithm == "ppo": + tensors.update( + all_latents=values(4, steps + 1, 3, 2), + all_timesteps=timesteps, + old_log_probs=values(4, steps), + advantages=values(4, steps), + ) + for key in PPO_OPTIONAL: + tensors[key] = values(4, steps, 3, 2) if key.endswith("sample_mean") else values(4, steps) + else: + timesteps = timesteps + torch.arange(4, dtype=torch.float64)[:, None] * 100 + tensors.update(latents_clean=values(4, 3, 2), train_timesteps=timesteps, reward_prob=values(4, steps)) + if noise_mode != "generated": + tensors["forward_noise"] = values(4, steps, 3, 2) if noise_mode == "per_step" else values(4, 3, 2) + batch = TensorDict(tensors, batch_size=[4]) + tu.assign_non_tensor(batch, micro_batch_size_per_gpu=2, height=16, width=24, vae_scale_factor=8) + return batch + + +@pytest.fixture +def transfer_spy(monkeypatch): + """Model device-copy ownership with CPU clones, without claiming hardware coverage.""" + observed = SimpleNamespace(shared=[], steps=[]) + original_to = torch.Tensor.to + + def tensor_to(tensor, *args, **kwargs): + if args and args[0] == "cpu": + copied = tensor.clone() + observed.steps.append((tuple(tensor.shape), tensor.dtype, weakref.ref(copied))) + return copied + return original_to(tensor, *args, **kwargs) + + def tensordict_to(data, device, **kwargs): + assert device == "cpu" + observed.shared.append(set(data.keys())) + return data.clone() + + monkeypatch.setattr(torch.Tensor, "to", tensor_to) + monkeypatch.setattr(TensorDict, "to", tensordict_to) + monkeypatch.setattr(diffusers_impl, "get_device_id", lambda: "cpu") + return observed + + +def _engine(algorithm, staging): + cls = diffusers_impl.PPODiffusersFSDPEngine if algorithm == "ppo" else diffusers_impl.NFTDiffusersFSDPEngine + engine = object.__new__(cls) + train_batch = engine.train_batch + + def train_with_staging(data, loss_function): + tu.assign_non_tensor(data, enable_timestep_staging=staging) + return train_batch(data, loss_function) + + engine.train_batch = train_with_staging + engine.ulysses_sequence_parallel_size = 1 + engine.ulysses_device_mesh = None + engine.get_data_parallel_group = lambda: None + engine.module = torch.nn.Linear(2, 2, bias=False, dtype=torch.float64) + with torch.no_grad(): + engine.module.weight.copy_(torch.tensor([[0.1, 0.2], [0.3, 0.4]], dtype=torch.float64)) + engine.optimizer = torch.optim.SGD(engine.module.parameters(), lr=0.01) + engine.optimizer_config = SimpleNamespace(clip_grad=1000.0) + observed = SimpleNamespace(steps=[], prompts=[], prior_inputs=[], previous=[], backward=0) + + def forward_step(data, loss_function, forward_only, step): + observed.prior_inputs.append(sum(ref() is not None for ref in observed.previous)) + observed.prompts.append(id(data["prompt_embeds"])) + if algorithm == "ppo": + latent = data["all_latents"] + timestep = data["all_timesteps"][:, step] + features = latent[:, step, 0] + 0.3 * latent[:, step + 1, 0] + for key in ("old_log_probs", "advantages", *PPO_OPTIONAL): + if key in data: + features = features + data[key][:, step].reshape(len(data), -1).mean(-1, keepdim=True) + observed.previous = [weakref.ref(latent)] + else: + x0 = data["latents_clean"] + noise = data.get("forward_noise", None) + if noise is None: + noise = torch.randn_like(x0) + elif noise.ndim == x0.ndim + 1: + noise = noise[:, step] + timestep = data["train_timesteps"][:, step] + features = x0[:, 0] + noise[:, 0] + data["reward_prob"][:, step, None] + observed.previous = [weakref.ref(data["train_timesteps"])] + observed.steps.append(timestep.tolist()) + features = features + data["prompt_embeds"][:, 0] - 0.2 * data["negative_prompt_embeds"][:, 0] + prediction = engine.module(features + timestep[:, None] / 1000) + outputs = {"prediction": prediction} + loss = prediction.square().mean() / tu.get_non_tensor_data(data, "gradient_accumulation_steps", default=1) + if loss.requires_grad: + + def count_backward(gradient): + assert all(ref() is not None for ref in observed.previous) + observed.backward += 1 + return gradient + + loss.register_hook(count_backward) + return loss, {"model_output": outputs, "loss": loss.detach().item(), "metrics": {}} + + engine.forward_step = forward_step + return engine, observed + + +@pytest.mark.parametrize( + "algorithm,noise_mode", [("ppo", "per_step"), ("nft", "per_step"), ("nft", "static"), ("nft", "generated")] +) +@pytest.mark.parametrize("steps", [1, 4, 12]) +@pytest.mark.parametrize("retain", [False, True]) +def test_staging_matches_updates_and_preserves_inputs(transfer_spy, algorithm, noise_mode, steps, retain): + candidate, observed = _engine(algorithm, staging=True) + reference, _ = _engine(algorithm, staging=False) + source = _batch(algorithm, steps, noise_mode) + frozen = source.clone() + for _ in range(2): + batch = source.clone() + tu.assign_non_tensor(batch, return_model_output=retain) + torch.manual_seed(99) + actual = candidate.train_batch(batch, loss_function=lambda **kwargs: None) + torch.manual_seed(99) + expected = reference.train_batch(batch.clone(), loss_function=lambda **kwargs: None) + assert actual["loss"] == expected["loss"] + assert actual["metrics"] == expected["metrics"] + assert bool(actual["model_output"]) is retain + if retain: + torch.testing.assert_close( + actual["model_output"]["prediction"], expected["model_output"]["prediction"], rtol=0, atol=0 + ) + torch.testing.assert_close(candidate.module.weight, reference.module.weight, rtol=0, atol=0) + torch.testing.assert_close(candidate.module.weight.grad, reference.module.weight.grad, rtol=0, atol=0) + for key, value in frozen.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(batch[key], value, rtol=0, atol=0) + assert observed.backward == 4 * steps + expected_steps = [ + [float(step + (row * 100 if algorithm == "nft" else 0)) for row in range(start, start + 2)] + for _ in range(2) + for start in (0, 2) + for step in reversed(range(steps)) + ] + assert observed.steps == expected_steps + assert observed.prior_inputs == [0] * (4 * steps) + assert all( + observed.prompts[start : start + steps] == [observed.prompts[start]] * steps + for start in range(0, 4 * steps, steps) + ) + staged_copies = [keys for keys in transfer_spy.shared if "unused_trajectory" not in keys] + assert len(staged_copies) == 4 + assert all("prompt_embeds" in keys for keys in staged_copies) + assert all("all_latents" not in keys and "train_timesteps" not in keys for keys in staged_copies) + assert all(ref() is None for _, _, ref in transfer_spy.steps) + assert all(shape[1] <= 2 for shape, _, _ in transfer_spy.steps) + + +@pytest.mark.parametrize("algorithm", ["ppo", "nft"]) +def test_inference_bypasses_staging(transfer_spy, algorithm): + candidate, observed = _engine(algorithm, staging=True) + reference, _ = _engine(algorithm, staging=False) + batch = _batch(algorithm) + tu.assign_non_tensor(batch, return_model_output=False, enable_timestep_staging=True) + actual = candidate.infer_batch(batch) + expected = reference.infer_batch(batch.clone()) + torch.testing.assert_close( + actual["model_output"]["prediction"], expected["model_output"]["prediction"], rtol=0, atol=0 + ) + assert observed.backward == 0 + assert transfer_spy.steps == [] + assert all("unused_trajectory" in keys for keys in transfer_spy.shared) + + +@pytest.mark.parametrize("algorithm", ["ppo", "nft"]) +def test_staging_recreated_after_failed_forward(transfer_spy, algorithm): + engine, observed = _engine(algorithm, staging=True) + original_forward = engine.forward_step + + def failing_forward(*args, **kwargs): + raise RuntimeError("injected forward failure") + + engine.forward_step = failing_forward + with pytest.raises(RuntimeError, match="injected forward failure"): + engine.train_batch(_batch(algorithm), loss_function=lambda **kwargs: None) + assert all(ref() is None for _, _, ref in transfer_spy.steps) + engine.forward_step = original_forward + result = engine.train_batch(_batch(algorithm), loss_function=lambda **kwargs: None) + assert result["model_output"] == {} + assert observed.backward == 8 + + +def test_ppo_optional_fields_can_be_absent(transfer_spy): + engine, observed = _engine("ppo", staging=True) + batch = _batch("ppo") + for key in PPO_OPTIONAL: + del batch[key] + assert engine.train_batch(batch, loss_function=lambda **kwargs: None)["model_output"] == {} + assert observed.backward == 8 + + +@pytest.mark.parametrize( + "case", ["empty", "latent_length", "loss_length", "noise_shape", "missing_mask", "input_grad", "device"] +) +def test_invalid_staging_input_fails_before_forward(transfer_spy, case): + algorithm = "nft" if case in {"noise_shape", "missing_mask"} else "ppo" + engine, observed = _engine(algorithm, staging=True) + batch = _batch(algorithm) + if case == "empty": + batch["all_timesteps"] = torch.empty(4, 0) + elif case == "latent_length": + batch["all_latents"] = batch["all_latents"][:, :3] + elif case == "loss_length": + batch["old_log_probs"] = batch["old_log_probs"][:, :3] + elif case == "noise_shape": + batch["forward_noise"] = batch["forward_noise"][:, :3] + elif case == "missing_mask": + del batch["prompt_embeds_mask"] + elif case == "input_grad": + batch["all_latents"].requires_grad_() + else: + batch["all_latents"] = torch.empty(4, 5, 3, 2, device="meta") + with pytest.raises(ValueError): + engine.train_batch(batch, loss_function=lambda **kwargs: None) + assert observed.steps == [] + assert observed.backward == 0 + assert not transfer_spy.shared and not transfer_spy.steps + + +@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2"]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_hydra_actor_forwards_timestep_staging(strategy, enabled): + config_dir = Path(verl_omni.__file__).parent / "trainer/config/diffusion/actor" + with initialize_config_dir(config_dir=str(config_dir), version_base=None): + cfg = compose( + config_name="dp_diffusion_actor", + overrides=[ + f"strategy={strategy}", + "ppo_micro_batch_size_per_gpu=2", + f"enable_timestep_staging={str(enabled).lower()}", + ], + ) + actor = omega_conf_to_dataclass(cfg) + assert isinstance(actor, FSDPDiffusionActorConfig) + assert type(actor.engine) is FSDPEngineConfig + assert actor.engine is actor.fsdp_config + assert actor.enable_timestep_staging is enabled + assert not hasattr(actor.engine, "enable_timestep_staging") + assert actor.engine.strategy == strategy + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_public_trainer_override_reaches_actor_worker(enabled): + config_dir = Path(verl_omni.__file__).parent / "trainer/config" + overrides = ["actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2"] + if enabled: + overrides.append("actor_rollout_ref.actor.enable_timestep_staging=true") + with initialize_config_dir(config_dir=str(config_dir), version_base=None): + cfg = compose(config_name="diffusion_trainer", overrides=overrides) + validate_config(cfg) + actor = omega_conf_to_dataclass(cfg.actor_rollout_ref.actor) + assert actor.enable_timestep_staging is enabled + assert type(actor.engine) is FSDPEngineConfig + ref = omega_conf_to_dataclass(cfg.actor_rollout_ref.ref) + assert not ref.enable_timestep_staging + assert type(ref.engine) is FSDPEngineConfig + + received = [] + + def train_mini_batch(data): + received.append(tu.get_non_tensor_data(data, "enable_timestep_staging", default=None)) + return None + + worker = SimpleNamespace(config=cfg.actor_rollout_ref, actor=SimpleNamespace(train_mini_batch=train_mini_batch)) + batch = TensorDict({}, batch_size=[2]) + tu.assign_non_tensor(batch, enable_timestep_staging=not enabled) + unwrap(ActorRolloutRefWorker.update_actor)(worker, batch) + assert received == [enabled] + + +@pytest.mark.parametrize("entrypoint", ["main_diffusion", "main_diffusion_v1"]) +def test_public_config_rejects_staging_with_sequence_parallel(monkeypatch, entrypoint): + import importlib + + config_dir = Path(verl_omni.__file__).parent / "trainer/config" + with initialize_config_dir(config_dir=str(config_dir), version_base=None): + cfg = compose( + config_name="diffusion_trainer", + overrides=[ + "actor_rollout_ref.actor.enable_timestep_staging=true", + "actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=2", + ], + ) + with pytest.raises(ValueError, match="sequence_parallel_size=1"): + validate_config(cfg) + module = importlib.import_module(f"verl_omni.trainer.{entrypoint}") + monkeypatch.setattr(module, "auto_set_device", lambda config: None) + with pytest.raises(ValueError, match="sequence_parallel_size=1"): + module.main.__wrapped__(cfg) + + +def test_worker_without_staging_config_resets_batch_flag(): + received = [] + worker = SimpleNamespace( + config=OmegaConf.create({"actor": {}}), + actor=SimpleNamespace( + train_mini_batch=lambda data: received.append( + tu.get_non_tensor_data(data, "enable_timestep_staging", default=None) + ) + ), + ) + batch = TensorDict({}, batch_size=[1]) + tu.assign_non_tensor(batch, enable_timestep_staging=True) + unwrap(ActorRolloutRefWorker.update_actor)(worker, batch) + assert received == [False] + + +@pytest.mark.parametrize( + "algorithm,noise_mode", [("ppo", "per_step"), ("nft", "per_step"), ("nft", "static"), ("nft", "generated")] +) +def test_real_qwen_adapter_and_scheduler_contract(transfer_spy, algorithm, noise_mode): + """Exercise real forward_step/adapters with a tiny CPU projection, not a transformer or FSDP.""" + engines = [] + for staging in (False, True): + engine, _ = _engine(algorithm, staging) + del engine.forward_step # Use the real algorithm-specific class method. + engine.use_ulysses_sp = False + engine.model_config = SimpleNamespace( + architecture="QwenImagePipeline", + algorithm="flow_grpo" if algorithm == "ppo" else "diffusion_nft", + external_lib=None, + pipeline=SimpleNamespace(guidance_scale=1.0, true_cfg_scale=2.0), + algo=SimpleNamespace(noise_level=0.8, sde_type="sde"), + ) + engine.module.config = SimpleNamespace(guidance_embeds=False) + projection = engine.module.forward + + def project(hidden_states, encoder_hidden_states, _projection=projection, **kwargs): + return (_projection(hidden_states) + encoder_hidden_states[:, :1] * 0.1,) + + engine.module.forward = project + engine.scheduler = FlowMatchSDEDiscreteScheduler() + engine.scheduler.set_timesteps(5, device="cpu") + engine.use_adapter = lambda name: nullcontext() + engine.disable_adapter = nullcontext + engine._set_adapter = lambda name: None + engines.append(engine) + + batch = _batch(algorithm, noise_mode=noise_mode) + if algorithm == "ppo": + batch["all_timesteps"] = engines[0].scheduler.timesteps[:4].expand(4, -1).clone() + tu.assign_non_tensor(batch, return_model_output=True) + seen_loss_data = [] + + def loss_function(model_output, data, dp_group): + seen_loss_data.append(data.clone()) + prediction = model_output["prev_sample_mean" if algorithm == "ppo" else "forward_prediction"] + weight = sum(value.float().mean() for value in data.values() if isinstance(value, torch.Tensor)) + loss = prediction.square().mean() * (1 + weight.square()) + return loss / tu.get_non_tensor_data(data, "gradient_accumulation_steps", default=None), {} + + results = [] + for engine in engines: + torch.manual_seed(99) + results.append(engine.train_batch(batch.clone(), loss_function)) + reference, candidate = results + torch.testing.assert_close(torch.tensor(candidate["loss"]), torch.tensor(reference["loss"]), rtol=0, atol=0) + for key in reference["model_output"]: + torch.testing.assert_close(candidate["model_output"][key], reference["model_output"][key], rtol=0, atol=0) + torch.testing.assert_close(engines[1].module.weight, engines[0].module.weight, rtol=0, atol=0) + torch.testing.assert_close(engines[1].module.weight.grad, engines[0].module.weight.grad, rtol=0, atol=0) + assert len(seen_loss_data) == 16 + for expected, actual in zip(seen_loss_data[:8], seen_loss_data[8:], strict=True): + assert set(expected.keys()) == set(actual.keys()) + for key, value in expected.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(actual[key], value, rtol=0, atol=0) + else: + assert tu.get_non_tensor_data(actual, key, default=None) == tu.get_non_tensor_data( + expected, key, default=None + ) diff --git a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml index 85ca277be..3aea58775 100644 --- a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml @@ -106,6 +106,7 @@ actor_rollout_ref: - extra load_contents: ${.save_contents} async_save: false + enable_timestep_staging: false grad_clip: 1.0 ref: rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1} diff --git a/verl_omni/trainer/config/diffusion/actor/dp_diffusion_actor.yaml b/verl_omni/trainer/config/diffusion/actor/dp_diffusion_actor.yaml index d24b37359..159ed69d2 100644 --- a/verl_omni/trainer/config/diffusion/actor/dp_diffusion_actor.yaml +++ b/verl_omni/trainer/config/diffusion/actor/dp_diffusion_actor.yaml @@ -28,5 +28,8 @@ _target_: verl_omni.workers.config.diffusion.FSDPDiffusionActorConfig # Training strategy: fsdp, fsdp2 strategy: fsdp +# Stage training inputs from CPU per timestep; see the Qwen-Image recipe README. +enable_timestep_staging: false + # Gradient clipping for actor updates, specific to the strategy. grad_clip: 1.0 diff --git a/verl_omni/utils/config.py b/verl_omni/utils/config.py index 90e53c1e6..1de4a60a9 100644 --- a/verl_omni/utils/config.py +++ b/verl_omni/utils/config.py @@ -23,6 +23,11 @@ def _select(config: Any, path: str, default: Any = None) -> Any: def validate_config(config: Any) -> None: """Validate configuration values that otherwise trigger silent fallbacks.""" + if _select(config, "actor_rollout_ref.actor.enable_timestep_staging", False): + sp_size = _select(config, "actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size", 1) + if sp_size != 1: + raise ValueError("Timestep staging requires ulysses_sequence_parallel_size=1.") + resume_mode = _select(config, "trainer.resume_mode") valid_resume_modes = ("disable", "auto", "resume_path") if resume_mode not in valid_resume_modes: diff --git a/verl_omni/workers/config/diffusion/actor.py b/verl_omni/workers/config/diffusion/actor.py index 2cd1261af..bffc6ee46 100644 --- a/verl_omni/workers/config/diffusion/actor.py +++ b/verl_omni/workers/config/diffusion/actor.py @@ -186,6 +186,8 @@ class FSDPDiffusionActorConfig(DiffusionActorConfig): strategy: str = "fsdp" grad_clip: float = 1.0 fsdp_config: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) + # Stage training inputs from CPU one timestep at a time. + enable_timestep_staging: bool = False def __post_init__(self): """Validate diffusion FSDP actor configuration parameters.""" diff --git a/verl_omni/workers/engine/fsdp/diffusers_impl.py b/verl_omni/workers/engine/fsdp/diffusers_impl.py index 928141b8c..cc2d8345a 100644 --- a/verl_omni/workers/engine/fsdp/diffusers_impl.py +++ b/verl_omni/workers/engine/fsdp/diffusers_impl.py @@ -866,6 +866,72 @@ def get_per_tensor_param( peft_config_dict = peft_config.to_dict() if peft_config is not None else None return per_tensor_param, peft_config_dict + def _prepare_timestep_staging(self, data: TensorDict, timesteps_key: str) -> tuple[dict[str, int], list[str]]: + """Validate CPU inputs and describe the Qwen-Image fields consumed by each step.""" + timesteps = data.get(timesteps_key, None) + if not isinstance(timesteps, torch.Tensor) or timesteps.ndim != 2 or 0 in timesteps.shape: + raise ValueError("Timestep staging requires a nonempty (batch, timesteps) tensor.") + num_steps = timesteps.shape[1] + shared_keys = [ + "prompt_embeds", + "prompt_embeds_mask", + "negative_prompt_embeds", + "negative_prompt_embeds_mask", + "height", + "width", + "vae_scale_factor", + "gradient_accumulation_steps", + "sp_size", + ] + if timesteps_key == "all_timesteps": + latents = data.get("all_latents", None) + if not isinstance(latents, torch.Tensor) or latents.ndim != 4 or latents.shape[1] != num_steps + 1: + raise ValueError("Staged PPO all_latents must have shape (batch, timesteps + 1, tokens, channels).") + required = ("all_timesteps", "all_latents", "old_log_probs", "advantages", "prompt_embeds") + step_fields = dict.fromkeys( + ( + "all_timesteps", + "old_log_probs", + "advantages", + "ref_log_prob", + "ref_prev_sample_mean", + "teacher_prev_sample_mean", + "old_prev_sample_mean", + "rollout_is_weights", + ), + 1, + ) + step_fields["all_latents"] = 2 + else: + latents = data.get("latents_clean", None) + if not isinstance(latents, torch.Tensor) or latents.ndim != 3: + raise ValueError("Staged NFT latents_clean must have shape (batch, tokens, channels).") + required = ("train_timesteps", "reward_prob", "latents_clean", "prompt_embeds", "prompt_embeds_mask") + step_fields = {"train_timesteps": 1, "reward_prob": 1} + shared_keys.append("latents_clean") + noise = data.get("forward_noise", None) + if noise is not None: + if not isinstance(noise, torch.Tensor): + raise ValueError("Staged NFT forward_noise must be a tensor.") + if noise.shape == latents.shape: + shared_keys.append("forward_noise") + elif noise.shape == (latents.shape[0], num_steps, *latents.shape[1:]): + step_fields["forward_noise"] = 1 + else: + raise ValueError("Staged NFT forward_noise must match latents_clean or add a timestep dimension.") + for key in required: + if not isinstance(data.get(key, None), torch.Tensor): + raise ValueError(f"Timestep staging requires tensor input {key!r}.") + step_fields = {key: width for key, width in step_fields.items() if data.get(key, None) is not None} + for key, width in step_fields.items(): + value = data[key] + if not isinstance(value, torch.Tensor) or value.ndim < 2 or value.shape[1] != num_steps + width - 1: + raise ValueError(f"Staged input {key!r} has an incompatible timestep dimension.") + for key, value in data.select(*shared_keys, *step_fields, strict=False).items(): + if isinstance(value, torch.Tensor) and (value.device.type != "cpu" or value.requires_grad): + raise ValueError(f"Timestep staging requires CPU inputs without gradients, got {key!r}.") + return step_fields, shared_keys + def _merged_lora_per_tensor_param(self): """Stream merged (base + LoRA) weights for rollout weight sync. @@ -907,7 +973,11 @@ def _run_forward_backward_batch( *, timesteps_key: str, ) -> dict: + stage_inputs = tu.get_non_tensor_data(data, "enable_timestep_staging", default=False) and not forward_only + if stage_inputs: + step_fields, shared_keys = self._prepare_timestep_staging(data, timesteps_key) num_timesteps = int(data[timesteps_key].shape[1]) + return_model_output = tu.get_non_tensor_data(data, "return_model_output", default=False) tu.assign_non_tensor(data, sp_size=self.ulysses_sequence_parallel_size) tu.assign_non_tensor(data, use_dynamic_bsz=False) @@ -920,20 +990,38 @@ def _run_forward_backward_batch( ctx = torch.no_grad() if forward_only else nullcontext() for micro_batch in micro_batches: - micro_batch = micro_batch.to(get_device_id()) tu.assign_non_tensor(micro_batch, gradient_accumulation_steps=gradient_accumulation_steps) + if stage_inputs: + shared_batch = micro_batch.select(*shared_keys, strict=False).to(get_device_id()) + else: + micro_batch = micro_batch.to(get_device_id()) meta_info_lst = {"model_output": [], "loss": [], "metrics": []} # Forward and backward for each timestep with ctx: for step in range(num_timesteps): + if stage_inputs: + step_batch = shared_batch.clone(recurse=False) + for key, width in step_fields.items(): + step_batch[key] = micro_batch[key][:, step : step + width].to(get_device_id()) + else: + step_batch = micro_batch loss, meta_info = self.forward_step( - micro_batch, loss_function=loss_function, forward_only=forward_only, step=step + step_batch, + loss_function=loss_function, + forward_only=forward_only, + step=0 if stage_inputs else step, ) if not forward_only: loss.backward() + if not return_model_output: + # Training consumers only need metrics; do not retain every timestep's latents. + meta_info.pop("model_output", None) for key, val in meta_info.items(): meta_info_lst[key].append(val) + del step_batch output_lst.append(meta_info_lst) + if stage_inputs: + del shared_batch # postprocess and return return self.postprocess_batch_func(output_lst=output_lst, indices=indices, data=data) diff --git a/verl_omni/workers/engine_workers.py b/verl_omni/workers/engine_workers.py index ecf9139d5..f4b0c766c 100644 --- a/verl_omni/workers/engine_workers.py +++ b/verl_omni/workers/engine_workers.py @@ -895,6 +895,7 @@ def compute_log_prob(self, data: TensorDict) -> TensorDict: @DistProfiler.annotate(color="red", role="actor_update") @_with_routing_replay_flag(enabled=True) def update_actor(self, data: TensorDict) -> TensorDict: + tu.assign_non_tensor(data, enable_timestep_staging=self.config.actor.get("enable_timestep_staging", False)) output = self.actor.train_mini_batch(data=data) return output.cpu() if output is not None else None