Wan2.1 finetuning script seems to only support bs = 1

[connorxian]

This issue aims to raise the attention of the abnormal training setting. I found that serval settings in the model and training files only code for batch size = 1.

train_wan_t2v.py function training_step:
torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))

models.wan_video_dit.py: Line 252:
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)

[cfeng16]

I am also wondering whether code can support bs > 1

[connorxian]

I am also wondering whether code can support bs > 1

I have solved this problem with some modification on the mentioned code.

DiffSynth-Studio/diffsynth/models/wan_video_dit.py

Line 255 in 3b3e1e4

def forward(self, x, t_mod):

def forward(self, x, t_mod):
    batch_size = x.shape[0]
    shift, scale = (self.modulation.repeat(batch_size,1,1).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1)
    x = (self.head(self.norm(x) * (1 + scale) + shift))
    return x

DiffSynth-Studio/diffsynth/schedulers/flow_match.py

Line 65 in 3b3e1e4

def add_noise(self, original_samples, noise, timestep):

def add_noise(self, original_samples, noise, timestep):
    if isinstance(timestep, torch.Tensor):
        timestep = timestep.cpu()
    timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim = 0) 
    sigma = self.sigmas[timestep_id].to(device=original_samples.device, dtype=original_samples.dtype)
    while len(sigma.shape) < len(original_samples.shape):
        sigma = sigma.unsqueeze(-1)
    sample = (1 - sigma) * original_samples + sigma * noise
    return sample

[connorxian]

the training script seem to has been updated serval days ago, please check this line:

DiffSynth-Studio/diffsynth/pipelines/wan_video_new.py

Line 235 in 3b3e1e4

timestep_id = torch.randint(0, self.scheduler.num_train_timesteps, (1,))

timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (batch_size,))

[Artiprocher]

@connorxian We truly understand your request. However, we believe that the engineering effort required to support a batch size greater than 1 is not justified by the limited benefits it brings.

On one hand, most GPUs do not have enough memory to support training with a batch size greater than 1. On the other hand, using gradient accumulation can effectively achieve the same result. Additionally, most developers of the Wan series models we've integrated did not implement support for batch sizes larger than 1, which would require us to conduct extensive testing to enable this feature.

For these reasons, we currently have no plans to support batch sizes greater than 1.

[ShaoDingBao]

Generally, fine-tuning video generation models seems to require joint training with both images and videos. For images, isn't a batch size of 1 too small? If we use gradient accumulation, it seems like that would lead to very low training efficiency. Also, according to what you mentioned in the training documentation, it's currently not recommended to mix images and videos in the same dataset. Does that mean we can only train with images first, and then with videos? What's the difference in training effectiveness compared to joint image-video training? Have you conducted any related experiments on this? @Artiprocher

[Artiprocher]

@ShaoDingBao We do not recommend mixing images and videos in training because of load balancing issues in multi-GPU distributed training. At each step, gradient synchronization is required, but the size of the video or image assigned to each process may vary, causing each process to wait for the slowest one, resulting in resource wastage. However, this issue does not occur when using single GPU training.

[Artiprocher]

@ShaoDingBao For images, a batch size of 1 does not cause efficiency issues. Modern model architectures have fully utilized the parallelism of GPUs, and even with a batch size of 1, the GPU operates at full capacity. Gradient accumulation does not introduce additional tensor computations, making it a reliable alternative for achieving large effective batch sizes.

The original Wan model was trained first on image data and then on video data during the pre-training phase. This approach allows for better utilization of multi-GPU resources. The initial intention behind designing the image-video joint training feature was to make it more convenient for users to leverage their training data, especially since high-quality datasets are difficult to collect.

[ShaoDingBao]

@Artiprocher My previous understanding is that a batch size of 1 might lead to inefficient GPU utilization and potentially waste VRAM, as GPUs are designed for parallel processing. In my experience, inference with a batch size of 8 in one go is generally faster than inference with a batch size of 1 repeated 8 times.

You mentioned that for training, a batch size of 1 doesn't cause efficiency issues because modern architectures fully utilize GPU parallelism. While I understand that gradient accumulation can achieve a large effective batch size, I'm still unclear about the speed implications for training iterations.

Specifically, will training with a physical batch size of 1 for 8 steps (with gradient accumulation) be slower than training with a physical batch size of 8 for 1 step? My intuition suggests the latter would be faster due to reduced overhead and better GPU saturation. Could you elaborate on this aspect of training speed, particularly concerning the iteration time difference?
Thank you very much for your detailed previous response; it truly helped me understand many complex details!

[Artiprocher]

@ShaoDingBao Let me explain the training process, which consists of three steps:

1. Forward pass (loss = xxx)
2. Backward pass (loss.backward())
3. Update model (optimizer.step())

When batch size = 8

1. Forward pass (process 8 videos in parallel)
2. Backward pass (process 8 videos in parallel)
3. Update model

When using gradient accumulation to achieve batch size = 8

1. [Forward - Backward - Gradient Sync] - [Forward - Backward - Gradient Sync] - ... (repeated 8 times)
2. Update model

The key question is: will performing 8 iterations of [Forward - Backward - Gradient Sync] be slower?

Strictly speaking, there are three main components contributing to any performance difference:

1. Parallelism: Thanks to modern model architectures and GPU optimizations, a batch size of 1 is often sufficient to fully utilize GPU computing power. Therefore, the difference in tensor computation time between batch sizes of 1 and 8 is minimal.

2. Computation graph metadata inference: This involves determining where gradients should be propagated. It is not a tensor-based operation and is extremely fast — so much so that its impact can be ignored.

3. Gradient synchronization: This is where performance issues arise during multi-process training. Gradient sync is essentially redundant computation, and modern training frameworks like Hugging Face Accelerate optimize this part by skipping unnecessary synchronization unless explicitly needed. For more details, see: [Gradient Synchronization](https://huggingface.co/docs/accelerate/concept_guides/gradient_synchronization)

Conclusion

Gradient accumulation will definitely be slower than directly using a large batch size. However, the difference is negligible in practice. Moreover, as GPU hardware improves and training frameworks continue to evolve, this performance gap will only get smaller over time.

[ShaoDingBao]

Thank you so much for your patient and detailed explanation! Your clear breakdown of the mechanisms, especially regarding GPU parallelism, computation graph metadata inference, and gradient synchronization, truly deepened my understanding of these complex concepts.

I completely agree with your conclusion that while gradient accumulation might be theoretically slower, the difference is negligible in practice and will continue to shrink. This insight is incredibly helpful for optimizing my training strategies.

Thanks again for your valuable time and expertise!

[yfpeng1234]

the training script seem to has been updated serval days ago, please check this line:

DiffSynth-Studio/diffsynth/pipelines/wan_video_new.py

Line 235 in 3b3e1e4

timestep_id = torch.randint(0, self.scheduler.num_train_timesteps, (1,))
timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (batch_size,))

Hi @connorxian , thanks for sharing this, I try that but the loss can not converge

[connorxian]

@yfpeng1234 Thanks for attention. From my experience, diffusion-based fine-tuning doesn't result in a loss converge, although the loss goes down initially. Maybe you can focus more on the samples from the intermediate checkpoint.

[yfpeng1234]

Thanks a lot for your prompt reply @connorxian , I will try that