docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:c39107c4160e749b7c4bac18862c6c1b6d56e1aa60644a4fe323e315ffba0a0b" /work/xla-tools-dir/hlo_runner_main --hlo_file=/work/xla_hlo_before_optimizations.txt --device_type=gpu --num_repeats=50 --input_format=text --num_replicas=1 --num_partitions=1 --logtostderr
2023-08-04 19:15:21.721351: I xla/service/service.cc:168] XLA service 0x5640370dddd0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-08-04 19:15:21.721415: I xla/service/service.cc:176]   StreamExecutor device (0): NVIDIA A100-SXM4-40GB, Compute Capability 8.0
2023-08-04 19:15:21.721767: I xla/pjrt/gpu/se_gpu_pjrt_client.cc:633] Using BFC allocator.
2023-08-04 19:15:21.721826: I xla/pjrt/gpu/gpu_helpers.cc:105] XLA backend allocating 31753961472 bytes on device 0 for BFCAllocator.
2023-08-04 19:15:31.158463: I xla/stream_executor/cuda/cuda_dnn.cc:442] Loaded cuDNN version 8900
2023-08-04 19:15:34.067278: I xla/stream_executor/gpu/asm_compiler.cc:328] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_295', 996 bytes spill stores, 1108 bytes spill loads

2023-08-04 19:15:36.668819: W xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: Unexpected GEMM dtype: f32 f32 f16
2023-08-04 19:15:36.699421: F xla/tools/multihost_hlo_runner/hlo_runner_main.cc:121] Non-OK-status: xla::FunctionalHloRunner::LoadAndRunAndDump( *client.value(), preproc_options, raw_compile_options, running_options, {hlo_file}, input_format, dump_output_literal_to, task_id) status: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.gemm' failed: Unexpected GEMM dtype: f32 f32 f16; current tracing scope: custom-call; current profiling annotation: XlaModule:#hlo_module=extracted,program_id=131#.

@tf.function(jit_compile=True)
def forward(self, inputs: Tuple[Any, ...]) -> Tuple[Any, ...]:
  input_ids, attention_mask = inputs
  output = self.model(input_ids, attention_mask,
                      training=False).last_hidden_state
  return (output,)

@tf.function(jit_compile=True)
def forward_sm(self, input_ids, attention_mask):
  return self.model(input_ids, attention_mask,
                    training=False).last_hidden_state

# We only need to have one definition for each input data source. And it only needs to tell where to get the source raw data.
APPLE_IMAGE = ModelTestData(
  name="APPLE_IMAGE",
  source_url="url to download the image"
)
ORANGE_IMAGE = ModelTestData(
  name="ORANGE_IMAGE",
  source_url="url to download the image"
)

# Generate the models with different batch sizes, as what we do today.
# (we can actually expand the template to support different data types)
MODEL_RESNET50_FP32_BATCHES = generate_batch_models(...)
MODEL_RESNET50_FP16_BATCHES = generate_batch_models(...)

# Generate the combinations of all models with the apple image input.
BENCHMARK_LIST = generate_benchmark_cases(
  models=MODEL_RESNET50_FP32_BATCHES + MODEL_RESNET50_FP16_BATCHES,
  input_data=APPLE_IMAGE,
  output_verification="the tolerance to verify the output"
)

 ~/benchmark > git clone https://github.com/openxla/iree-comparative-benchmark                                                                                     
Cloning into 'iree-comparative-benchmark'...                                                                                                                                                                                                                                                                                                            
remote: Enumerating objects: 2003, done.                                                                                                                                                                                                                                                                                                                
remote: Counting objects: 100% (1051/1051), done.                                                                                                                           
remote: Compressing objects: 100% (440/440), done.                   
remote: Total 2003 (delta 770), reused 683 (delta 585), pack-reused 952                                                                                                                                                                                                                                                                                 
Receiving objects: 100% (2003/2003), 1.01 MiB | 3.58 MiB/s, done.                                                                                                                                                                                                                                                                                       
Resolving deltas: 100% (1148/1148), done.                                                                                                                                                                                                                                                                                                               
 ~/benchmark > cd iree-comparative-benchmark/comparative_benchmark/jax                                                                                                                                                                                                                                                                         
 ~/benchmark/iree-comparative-benchmark/comparative_benchmark/jax > ./setup_venv.sh      
...
 ~/benchmark/iree-comparative-benchmark/comparative_benchmark/jax > source jax-benchmarks.venv/bin/activate                                                        
(jax-benchmarks.venv)  ~/benchmark/iree-comparative-benchmark/comparative_benchmark/jax > ./run_benchmarks.py  -o test -device host-cpu -name models/BERT_BASE_FP32_JAX_I32_SEQLEN32/inputs/INPUT_DATA_MODEL_DEFAULT


--- models/BERT_BASE_FP32_JAX_I32_SEQLEN32/inputs/INPUT_DATA_MODEL_DEFAULT ---
/usr/lib/python3.9/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  self.pid = os.fork()
Some weights of FlaxBertModel were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: {('pooler', 'dense', 'bias'), ('pooler', 'dense', 'kernel')}
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
(jax-benchmarks.venv) ~/benchmark/iree-comparative-benchmark/comparative_benchmark/jax > cat test
{"benchmarks": [{"definition": {"benchmark_name": "models/BERT_BASE_FP32_JAX_I32_SEQLEN32/inputs/INPUT_DATA_MODEL_DEFAULT", "framework": "ModelFrameworkType.JAX", "data_type": "fp32", "batch_size": 1, "compiler": "xla", "device": "host-cpu", "tags": ["transformer-encoder", "bert", "seqlen-32"]}, "metrics": {"framework_level": {"error": "['Output 0 exceeds tolerance. Max diff: 8.474491119384766, atol: 0.5, rtol: 0']"}}}]}

"error": "Failed to import transformers.models.gpt2.modeling_flax_gpt2 because of the following error (look up to see its traceback):\nmodule 'jax.numpy' has no attribute 'DeviceArray'"

--- models/BERT_LARGE_FP16_PT_384XI32_BATCH16/inputs/INPUT_DATA_BERT_LARGE_FP16_PT_384XI32_BATCH16/expected_outputs/OUTPUT_DATA_BERT_LARGE_FP16_PT_384X1024XF16_BATCH16/target_devices/a2-highgpu-1g ---
Some weights of the model checkpoint at bert-large-uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.bias', 'cls.predictions.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.bias']
- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Failed to benchmark model BERT_LARGE_FP16_PT_384XI32_BATCH16.
{
  "definition": {
    "benchmark_id": "models/47cb0d3a-5eb7-41c7-9d7c-97aae7023ecf-MODEL_BERT_LARGE-fp16-PT-384xi32-batch16/inputs/2bbb87cf-a910-4262-a9d8-ceff295f1c24-fp16-batch16/expected_outputs/cd625ba7-fc70-4a87-92eb-5acab0c77beb-fp16-batch16/target_devices/78c56b95-2d7d-44b5-b5fd-8e47aa961108",
    "benchmark_name": "models/BERT_LARGE_FP16_PT_384XI32_BATCH16/inputs/INPUT_DATA_BERT_LARGE_FP16_PT_384XI32_BATCH16/expected_outputs/OUTPUT_DATA_BERT_LARGE_FP16_PT_384X1024XF16_BATCH16/target_devices/a2-highgpu-1g",
    "framework": "ModelFrameworkType.PYTORCH",
    "data_type": "fp16",
    "batch_size": 16,
    "inputs": [
      "16x384xi32",
      "16x384xi32"
    ],
    "outputs": [
      "16x384x1024xf16"
    ],
    "compiler": "xla",
    "device": "a2-highgpu-1g",
    "tags": [
      "transformer-encoder",
      "bert",
      "batch-16"
    ]
  },
  "metrics": {
    "framework_level": {
      "error": "quantile() input tensor must be either float or double dtype"
    }
  }
}