[diffusion, fsdp, cfg] feat: stage Qwen-Image training inputs by timestep

[WenzheWang]

What does this PR do?

Reduce retained diffusion-training tensors in two bounded steps:

• Release unused PPO/NFT model_output after backward by default; explicit return_model_output=True and inference keep their outputs.
• Add default-off synchronous timestep input staging for QwenImagePipeline, FlowGRPO or DiffusionNFT, FSDP/FSDP2, GPU, SP=1. Keep the caller's trajectory on CPU, copy prompt conditions once per micro-batch, and copy only current-step fields for each forward/backward.

PPO stages the current/next latent pair and matching loss inputs. NFT shares the clean latent/static noise and stages per-step noise when supplied. Dtypes, ordering, gradient accumulation and loss implementations are unchanged. Unsupported opt-in combinations fail before model initialization. DPO, other models/backends, sequence parallelism, async prefetch/pinned buffers and rollout transport are out of scope.

Checklist Before Starting

• Checked staging / trajectory / offload PRs, exact head and the author's open PRs. No equivalent Qwen training-side timestep-staging PR was found. #354 owns rollout SHM/GenRM overlap, #479 owns an SD3 train-side sampler, and #545 owns distillation control-plane scheduling.
• Title follows [{modules}] feat: ....

API and Usage Example

actor_rollout_ref.actor.fsdp_config.enable_timestep_staging=true

Consumed tensor inputs must already be CPU tensors without input gradients; there is no implicit full-trajectory GPU-to-CPU round trip. Inference bypasses staging. Explicitly retaining training outputs still incurs timestep-dependent output memory. Transfers are synchronous, so this is a memory/transfer-cost tradeoff, not a throughput guarantee.

Test

Main refresh (2026-09-08)

Merged main 23a5fc4f26a525dc559c968cf9807be00ae9eb27 into this branch as ba66c044cbca5eb9314bc184fb3de34c6d4b49c9; exact tested tree 8b32b8117690c830829faef112478d61f3e54118. Both the upstream merged-LoRA synchronization (#553) and the DAPO performance reference (#456) are retained alongside this PR's original timestep-staging/output-lifetime changes.

On the refreshed tree, 163 CPU tests passed, all 12 literal pre-commit run --all-files hooks passed, and all 4 trainer configs regenerated with no residual diff. CPU-only job completed 0:0 in 3m04s. The focused pytest command below was rerun with tests/workers/test_diffusers_fsdp_merged_lora_on_cpu.py additionally included. AST comparisons verify unchanged timestep/output-loss methods against the original PR and unchanged merged-LoRA methods against main; both documentation/evidence sections are preserved verbatim. Independent review: SAFE.

These refreshed CPU checks reused the runtime documented below. The original GPU matrix and measurements remain attached to 678d855d / tree 8b4000b3; they were not rerun or relabeled as measurements of this merge. No new algorithm, runtime, or feature scope is introduced. Official CI remains separate from these standalone checks.

Original implementation and GPU measurements (unchanged)

Base: 03655dc422b0437371dd856de00f3fd9fce91af9; exact declared verl pin: fefb080262e1c015a0ea05f958822a6a512dc795.
Validated candidate tree: 8b4000b353db491c7ed40351b33b9414ecb689bb (local tree independently matched).
Commit 678d855d8af837393a3cb836ccc1529e98027709 was created after validation and has exactly this tree.

bash scripts/generate_trainer_config.sh
pre-commit run --all-files
TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1 python3 -m pytest -q -x \
  tests/workers/test_diffusers_timestep_staging_on_cpu.py \
  tests/workers/test_diffusers_output_lifetime_on_cpu.py \
  tests/workers/test_diffusers_dtype_on_cpu.py \
  tests/workers/config/test_diffusion_config_on_cpu.py \
  tests/trainer/diffusion/test_worker_batch_projection_on_cpu.py \
  tests/special_sanity/test_config_docs.py

156 passed, all 12 literal all-file pre-commit hooks passed, all 4 flattened configs generated with no skip or residual diff. Coverage includes weakref ownership through backward, repeated updates, failure recovery, default/retention/inference behavior, row-distinct NFT timesteps and all three noise layouts, optional PPO loss inputs, invalid scope/input rejection, and documented top-level Hydra override propagation. A focused CPU test exercises real engine forward_step, registered Qwen adapters and PPO scheduler with a tiny linear projection; it is not transformer/FSDP hardware evidence.

GPU: all four profiles passed on 2 x NVIDIA A800 80 GB. Real Qwen transformer, adapters, scheduler, diffusion losses, FSDP/FSDP2 backward and AdamW updates. The immutable randomly initialized fixture has two layers, two heads, head dimension 32; native attention, BF16 model, SP=1, per-rank batch 4 / micro-batch 2, 1024x64 latent tokens (512x512 equivalent), positive/negative prompts and True-CFG 2.0. FP32 LoRA except explicitly BF16 for FSDP2/NFT.

Three arms (actual main, A-only output release, A+B input staging), four backend/algorithm profiles, and 1/8/32 timesteps: 36 arm-cases; 108 measured distributed optimizer updates plus 36 warmups. Both ranks passed all comparisons: 216 measured rank observations and 18 completion markers, not 216 independent replicates. Each arm consumes the same main inference outputs/old_log_probs; all inference tensors, training losses, local gradients and updated parameters compare at rtol=1e-5, atol=1e-6. Updates are finite/nonzero. State/optimizer/RNG are reset between paired observations. The exact-tree job completed 0:0 in 28m13s.

At 32 timesteps, values are descriptive medians of three repetitions, taking the maximum rank for each repetition. Starting allocated bytes match across arms within every profile/length.

Backend / algorithm	Main peak MiB	A peak MiB	A+B peak MiB	A seconds	A+B seconds	B vs A time
fsdp/flow_grpo	179.27	118.12	102.62	20.11	18.69	-7.1%
fsdp/diffusion_nft	308.93	121.61	106.11	22.82	23.08	1.2%
fsdp2/flow_grpo	179.28	118.06	102.56	8.43	8.60	2.1%
fsdp2/diffusion_nft	307.67	120.79	105.29	12.57	14.03	11.7%

The observed step-input CUDA storage is 16.50 MiB -> 1.00 MiB from A to A+B (about 93.9% lower) and stays near 1.00 MiB at 1/8/32 timesteps in this fixed-shape grid. B's incremental peak-allocated reduction is 15.50 MiB here. NFT uses supplied per-step noise in this GPU grid; that input-saving number does not apply to NFT runs that already generate noise on demand.

Tradeoff: FSDP2/NFT update time increases about 11.7% relative to A. Other timings also vary; arms ran in fixed order, so apparent speedups are not a controlled production-throughput claim. This small random model establishes the input-residency mechanism and numerical integration, not large-model peak-memory savings. The default-off flag lets users make a workload-specific memory/transfer-cost choice.

All measured-case aggregates (36 rows; three repeated max-rank observations per row)

variant,strategy,algorithm,steps,peak_allocated_MiB,start_allocated_MiB,peak_reserved_MiB,max_input_storage_MiB,median_max_rank_seconds,min_max_rank_seconds,max_max_rank_seconds
main,fsdp,flow_grpo,1,103.120605,66.664551,130.000000,1.004173,0.576789,0.572853,0.580672
a,fsdp,flow_grpo,1,102.619141,66.664551,130.000000,1.004173,0.633456,0.632962,0.634061
ab,fsdp,flow_grpo,1,102.619141,66.664551,130.000000,1.004173,0.630626,0.630528,0.635539
main,fsdp,flow_grpo,8,113.641113,66.664551,158.000000,4.504333,4.641108,4.636991,4.650221
a,fsdp,flow_grpo,8,106.119141,66.664551,128.000000,4.504333,4.645812,4.636187,4.649722
ab,fsdp,flow_grpo,8,102.619141,66.664551,130.000000,1.004173,5.080342,5.075288,5.102016
main,fsdp,flow_grpo,32,179.267578,66.664551,226.000000,16.504883,18.764022,18.597090,18.785161
a,fsdp,flow_grpo,32,118.119141,66.664551,126.000000,16.504883,20.112818,20.110825,20.169019
ab,fsdp,flow_grpo,32,102.619141,66.664551,130.000000,1.004173,18.687138,18.668437,18.701372
main,fsdp,diffusion_nft,1,107.863281,67.895508,136.000000,1.004166,0.717446,0.715278,0.717489
a,fsdp,diffusion_nft,1,106.112793,67.895508,134.000000,1.004166,0.803190,0.802838,0.804875
ab,fsdp,diffusion_nft,1,106.112793,67.895508,134.000000,1.004166,0.718866,0.717897,0.720760
main,fsdp,diffusion_nft,8,128.910645,67.895508,194.000000,4.504272,5.776797,5.774918,5.778564
a,fsdp,diffusion_nft,8,109.612793,67.895508,134.000000,4.504272,5.705620,5.702425,5.711195
ab,fsdp,diffusion_nft,8,106.112793,67.895508,134.000000,1.004166,6.468721,6.463199,6.473139
main,fsdp,diffusion_nft,32,308.934082,67.895508,430.000000,16.504639,25.628432,25.626415,25.659985
a,fsdp,diffusion_nft,32,121.612793,67.895508,150.000000,16.504639,22.815814,22.798162,22.829415
ab,fsdp,diffusion_nft,32,106.112793,67.895508,134.000000,1.004166,23.084161,23.032093,23.085334
main,fsdp2,flow_grpo,1,103.061523,66.579590,128.000000,1.004173,0.255264,0.249817,0.260705
a,fsdp2,flow_grpo,1,102.560059,66.579590,128.000000,1.004173,0.260484,0.255807,0.261830
ab,fsdp2,flow_grpo,1,102.560059,66.579590,128.000000,1.004173,0.276065,0.275778,0.276744
main,fsdp2,flow_grpo,8,113.582031,66.579590,156.000000,4.504333,2.137726,2.082971,2.169007
a,fsdp2,flow_grpo,8,106.060059,66.579590,128.000000,4.504333,2.119337,2.041742,2.129787
ab,fsdp2,flow_grpo,8,102.560059,66.579590,128.000000,1.004173,2.239548,2.112779,2.276339
main,fsdp2,flow_grpo,32,179.283691,66.579590,224.000000,16.504883,8.543005,8.445626,8.688514
a,fsdp2,flow_grpo,32,118.060059,66.579590,126.000000,16.504883,8.426803,8.425714,8.441801
ab,fsdp2,flow_grpo,32,102.560059,66.579590,128.000000,1.004173,8.602956,8.547832,8.604706
main,fsdp2,diffusion_nft,1,107.041016,66.580078,134.000000,1.004166,0.405020,0.404238,0.406048
a,fsdp2,diffusion_nft,1,105.290527,66.580078,132.000000,1.004166,0.453305,0.448536,0.454265
ab,fsdp2,diffusion_nft,1,105.290527,66.580078,132.000000,1.004166,0.401611,0.400806,0.402562
main,fsdp2,diffusion_nft,8,128.047852,66.580078,190.000000,4.504272,3.160905,3.158372,3.168793
a,fsdp2,diffusion_nft,8,108.790527,66.580078,132.000000,4.504272,3.153288,3.145452,3.155891
ab,fsdp2,diffusion_nft,8,105.290527,66.580078,132.000000,1.004166,3.159030,3.157298,3.162322
main,fsdp2,diffusion_nft,32,307.669434,66.580078,428.000000,16.504639,12.635724,12.631546,12.642104
a,fsdp2,diffusion_nft,32,120.790527,66.580078,148.000000,16.504639,12.565981,12.558570,12.567242
ab,fsdp2,diffusion_nft,32,105.290527,66.580078,132.000000,1.004166,14.030557,14.029209,14.036057

The reused Linux validation runtime has PyTorch 2.11.0+cu130, diffusers 0.38.0 and older vLLM-Omni compatibility dependencies; it is not full current official-CI dependency parity (current diffusers requirement is >=0.40). These checks do not validate rollout transport, reward-training convergence, pretrained numerical alignment, SP>1, NPU or other architectures. The existing FSDP2/NFT FP32-LoRA backward dtype failure was reproduced on unmodified main in this runtime; the declared FSDP2/NFT profile uses BF16 LoRA. This contribution does not fix or claim the FP32 profile.

Design & Code Changes

The diffusion-owned FSDP config subclass exposes the flag through the existing actor engine config. The shared PPO/NFT loop validates an explicit field map before training, selects CPU slices before device transfer, rebases all temporal fields together to local step zero, and releases step-owned inputs after backward. Shared inputs live only for the current micro-batch; no persistent cache or new adapter abstraction is added.

Checklist Before Submitting

• Read the contributing guide and checked overlap.
• Ran complete pre-commit hooks and updated generated configuration.
• Added usage/boundary documentation.
• Added *_on_cpu.py regression tests discovered by the existing CPU workflow. GPU integration is manual; no CI workflow or runner policy changed.
• Independent code/evidence review and exact-tree CPU/GPU validation passed.
• The human submitter confirmed review of every changed line, including the newly added B scope, before publication.

AI assistance was used. The human submitter has reviewed every changed line and authorized publication.

Complete standalone GPU reproducer and A-only control patch

Save the Python script as gpu_validate.py and the diff as a-control.patch outside the repository. Use a compatible Linux CUDA environment with the exact verl pin and versions above; this does not install dependencies or download model weights. From this PR checkout:

repro_script=$(realpath gpu_validate.py)
control_patch=$(realpath a-control.patch)
candidate_checkout=$PWD
experiment_root=$(mktemp -d)
git worktree add --detach "$experiment_root/main" 03655dc422b0437371dd856de00f3fd9fce91af9
git worktree add --detach "$experiment_root/a" 03655dc422b0437371dd856de00f3fd9fce91af9
(cd "$experiment_root/a" && git apply "$control_patch")
mkdir -p "$experiment_root/reference"
export TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export OMP_NUM_THREADS=8 MKL_NUM_THREADS=8
python3 "$repro_script" --build-only --variant main \
  --reference "$experiment_root/reference" --model "$experiment_root/model" \
  --output "$experiment_root/reference"
for arm in main a ab; do
  source_dir="$experiment_root/$arm"
  if [ "$arm" = ab ]; then source_dir="$candidate_checkout"; fi
  for profile in fsdp fsdp2-ppo fsdp2-nft; do
    case "$profile" in
      fsdp) strategy=fsdp; algorithms=(flow_grpo diffusion_nft); lora=float32 ;;
      fsdp2-ppo) strategy=fsdp2; algorithms=(flow_grpo); lora=float32 ;;
      fsdp2-nft) strategy=fsdp2; algorithms=(diffusion_nft); lora=bfloat16 ;;
    esac
    result_dir="$experiment_root/results/$arm/$profile"
    mkdir -p "$result_dir"
    (cd "$source_dir" && PYTHONPATH="$source_dir:${PYTHONPATH:-}" \
      torchrun --standalone --nnodes=1 --nproc-per-node=2 "$repro_script" \
      --variant "$arm" --model "$experiment_root/model" \
      --reference "$experiment_root/reference" --output "$result_dir" \
      --strategies "$strategy" --algorithms "${algorithms[@]}" --lora-dtype "$lora")
  done
done

A-only production control (documentation/tests do not affect this GPU harness):

diff --git a/verl_omni/workers/engine/fsdp/diffusers_impl.py b/verl_omni/workers/engine/fsdp/diffusers_impl.py
index 513af77..bc72f2c 100644
--- a/verl_omni/workers/engine/fsdp/diffusers_impl.py
+++ b/verl_omni/workers/engine/fsdp/diffusers_impl.py
@@ -825,6 +825,7 @@ class DiffusersFSDPEngine(LoRAAdapterMixin, BaseEngine, ABC):
         timesteps_key: str,
     ) -> dict:
         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)
 
@@ -848,6 +849,9 @@ class DiffusersFSDPEngine(LoRAAdapterMixin, BaseEngine, ABC):
                     )
                     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)
             output_lst.append(meta_info_lst)

Harness SHA256: acc900870c96dd97b743cb5cbbf67a9e3196f3ceca31aa40f73674d0d388b063.

"""Frozen main/A/A+B input-residency integration; not a production benchmark."""

import argparse
import gc
import hashlib
import json
import os
import time
from dataclasses import asdict
from functools import partial
from pathlib import Path

import torch
import torch.distributed as dist
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageTransformer2DModel
from tensordict import TensorDict
from torch.distributed.tensor import DTensor
from verl.trainer.config import CheckpointConfig
from verl.utils import tensordict_utils as tu
from verl.utils.fsdp_utils import fsdp_version
from verl.workers.config import FSDPEngineConfig, FSDPOptimizerConfig

from verl_omni.workers.config import DiffusionActorConfig, DiffusionLossConfig, DiffusionModelConfig
from verl_omni.workers.config.diffusion.rollout import DiffusionPipelineConfig, DiffusionRolloutAlgoConfig
from verl_omni.workers.engine.fsdp.diffusers_impl import NFTDiffusersFSDPEngine, PPODiffusersFSDPEngine
from verl_omni.workers.utils.losses import diffusion_loss

SEED = 20260906


def atomic_json(path, value):
    path = Path(path)
    pending = path.with_suffix(path.suffix + ".pending")
    pending.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n")
    pending.replace(path)


def make_fixture(path):
    torch.manual_seed(SEED)
    path = Path(path)
    path.mkdir(exist_ok=False)
    model = QwenImageTransformer2DModel(
        num_attention_heads=2, attention_head_dim=32, num_layers=2,
        in_channels=64, out_channels=16, patch_size=2, joint_attention_dim=32,
        axes_dims_rope=(8, 12, 12), guidance_embeds=False,
    )
    model.save_pretrained(path / "transformer")
    FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True).save_pretrained(path / "scheduler")
    atomic_json(path / "model_index.json", {"_class_name": "QwenImagePipeline"})
    atomic_json(path / "fixture.json", {
        "seed": SEED, "weights": "random", "scope": "GPU integration only",
        "files": {str(p.relative_to(path)): hashlib.sha256(p.read_bytes()).hexdigest()
                  for p in sorted(path.rglob("*")) if p.is_file()},
    })


def local(tensor):
    return tensor.to_local() if isinstance(tensor, DTensor) else tensor


def snapshot(engine, gradients=False):
    return {name: local(p.grad if gradients else p).detach().float().cpu().clone()
            for name, p in engine.module.named_parameters()
            if p.requires_grad and (not gradients or p.grad is not None)}


def restore(engine, state):
    engine.optimizer_zero_grad()
    engine.optimizer.state.clear()
    with torch.no_grad():
        for name, p in engine.module.named_parameters():
            if name in state:
                local(p).copy_(state[name])
    torch.manual_seed(SEED)
    torch.cuda.manual_seed_all(SEED)


def check_equal(actual, expected):
    assert actual.keys() == expected.keys()
    for name in actual:
        torch.testing.assert_close(actual[name], expected[name], rtol=1e-5, atol=1e-6)


def make_batch(engine, algorithm, steps):
    gen = torch.Generator().manual_seed(SEED + dist.get_rank())
    common = {
        "prompt_embeds": torch.randn(4, 16, 32, generator=gen, dtype=torch.bfloat16),
        "prompt_embeds_mask": torch.ones(4, 16, dtype=torch.int32),
        "negative_prompt_embeds": torch.randn(4, 16, 32, generator=gen, dtype=torch.bfloat16),
        "negative_prompt_embeds_mask": torch.ones(4, 16, dtype=torch.int32),
    }
    if algorithm == "flow_grpo":
        common.update(
            all_latents=torch.randn(4, steps + 1, 1024, 64, generator=gen),
            all_timesteps=engine.scheduler.timesteps[:steps].cpu().expand(4, -1).clone(),
            advantages=torch.tensor([0.25, 0.75, -0.25, -0.75])[:, None].expand(-1, steps).clone(),
            old_log_probs=torch.zeros(4, steps),
        )
    else:
        common.update(
            latents_clean=torch.randn(4, 1024, 64, generator=gen),
            forward_noise=torch.randn(4, steps, 1024, 64, generator=gen),
            train_timesteps=torch.stack([torch.linspace(900, 100, steps)[torch.randperm(steps, generator=gen)] for _ in range(4)]),
            reward_prob=torch.tensor([0.25, 0.75, 0.25, 0.75])[:, None].expand(-1, steps).clone(),
        )
    data = TensorDict(common, batch_size=[4])
    tu.assign_non_tensor(data, micro_batch_size_per_gpu=2, height=512, width=512, vae_scale_factor=8)
    return data


def run_case(args, strategy, algorithm):
    torch.manual_seed(SEED)
    model_config = DiffusionModelConfig(
        path=args.model, algorithm=algorithm, load_tokenizer=False,
        attn_backend="native", enable_gradient_checkpointing=False,
        lora_rank=8, lora_alpha=16, lora_dtype=args.lora_dtype,
        policy_state_adapters=("default", "old") if algorithm == "diffusion_nft" else ("default",),
        pipeline=DiffusionPipelineConfig(height=512, width=512, num_inference_steps=40, true_cfg_scale=2.0),
        algo=DiffusionRolloutAlgoConfig(noise_level=0.8, sde_type="sde"),
    )
    engine_config = FSDPEngineConfig(
        strategy=strategy, fsdp_size=dist.get_world_size(), ulysses_sequence_parallel_size=1,
        model_dtype="bfloat16", dtype="bfloat16", use_orig_params=True,
        mixed_precision={"param_dtype": "bfloat16", "reduce_dtype": "float32", "buffer_dtype": "float32"},
        param_offload=False, optimizer_offload=False, forward_only=False,
    )
    if args.variant == "ab":
        from verl_omni.workers.config import FSDPDiffusionEngineConfig
        engine_config = FSDPDiffusionEngineConfig(**asdict(engine_config), enable_timestep_staging=True)
    optim = FSDPOptimizerConfig(lr=1e-4, clip_grad=1.0, total_training_steps=100)
    actor = DiffusionActorConfig(
        strategy=strategy, ppo_micro_batch_size_per_gpu=2, rollout_n=1,
        diffusion_loss=DiffusionLossConfig(loss_mode=algorithm, clip_ratio=0.2),
    )
    cls = PPODiffusersFSDPEngine if algorithm == "flow_grpo" else NFTDiffusersFSDPEngine
    engine = cls(model_config, engine_config, optim, CheckpointConfig())
    engine.initialize()
    assert fsdp_version(engine.module) == (1 if strategy == "fsdp" else 2)
    if algorithm == "diffusion_nft":
        engine.copy_adapter("default", "old")
    initial = snapshot(engine)
    assert initial
    loss_fn = partial(diffusion_loss, config=actor)
    records = []
    for steps in (1, 8, 32):
        data = make_batch(engine, algorithm, steps)
        with engine.eval_mode():
            inference_a = engine.infer_batch(data.clone())
            flagged = data.clone()
            tu.assign_non_tensor(flagged, return_model_output=False)
            inference_b = engine.infer_batch(flagged)
        assert inference_a["model_output"]
        for name, tensor in inference_a["model_output"].items():
            assert tensor.shape[:2] == (4, steps), (name, tensor.shape)
            torch.testing.assert_close(tensor, inference_b["model_output"][name], rtol=0, atol=0)
        inference_path = Path(args.reference) / f"inference-{strategy}-{algorithm}-t{steps}-rank{dist.get_rank()}.pt"
        actual_inference = {name: tensor.detach().cpu() for name, tensor in inference_a["model_output"].items()}
        if args.variant == "main":
            torch.save(actual_inference, inference_path)
            reference_inference = actual_inference
        else:
            reference_inference = torch.load(inference_path, weights_only=True)
            check_equal(actual_inference, reference_inference)
        if algorithm == "flow_grpo":
            data["old_log_probs"] = reference_inference["log_probs"]
        del actual_inference, reference_inference
        del inference_a, inference_b, flagged

        reference_path = Path(args.reference) / f"{strategy}-{algorithm}-t{steps}-rank{dist.get_rank()}.pt"
        expected = None if args.variant == "main" else torch.load(reference_path, weights_only=True)
        original_forward_step = engine.forward_step
        input_bytes = []

        def observed_forward_step(micro_batch, *forward_args, **forward_kwargs):
            storages = {value.untyped_storage().data_ptr(): value.untyped_storage().nbytes()
                        for value in micro_batch.values()
                        if isinstance(value, torch.Tensor) and value.is_cuda}
            input_bytes.append(sum(storages.values()))
            if args.variant == "ab":
                assert forward_kwargs["step"] == 0
                if algorithm == "flow_grpo":
                    assert micro_batch["all_latents"].shape[1] == 2
                else:
                    assert micro_batch["forward_noise"].shape[1] == 1
                assert micro_batch["all_timesteps" if algorithm == "flow_grpo" else "train_timesteps"].shape[1] == 1
            return original_forward_step(micro_batch, *forward_args, **forward_kwargs)

        engine.forward_step = observed_forward_step
        # One discarded warm-up then three fixed observations for every arm.
        for repetition in range(-1, 3):
            restore(engine, initial)
            batch = data.clone()
            input_bytes.clear()
            gc.collect()
            torch.cuda.empty_cache()
            dist.barrier()
            with engine.train_mode():
                torch.cuda.synchronize()
                start_allocated = torch.cuda.memory_allocated()
                torch.cuda.reset_peak_memory_stats()
                started = time.perf_counter()
                output = engine.train_batch(batch, loss_fn)
                torch.cuda.synchronize()
                seconds = time.perf_counter() - started
                measurement = {
                    "variant": args.variant, "rank": dist.get_rank(), "strategy": strategy,
                    "algorithm": algorithm, "steps": steps, "repetition": repetition,
                    "loss": output["loss"], "grad_norm": output["metrics"]["grad_norm"],
                    "seconds": seconds, "start_allocated": start_allocated,
                    "end_allocated": torch.cuda.memory_allocated(),
                    "peak_allocated": torch.cuda.max_memory_allocated(),
                    "peak_reserved": torch.cuda.max_memory_reserved(),
                    "max_input_storage_bytes": max(input_bytes),
                }
                gradients = snapshot(engine, gradients=True)
                updated = snapshot(engine)
                assert torch.isfinite(torch.tensor(measurement["loss"])).all()
                assert 0 < measurement["grad_norm"] < float("inf")
                changed = sum(torch.count_nonzero(updated[n] - initial[n]).item() for n in initial)
                global_changed = torch.tensor(changed, device="cuda", dtype=torch.long)
                dist.all_reduce(global_changed)
                assert global_changed.item() > 0
                measurement["changed_local_parameters"] = changed
                assert bool(output["model_output"]) is (args.variant == "main")
                assert all(v.device.type == "cpu" for v in batch.values() if isinstance(v, torch.Tensor))
                reference_value = {"loss": torch.tensor(output["loss"]), "gradients": gradients, "updated": updated}
                if expected is None:
                    expected = reference_value
                    torch.save(expected, reference_path)
                torch.testing.assert_close(reference_value["loss"], expected["loss"], rtol=1e-5, atol=1e-6)
                check_equal(gradients, expected["gradients"])
                check_equal(updated, expected["updated"])
                del output, batch
            if repetition >= 0:
                records.append(measurement)
                atomic_json(Path(args.output) / f"{strategy}-{algorithm}-rank{dist.get_rank()}.json", records)
            if dist.get_rank() == 0:
                print("CASE_PASS", args.variant, strategy, algorithm, steps, repetition, flush=True)
        engine.forward_step = original_forward_step
        restore(engine, initial)
    del engine, initial
    gc.collect()
    torch.cuda.empty_cache()
    dist.barrier()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--variant", choices=("main", "a", "ab"), required=True)
    parser.add_argument("--reference", required=True)
    parser.add_argument("--build-only", action="store_true")
    parser.add_argument("--strategies", nargs="+", choices=("fsdp", "fsdp2"), default=("fsdp", "fsdp2"))
    parser.add_argument("--algorithms", nargs="+", choices=("flow_grpo", "diffusion_nft"), default=("flow_grpo", "diffusion_nft"))
    parser.add_argument("--lora-dtype", choices=("float32", "bfloat16"), default="float32")
    args = parser.parse_args()
    if args.build_only:
        make_fixture(args.model)
        return
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    dist.init_process_group("nccl")
    assert dist.get_world_size() == 2
    import diffusers, verl
    from verl_omni.workers.engine.fsdp import diffusers_impl
    print("SOURCE", args.variant, diffusers_impl.__file__, verl.__file__, diffusers.__version__, flush=True)
    print("GPU_PREFLIGHT", dist.get_rank(), torch.__version__, torch.cuda.get_device_name(), flush=True)
    for strategy in args.strategies:
        for algorithm in args.algorithms:
            run_case(args, strategy, algorithm)
    atomic_json(Path(args.output) / f"complete-rank{dist.get_rank()}.json", {
        "status": "passed", "rank": dist.get_rank(), "world_size": 2,
        "cases": 3 * len(args.strategies) * len(args.algorithms),
        "measured_updates": 9 * len(args.strategies) * len(args.algorithms),
        "warmup_updates": 3 * len(args.strategies) * len(args.algorithms),
        "lora_dtype": args.lora_dtype,
        "scope": "tiny random Qwen-Image GPU integration; not production performance",
    })
    dist.destroy_process_group()


if __name__ == "__main__":
    main()