Trying to revive Validation loss, could use some insight.

[mcmanus82]

As the title says, i'm trying to fix up this old PR mostly for Wan2.1 / Skyreels. My first goal is to make it full functional and properly working, then fix up issues that the original had and add maybe progress bar, config options etc.

In its current state it seems to work, at least that is my general assumption. Val loss is calculated, reported and logged. The problem is that it does not really follow the expected curve and i don't know if that is because of Wan or my dataset/captions or the code has some flaws. It does follow a somewhat downwards trend but it is a lot more erratic.

So, here is the main def validate code part, i have also attached a zip with all the changed files if someone wants to try it. Be warned it only works with Wan2.1 / Skyreels (t2v/df) and you cannot turn it off :D you have to have a [val_datasets] in your toml
Any feedback is appreciated and if you can point out any flaws in the code / logic please do so.(i'm a noob at python)
Vall_loss.zip

    def validate(
        accelerator,
        transformer,
        val_dataloader,
        noise_scheduler,
        args,
    ):
        # Unwrap and switch to inference mode
        unwrapped_model = accelerator.unwrap_model(transformer)
        unwrapped_model.switch_block_swap_for_inference()
        unwrapped_model.eval()

        fixed_timesteps = [600, 700, 800, 900]
        fixed_seed = 42
        losses = []

        with torch.no_grad():
            for step1, batch in enumerate(val_dataloader):
                clip_fea = None
                latents = batch["latents"]
                context = [t.to(device=accelerator.device, dtype=network_dtype) for t in batch["t5"]]

                # Put everything on the correct device/dtype
                latents = latents.to(accelerator.device, dtype=torch.float32)
                lat_f, lat_h, lat_w = latents.shape[2:5]
                seq_len = lat_f * lat_h * lat_w // (self.config.patch_size[0] * self.config.patch_size[1] * self.config.patch_size[2])

                # Override RNG inside a local fork_rng context:
                with torch.random.fork_rng(devices=[accelerator.device]):
                    seed = fixed_seed + step1
                    torch.manual_seed(seed)
                    # If on CUDA:
                    if accelerator.device.type == "cuda":
                        torch.cuda.manual_seed(seed)

                    noise = torch.randn_like(latents)  # Now gives stable noise each time.

                # Cycle through multiple fixed timestep for each image, 
                for step2, ts in enumerate(fixed_timesteps, start=1):
                    timesteps = torch.full(
                        (latents.size(0),),
                        ts,
                        device=accelerator.device,
                        dtype=torch.float32
                    )

                    # Make "noisy_model_input" deterministically
                    sigma = get_sigmas(
                        noise_scheduler,
                        timesteps,
                        accelerator.device,
                        latents.dim(),
                        latents.dtype
                    )
                    noisy_model_input = sigma * noise + (1.0 - sigma) * latents

                    # Forward pass
                    pred = unwrapped_model(
                        noisy_model_input,
                        t=timesteps,                        
                        context=context,
                        clip_fea=clip_fea,
                        seq_len=seq_len,
                    )
                    pred_noise = torch.stack(pred, dim=0)

                    # MSE target
                    target = noise - latents
                    val_loss = torch.nn.functional.mse_loss(pred_noise, target, reduction="mean")
                    losses.append(val_loss.item())
                    print(f"[DEBUG] Validation Image:{step1} Timestep:{ts} Val_loss={val_loss}")

            # Switch back to train mode if necessary
            unwrapped_model.train()
            unwrapped_model.switch_block_swap_for_training()

            return float(np.mean(losses))

[Sarania]

Heya, I had previously had a look at this but I don't know /anything/ about validation loss. What exactly is it? AI said "Validation loss measures how well a machine learning model performs on a separate dataset that it hasn't seen during training." but I'm not quite sure what it's getting at. Just calculating loss for samples or something? Anyway nothing particularly stands out to me here but I have a lot less familiarity with the training parts of the code. I asked my ChatGPT o4-mini-high but that didn't produce anything helpful either. Sorry!

[Sarania]

Oh one thing I did just notice compared to the normal training loop is you aren't doing this:

        # ensure the hidden state will require grad
        if args.gradient_checkpointing:
            noisy_model_input.requires_grad_(True)
            for t in context:
                t.requires_grad_(True)
            if image_latents is not None:
                image_latents.requires_grad_(True)
            if clip_fea is not None:
                clip_fea.requires_grad_(True)

You might be able to skip the i2v stuff but im not sure if you can skip all the gradient checkpointing stuff. Or this might be irrelevant XD but i tried!

[mcmanus82]

It basically does a forward pass with a validation image (not included in training) and measures how "well" the current model measures against it. How seeding with a fixed noise at a given timestep measures against what the model would give for those parameters. The thing is i'm new to this and python too.. so its a miracle it even works to any extent.. :)

I'm testing on the 1.3b wan2.1 so i did not enable gradient checkpointing (3090) but yeah i figured that would be needed later.

[Sarania]

Do you have any prior coding experience? Python is not that hard to pick up and AI can help you a lot if you use it in the right way(don't put code you don't understand in your projects, take time to ask questions, understand it can make mistakes, treat it like a partner who works on your project together with you and the amount of work you can get done is OBSCENE.) I only started with Python about 1.5 years ago now, though I have almost 30 years experience as a coder in general. Overall I REALLY like Python, with the notable exception of dealing with environments/venvs and versions.

Oh okay I get it now thanks for the explanation! So you're just doing a step that doesn't include the backwards pass to compute a separate loss metric I see how that's useful! What's the reason for val_loss = torch.nn.functional.mse_loss(pred_noise, target, reduction="mean") in your code versus loss = torch.nn.functional.mse_loss(model_pred.to(network_dtype), target, reduction="none") in the main code e.g. why using mean instead of none? I'm basically playing "spot the difference" with the base training code at this point lol.

[mcmanus82]

I have about 15 years of IT experience, mostly testing, low level coding like unit tests. Mainly worked with C#. So coding in general is not that foreign to me. The part of python that gives me the creeps is that everything calls another function (def) which calls another one and i have to dig for like 15 minutes just to find anything :D

The current version i'm testing now has "none" and i manually reduce the values except batch, otherwise it would not be translatable to a scalar.
Atm it looks like this, and the results are almost identical with both methods:
val_loss = torch.nn.functional.mse_loss(pred_noise, target, reduction="none")
val_loss = val_loss.mean(dim=list(range(1, len(val_loss.shape))))

[mcmanus82]

This is the last run i did. 24 image dataset with a repeat of 2, so 48 samples per epoch. 3 validation images. Wan2.1 1.3b t2v. 5e-5 lr 64/1 rank/alpha.

According to this it bottomed out at 20-25 epoch. (left graph is the mean val loss, the 3 on the right are for the separate images)
But when i mage actual 3 sec videos with the lora at 20-25 epich, it's not even half done and only vaguely looks like the dataset person.

[Sarania]

Well your loss for image_1 looks wonky af while the other two look reasonable but I don't have experience with Wan 1.3b either. Usually when you see wonky loss like that you either have an image that's tripping it up or your captions are off so... What are your captions like? I know that can make ALL the difference in the world when training. For Wan I get the best results with verbose, similar structure captions(by that I mean describing the same parts of the scene in roughly the same area of the caption each time). I tend to do like:

"[wide shot/medium shot/medium close up/close up] of [a man/a woman/Name] doing [thing] in [place]. They are [appearance] and wearing [outfit]. [Optional second character description]. [Reinforce action/movement with more details]. [Describe setting, lighting, and additional camera details]"

I'm sorry I can't be more help with this but I'm definitely interested in making use of it when it gets sorted out. Hopefully kohya will merge it to Musubi, but I will definitely merge it to https://github.com/Sarania/blissful-tuner/ either way!

[mcmanus82]

Tried different datasets, and i used these exact same ones for flux training where they did their job well.
Here is another run, this was 1e-5 64/1 R/A and looks a bit more normal but it was still off by a mile.

I also noticed that commenting out validation, resulted in a slightly different loss/average & loss/epoch curve, not huge but noticeable. Which could be a really bad sign :)

To be honest i'm having a hard time making a decent character lora for wan. My usual evaluation method seem to fall on their face :/
Sample images made during training, have little relation to the ones i make in Comfy (KijaiWan) - step 25, cfg 6, shift 5, 81 frames at 480x832(no teacache or any other extra stuff)
I have read the discussion about Wan training, tried a few things from it but did not make much difference.
What lr, dim/alpha, loraplus_lr_ratio=X would you use for a dataset of 24 512x512 images and at what approximate step would that converge in your experience?
My last attempt had a somewhat decent result at epoch 65(1560 step) with 32/1, 2e-4 and loraplus_lr_ratio=4

[Sarania]

So the only time I've trained an image only dataset for Wan I used very similar settings to what I use for video:

LR 2e-5 with loraplus_lr_ratio 4, network_dim 16, network_alpha 16, constant_with_warmup 100 warmup steps out of 2400 total steps with came_pytorch optimizer in FP8 scaled precision with FP16 base weights, the rest default. This was with the 14B model I've never trained the 1.3B model.

I feel like your LR is maybe too high, especially with loraplus. I know when I was training Hunyuan on images and my LR was still too high I would get some what of the way there but never fully pick up like her freckles and fully converge. I've had /superb/ luck with 2e-5 and LoraPlus of 4 for videos with Wan when using network_dim=network_alpha(I almost always stick to dim 16 for motion LoRA, sometimes even dim 8 for simple character LoRA. These are BIG base models so you don't need as high a dim as like SD or SDXL. Note that this advice does not apply to 1.3B since it is small.) You can use a higher LR when using alpha=1 or != dim because of the way alpha scaling works, but I get better results turning that effect off(that's what dim=alpha does) and using lower LR(Edit: For Wan and Flux anyway. Hunyuan remains to be seen.) Hope it helps!