Skip to content

Fsdp2 support for activation checkpointing #359

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.

Sign up for GitHub

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

Jump to bottom
Draft
wants to merge 64 commits into
base: main
Choose a base branch
from
Draft

Fsdp2 support for activation checkpointing #359

wants to merge 64 commits into from

Conversation

le1nux
Copy link
Member

@le1nux le1nux commented Apr 18, 2025

What does this PR do?

This PR ..

General Changes

  • ..

Breaking Changes

  • ..

Checklist before submitting final PR

  • My PR is minimal and addresses one issue in isolation
  • I have merged the latest version of the target branch into this feature branch
  • I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
  • I have run a sample config for model training
  • I have checked that all tests run through (python tests/tests.py)
  • I have updated the internal changelog (CHANGELOG_DEV.md)

le1nux added 30 commits April 18, 2025 19:01
@le1nux
refactor: towards fsdp2 support for activation checkpointing
725a218
@le1nux
feat: added test configs for AC
8362e0d
@le1nux
feat: added selective activation checkpointing configs
a546189
@le1nux
feat: added selective activation checkpointing tests
61070d1
@le1nux
feat: added selective activation checkpointing strategies
9c7573b
@le1nux
feat: wired up selective activation checkpointing
034662a
@le1nux
refactor: experiment id can now be passed to Main constructor to allo…
a9fe173 
…w for testing with a distributed environment
@le1nux
chore: added doc strings to the selective AC entry point
f977525
@le1nux
refactor: renamed typing and logging module to fix name shadowing issues
ba3c9e2
@le1nux
refactor: reneamed typing module to typing_utils
391d4eb
@le1nux
feat: added SAC benchmarking script
c6b6e47
@le1nux
chore: added activation checkpointing benchmark config
742ab46
@le1nux
chore: removed legacy code
641fad1
@le1nux
refactor: extracted save_list to parent in AC
f940b0e
@le1nux
refactor: moved Main from __main__ to main
c346f0f
@le1nux
feat: added batch generator util
90c4bbe
@le1nux
refactor: split experiment_id syncing into multiple utity functions
c87688a
@le1nux
feat: implemented grid search setup for profiling
ebaf13e
@le1nux
refactor: added OOM error handling in CudaEnv
bcc0e7b
@le1nux
feat: added torchrun script for distributed profiling
899401a
@le1nux
feat: added profiling README
87296d1
@le1nux
feat: added profiler implementation
d6dddc0
@le1nux
feat: drafted profile logs analyzer
011e41b
@le1nux
chore: minor renamings
e6da67f
@le1nux
refactor: making sure that each compil
2d4a9ad
@le1nux
feat: added torchrun launcher
d990420
@le1nux
refactor: wrapped up the profiler_starter
f4a5b48
@le1nux
feat: added activation checkpoint profiling example
b2aa30b
@le1nux
feat: setup forward pass profiling
94025af
@le1nux
feat: ops in selective op activation checkpointing are now configurable
3ea7cf5
Base automatically changed from fsdp2_min_integration to main May 28, 2025 10:59
le1nux added 23 commits May 31, 2025 16:10
@le1nux
chore: drafted scaling up tutorial
d37c352
@le1nux
chore: extended profiling analysis notebook
0388145
@le1nux
chore: updated gitignore
cc39d11
@le1nux
feat: added slurm / sbatch scripts for the scale up tutorial
085995f
@le1nux
refactor: rendezvous timeout is now configurable
f64c8dc
@le1nux
feat: added jupyter notebook for scalability analysis
b12cefb
@le1nux
refactor: removed legacy rdzv_timeout
3a86ebe
@le1nux
chore: removed hardcoded stuff from sbatch script
8dc0fc0
@le1nux
feat: more work on scaling notebook
80c9e6a
@le1nux
refactor: grid search is now part of config
4fb2adf
@le1nux
feat: run_array_jobs.sh now accepts number of ranks as CMD arguments
7f3a1bf
@le1nux
feat: added README.md for scaling tutorial
5e4d08b
@le1nux
refactor: renamed sac_params to ac_params
a7d94f5
@le1nux
feat: added pretraining benchmarking config
87667ae
@le1nux
refactor: improved usability of submission scripts
ffe7d8f
@le1nux
chore: improved error handling / logging
da66a3c
@le1nux
feat: added training run optimization part
226dc44
@le1nux
refactor: defined grid search for training run optimizatino
6a0fa8e
@le1nux
feat: added peak memory tracking
ec56a05
@le1nux
feat: added results subscriber that writes JSONL to disc
cbab52c
@le1nux
chore: Merge remote-tracking branch 'refs/remotes/origin/fsdp2_activa…
c1561e2 
…tion_checkpointing' into fsdp2_activation_checkpointing
@le1nux
chore: replaced timer with perf_counter
edbda8c
@le1nux
chore: resolved merge conflict
dd9a453
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@le1nux