Add Automicrobatching for Non-Powers-of-2 + Fixes to FSDP deadlocks using Adaptive Sync Hooks - #3503

Draft

JackZ-db wants to merge 23 commits into

mosaicml:mainfrom

JackZ-db:jz/auto_non_powers_of_2

Commits on Jul 30, 2024

add automicrobatching for non-powers-of-2 + adaptive sync hooks
JackZ-db
committed
bbb9a66
include auto helpers in _all_
JackZ-db
committed
4228889
fix circular imports
JackZ-db
committed
a537c4c
remove circular import
JackZ-db
committed
ff806d1
remove import state
JackZ-db
committed
4025274
dist
JackZ-db
committed
9ad1719
fix imports
JackZ-db
committed
896b999
import defaultdict
JackZ-db
committed
7146f23
log for hook on off
JackZ-db
committed
cd2fe9f
fixed hook readd bug
JackZ-db
committed
b1b16cd
rename hooks to fsdp hooks, will only trigger if fsdp
JackZ-db
committed
476e028
only invoke hook logic if fsdp enabled
JackZ-db
committed
0693364
typo
JackZ-db
committed
df404ac
fix seq length warmup
JackZ-db
committed
0b6d6ce
only patch flat param handle unshard if > 2.3
JackZ-db
committed
153c413
fix version comparison
JackZ-db
committed
d98926d
mark unit test
JackZ-db
committed
3e82ef6

Commits on Jul 31, 2024

remove device mark
JackZ-db
committed
a09b844
filter user warnigns out
JackZ-db
committed
33840c5
fix
JackZ-db
committed
e87c9f6
dist sampler
JackZ-db
committed
86c32a0
ignore runtime warning
JackZ-db
committed
a193b76
only drop hooks after 3 consecutive successes with this microbatch size
JackZ-db
committed
0b6a30d