xla
use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on param sharding dim and pinning memory
#3830

Merged

use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on param sharding dim and pinning memory #3830

JackCaoG merged 4 commits into pytorch:master from ronghanghu:xla_fsdp_patched_linear

ronghanghu force pushed from 2ddc6a04 to c905ba8f 4 years ago

ronghanghu marked this pull request as ready for review 4 years ago

JackCaoG requested a review from

JackCaoG 4 years ago

hjm-aws requested changes on 2022-08-05

ronghanghu force pushed from 2e12e877 to c18e4865 4 years ago

ronghanghu force pushed from f5dd6e46 to c15b5453 4 years ago

ronghanghu requested a review from

hjm-aws 4 years ago

hjm-aws approved these changes on 2022-08-15

use XLA patched linear in FSDP; expose options for padding in all_gat…

74335297

directly patch `nn.Linear`'s forward method instead of `torch.nn.func…

38de3e24

introduce `shard_param_on_dim_0` to handle potential compiler issues

690947d3

ronghanghu force pushed from c15b5453 to 690947d3 4 years ago

ronghanghu changed the title ~~use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on padding in all-gather and pinning memory~~ use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on param sharding dim and pinning memory 4 years ago

JackCaoG commented on 2022-08-17

check `full_param.dim() == 1` instead of catching exceptions in `_con…

34519228

JackCaoG approved these changes on 2022-08-17

JackCaoG merged 3f47abab into master 4 years ago

Reviewers

JackCaoG

hjm-aws

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

xla use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on param sharding dim and pinning memory #3830 Merged

use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on param sharding dim and pinning memory #3830

xla
use XLA patched linear in FSDP (fix #3811 and #3718) and expose options on param sharding dim and pinning memory
#3830

Merged