xla
Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP
#3588

Merged

Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP #3588

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

allow specifying computation dtype other than float32 in XLA FSDP; al…

ebbb0696

minor update and lint

9c703ffd

hjm-aws commented on 2022-05-20

hjm-aws approved these changes on 2022-05-20

JackCaoG requested a review from

miladm 4 years ago

minor fix on model `extra_repr`

15353113

ronghanghu force pushed from 0b5e8864 to 15353113 4 years ago

ronghanghu changed the title ~~[WIP] Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP~~ Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP 4 years ago

JackCaoG approved these changes on 2022-05-26

JackCaoG merged f008ec88 into master 3 years ago

ronghanghu deleted the xla_fsdp_compute_dtype branch 3 years ago

Reviewers

JackCaoG

hjm-aws

miladm

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

xla Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP #3588 Merged

Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP #3588

xla
Allowing `float16` and `bfloat16` computation type and sharding groups in XLA FSDP
#3588

Merged