Add test_llama bf16 portable config to CI (pytorch#5472)

Summary:
Pull Request resolved: pytorch#5472

We should ensure that bfloat16 works.

bypass-github-export-checks

Reviewed By: mergennachin

Differential Revision: D62963924

fbshipit-source-id: 6fa2abaad083311ea1f1a280e6d27f1f0d59e367

Author: swolchok

.ci/scripts/test_llama.sh

@@ -11,7 +11,7 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
 MODEL_NAME=$1 # stories110M
 BUILD_TOOL=$2 # buck2 or cmake
-DTYPE=$3 # fp16 or fp32
+DTYPE=$3 # fp16, bf16, or fp32
 MODE=${4:-"xnnpack+custom"} # portable or xnnpack+custom or xnnpack+custom+qe
 UPLOAD_DIR=${5:-}
 if [[ $# -lt 4 ]]; then # Assuming 4 mandatory args
@@ -29,7 +29,7 @@ if [[ -z "${BUILD_TOOL:-}" ]]; then
 fi
 
 if [[ -z "${DTYPE:-}" ]]; then
-  echo "Missing dtype, choose fp16 or fp32, exiting..."
+  echo "Missing dtype, choose fp16, bf16, or fp32, exiting..."
   exit 1
 fi
 
@@ -174,6 +174,8 @@ fi
 EXPORTED_MODEL_NAME="llama2"
 if [[ "${DTYPE}" == "fp16" ]]; then
   EXPORTED_MODEL_NAME="${EXPORTED_MODEL_NAME}_h"
+elif [[ "${DTYPE}" == "bf16" ]]; then
+  EXPORTED_MODEL_NAME="${EXPORTED_MODEL_NAME}_bf"
 elif [[ "${DTYPE}" == "fp32" ]]; then
   :
 else

.github/workflows/pull.yml

@@ -91,6 +91,13 @@ jobs:
         dtype: [fp32]
         build-tool: [buck2, cmake]
         mode: [portable, xnnpack+custom, xnnpack+custom+qe]
+      include:
+        - dtype: bf16
+          build-tool: cmake
+          mode: portable
+        - dtype: bf16
+          build-tool: buck2
+          mode: portable
       fail-fast: false
     with:
       runner: linux.2xlarge

.github/workflows/trunk.yml

@@ -224,6 +224,9 @@ jobs:
       matrix:
         dtype: [fp32]
         mode: [portable, xnnpack+kv+custom, mps, coreml]
+        include:
+          - dtype: bf16
+            mode: portable
       fail-fast: false
     with:
       runner: macos-m1-stable