[WIP] Make FSDPv1 only perform cat() during last microbatch backward() within FlattenParamsWrapper#1178

Draft

chrisxcai wants to merge 11 commits intongoyal_changes_for_pp_fp8from chriscai_ngoyal_changes_for_pp_fp8

+105-5

Commits on Apr 19, 2024

Commits on Apr 20, 2024

try adding new_params into wrapper class field
chrisxcai
committed

Commits on Apr 23, 2024

use cloned tensor to accumulate grad
chrisxcai
committed

Commits on Apr 24, 2024

only trigger new_param.backward() in last microbatch
chrisxcai
committed
only trigger new_param.backward() in last microbatch
chrisxcai
committed
only run backward() for all new_params once per FSDP when the grad of last new_param_stop_grad is ready per, to reduce the call to SplitWithSizesBackward etc
chrisxcai
committed

Commits on Apr 25, 2024

Commits on Apr 27, 2024

remove .clone()
chrisxcai
committed