Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Draft
wants to merge 11 commits into
base: ngoyal_changes_for_pp_fp8
Choose a base branch
from

Conversation

chrisxcai
Copy link

@chrisxcai chrisxcai commented Apr 19, 2024

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

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 19, 2024
import functools

def _post_accumulation_hook(new_param_stop_grad, new_param):
# TODO: make it only call backward() only for last microbatch (within FSDP no_sync)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants