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

"""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()