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

Draft

chrisxcai wants to merge 27 commits into

ngoyal_changes_for_pp_fp8from

chriscai_ngoyal_changes_for_pp_fp8_v1

Commits on Apr 29, 2024

use torch.no_grad() to avoid calling cat() during FSDP backward except for last microbatch
chrisxcai
committed
d1102ce
remove logging
chrisxcai
committed
9a22628

Commits on Apr 30, 2024

logging
chrisxcai
committed
f787532

Commits on May 1, 2024

logging
chrisxcai
committed
3429f33

Commits on May 2, 2024

Commits on May 6, 2024

logging
chrisxcai
committed
901fb86

Commits on May 8, 2024

return grad in post_backward_hook()
chrisxcai
committed
ad40f24

Commits on May 9, 2024

correct param_index
chrisxcai
committed
14499fe
logging
chrisxcai
committed
ad7aa1f
add torch.testing.assert_allclose() to compare baseline and new grads
chrisxcai
committed
b835770
logging
chrisxcai
committed
d689f38

Commits on May 13, 2024

logging
chrisxcai
committed
e8df583

Commits on May 15, 2024

honor optimize_backward_concat flag
chrisxcai
committed
5926a79
documentation
chrisxcai
committed
5d08aa3
update documentation
chrisxcai
committed
c91cb72
update documentation
chrisxcai
committed
fd3f3fc
use grad instead of grad.data
chrisxcai
committed
7678503
clean up
chrisxcai
committed
c55a0d1
Added reshard hook for frozen params in backward

Andrew Gu
authored and
chrisxcai
committed
688b902
Avoid calling _free_fp16_param_shard() too early with PR 1159

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

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

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

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