xla
add `auto_wrap_policy` into XLA FSDP for automatic wrapping
#4318

Merged

add `auto_wrap_policy` into XLA FSDP for automatic wrapping #4318

JackCaoG merged 3 commits into pytorch:master from ronghanghu:xla_fsdp_auto_wrap

ronghanghu force pushed from 3a8429e7 to fa34a562 3 years ago

ronghanghu force pushed from fa34a562 to b5f09206 3 years ago

ronghanghu force pushed from b6dd40f3 to bd01aea4 3 years ago

JackCaoG requested a review from

AlexWertheim 3 years ago

JackCaoG requested a review from

JackCaoG 3 years ago

JackCaoG requested a review from

miladm 3 years ago

JackCaoG added fsdp

JackCaoG commented on 2022-12-19

JackCaoG approved these changes on 2022-12-19

ronghanghu force pushed from a3ee5f1b to bd7196f2 3 years ago

ronghanghu force pushed from 58f632cb to 8da27b5a 3 years ago

add `auto_wrap_policy` into XLA FSDP for automatic wrapping

8e87bf66

address review comments

0002f45d

ronghanghu force pushed from 8da27b5a to 0002f45d 3 years ago

default to pin_layout=True in MNIST and ImageNet FSDP

a5b4da77

JackCaoG merged d7d04792 into master 3 years ago

Reviewers

JackCaoG

AlexWertheim

miladm

Assignees

No one assigned

Labels

fsdp

Milestone

No milestone