add torch.no_grad() for block swap

[kohya-ss]

This may solve a potential issue and prevent loss=NaN in long training. Further testing is needed to see if existing training works as well as before.

[kohya-ss]

This fix appears to enable multi-GPU training of DDP, but further testing is required.

[FurkanGozukara]

This fix appears to enable multi-GPU training of DDP, but further testing is required.

with flux we were never able to do multi gpu training with block swap

so it is possible?

[kohya-ss]

so it is possible?

I think so. I will add a new branch to support this on sd-scripts repo today.

[FurkanGozukara]

so it is possible?

I think so. I will add a new branch to support this on sd-scripts repo today.

awesome

[kohya-ss]

This fix appears to enable multi-GPU training of DDP, but further testing is required.

Unfortunately this doesn't seem to work with nccl backend for now.

[kohya-ss]

This seems to be working fine on a Windows environment.
Testing on a Linux environment is welcome!

[xzuyn]

I've been using this on a single amd gpu on ubuntu for a bit, training qwen image. Haven't noticed any problems.