Update llama2-model.libsonnet

GoogleCloudPlatform · Aug 31, 2023 · 2ef583e · 2ef583e
1 parent 2ced52c
commit 2ef583e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/pytorch/nightly/llama2-model.libsonnet b/tests/pytorch/nightly/llama2-model.libsonnet
@@ -152,6 +152,8 @@ local utils = import 'templates/utils.libsonnet';
         export PROFILE_EPOCH=2
         export PROFILE_STEP=0
         export PROFILE_DURATION_MS=20000
+        export XLA_USE_SPMD=1
+        export PJRT_DEVICE=TPU
       |||,
       tpuVmExtraSetup: |||
         pip3 uninstall torch torch_xla torchvision libtpu-nightly -y