Improve Accelerate's FSDP2 preparation for language models whose embeddings and final normalization…
Improve Accelerate's FSDP2 preparation for language models whose embeddings and final normalization…: a task in HF ML Tasksmith (Harbor dataset). For untied weights, the input embedding should be an independent unit using the configured reshard policy. The final normalization and output embedding…
The task
For untied weights, the input embedding should be an independent unit using the configured reshard policy. The final normalization and output embedding should share a tail unit that stays unsharded after forward. For tied weights, the shared input embedding belongs to that tail unit; avoid double wrapping it.…
Part of FineEnvs/HF_ML_Tasksmith.