Enable XLA_USE_BF16 for hf-llm tests (GoogleCloudPlatform#1046)

RissyRan · Feb 12, 2024 · 5dcf21f · 5dcf21f
1 parent edafd46
commit 5dcf21f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/tests/pytorch/nightly/hf-llm.libsonnet b/tests/pytorch/nightly/hf-llm.libsonnet
@@ -106,6 +106,7 @@ local utils = import 'templates/utils.libsonnet';
         export LD_LIBRARY_PATH=/usr/local/lib/
         export PT_XLA_DEBUG=0
         export USE_TORCH=ON
+        export XLA_USE_BF16=1
       |||,
       tpuVmExtraSetup: |||
         pip install --upgrade accelerate

diff --git a/tests/pytorch/r2.2/hf-llm.libsonnet b/tests/pytorch/r2.2/hf-llm.libsonnet
@@ -106,6 +106,7 @@ local utils = import 'templates/utils.libsonnet';
         export LD_LIBRARY_PATH=/usr/local/lib/
         export PT_XLA_DEBUG=0
         export USE_TORCH=ON
+        export XLA_USE_BF16=1
       |||,
       tpuVmExtraSetup: |||
         pip install --upgrade accelerate