llama3 icon indicating copy to clipboard operation
llama3 copied to clipboard

Question RE FSDP usage in the paper

Open tsengalb99 opened this issue 6 months ago • 0 comments

Section 3.3.2 in the Llama 3.1 paper https://arxiv.org/pdf/2407.21783 says that Llama 3.1 was trained with FSDP on the parameters, gradients, and optimizer states. However, it also says that the parameters were not re-sharded for the backward pass to avoid another all gather reduction. Doesn't this mean that each DP rank needs to have enough memory to hold the entire model's parameters? If so, then why bother sharding parameters for the forward pass if you need enough memory to hold the whole model for the backward pass?

Screenshot 2024-08-01 at 10 50 08 AM

tsengalb99 avatar Aug 01 '24 17:08 tsengalb99