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

[chrisxcai]

Proposal from @GD06 : https://fburl.com/gdoc/byomncbg

sudo rm -r /tmp/chriscai/xldumps/* & CUDA_LAUNCH_BLOCKING=1 PYTHONPATH=~/benchmark/fairscale_repos/fairscale/ CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 torchrun --master_port 1024 --nproc_per_node=8 train.py --dump_dir /tmp/chriscai/xldumps --model_parallel_size 2 --pipeline_parallel_size 1 --seq_len=1024 --gpu_check_level=-1 --steps=10 --log_all_steps=True --profile_freq=10 --dump_profile_traces=True --profile_with_stack=True --use_microbatching=True --model.n_layers=1 --reshard_after_forward=False --batch_size=128 --model.efficient_attn=flash --model.attn_bias_type=causal --model=small --mem_snapshot_stop_step 5

trace before change (SplitWithSizesBackward triggered every microbatch per FSDP module):
https://fburl.com/perfdoctor/qdt32ibh

trace with applied change (SplitWithSizesBackward triggered only in last microbatch per FSDP module):
https://fburl.com/perfdoctor/mpjsoarn

[GD06]

Could you please try to split this into three logger.info entries? I think the printing of new_param failed as the new_param we used to build partial function was gone...