diffusion: FSDP weight loading speedup

[mickqian]

[mickqian]

@zyksir are there anything left to do here?

[zyksir]

Right Now Most of it is done. There are several things left when I try to run sgl-diffusion:

• In the case of cpu offload, lots of time is spent on creating pin memory. overlapping the creation of pin memory and memory copy might help, yet creating pin memory will still dominate the time. While in real world production I have to disable cpu offload since H2D takes too much time in the case when number of timesteps is small(after distillation).
• In the case of multiple gpus, since nvlink is much more faster then H2D copy. We can load part of the params and then broadcast it to other GPUs. The can save some time as well. but consider most of the model is small and we can already load it in 1 minutes, this can only gain another 20~30 seconds. And can be put at low priority