[FSDPv1] Only perform cat() during last microbatch backward() within FlattenParamsWrapper #1180

Draft

chrisxcai wants to merge 27 commits intongoyal_changes_for_pp_fp8from chriscai_ngoyal_changes_for_pp_fp8_v1

+346-25

Commits on Apr 29, 2024

Commits on Apr 30, 2024

logging
chrisxcai
committed

Commits on May 1, 2024

logging
chrisxcai
committed

Commits on May 2, 2024

use new field to accumulate per-parameter grads in fp32 and copy into flatten_parameter.unsharded_main_grad in last microbatch backward()
chrisxcai
committed
clean up accumulated fp32 grads between data batches
chrisxcai
committed
logging
chrisxcai
committed

Commits on May 6, 2024

logging
chrisxcai
committed

Commits on May 8, 2024

return grad in post_backward_hook()
chrisxcai
committed

Commits on May 9, 2024

Commits on May 13, 2024

logging
chrisxcai
committed

Commits on May 15, 2024

honor optimize_backward_concat flag
chrisxcai
committed
documentation
chrisxcai
committed
update documentation
chrisxcai
committed
update documentation
chrisxcai
committed
use grad instead of grad.data
chrisxcai
committed
clean up
chrisxcai
committed
Added reshard hook for frozen params in backward

awgu
authored and
chrisxcai
committed
Avoid calling _free_fp16_param_shard() too early with PR 1159

jiecaoyu
authored and
chrisxcai
committed
Added requires_grad check for params_with_grad method (#1171 )

whbldhwj
authored and
chrisxcai
committed
Changed to only run reshard hook if all gradients computed (#1166 )

awgu
authored and
chrisxcai
committed
Add cast input argument (#1175 )

whbldhwj
authored and
chrisxcai
committed
honor optimize_backward_concat flag
chrisxcai
committed
use grad instead of grad.data
chrisxcai
committed