xla
Allow running with bfloat16 on XLA:GPU autocast
#5598
Merged

Allow running with bfloat16 on XLA:GPU autocast #5598

yeounoh merged 19 commits into master from spmd_amp_gpu
yeounoh
yeounoh yeounoh added performance
yeounoh yeounoh requested a review from JackCaoG JackCaoG 2 years ago
yeounoh yeounoh assigned yeounoh yeounoh 2 years ago
yeounoh yeounoh force pushed from cbddc5d0 to 13834934 2 years ago
baoleai
baoleai commented on 2023-09-18
JackCaoG
JackCaoG commented on 2023-09-18
JackCaoG
JackCaoG commented on 2023-09-18
JackCaoG
JackCaoG commented on 2023-09-18
JackCaoG
JackCaoG approved these changes on 2023-09-18
yeounoh Enable autocast for XLA:GPU
66e1a015
yeounoh linter fix
83b5c712
yeounoh XLA autocast test for GPU and TPU
1311abec
yeounoh linter fix
1e90fe88
yeounoh Ensure that xla autocast is properly enabled for GPU and does not cra…
623a26b3
yeounoh linter fix
7abae8bb
yeounoh Add tests
2308541e
yeounoh Support bf16
77cf2b00
yeounoh linter fix
8cf78a0a
yeounoh exclude unsupported test cases
f92fd0ab
yeounoh Support bfloat16 for XLA:GPU
5275546a
yeounoh Revert "Support bfloat16 for XLA:GPU"
d7f0c06b
yeounoh Enable bfloat16 for XLA:GPU on cuda backend
35fdecaa
yeounoh Update test for bfloat16
c464e5a3
yeounoh linter fix
b5628b13
yeounoh rebase master
0163707a
yeounoh fix bug in autocast_mode.py
9637bd3a
yeounoh Rename _cuda_bfloat16 to _xla_bfloat16 since it is set when xla backe…
c24c1da7
yeounoh add comments
83580957
yeounoh yeounoh force pushed from 5c317bbb to 83580957 2 years ago
yeounoh
yeounoh yeounoh merged b214c104 into master 2 years ago

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
Labels
Milestone