Adding prefetching of first shards to train script when fsdp enabled

[chelsea0x3b]

If model is sharded calling .unshard() will prefetch the first shard. I placed this before the data loader & other preprocessing so it should overlap.

Sources:

• https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.FSDPModule.unshard

• https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html#forward-backward-with-prefetching

Issuing 1st all-gather earlier: Implicit prefetching happens at the time of calling model(x). The 1st all-gather gets exposed. We can call model.unshard() explicitly earlier to issue 1st all-gather earlier

[tianyu-l]

Curious how much benefit do we get from this? Could you show some traces?

I have the following concerns:

1. When we don't do logging (where getting loss to CPU incurs d2h sync), GPU is ahead of CPU and can already overlap with dataloading, so this is not saving anything.
2. In torchtitan there're other FSDP implementation, in in general we should avoid FSDP2-only code. There's a way around by testing if model_parts[0] is an FSDPModule, but it's not so clean.

Let's see if the benefit justifies the complexity. WDYT?

[chelsea0x3b]

Ah yeah turns out there is little benefit, you're right! I think i had logging higher when I was noticing an improvement. Will close, ty!