fix llama2-70b bug, add document (#18398)

1. fix dist setting bug for LLaMA2-70b distributed convert and benchmark
2. Add instruction in README for how to benchmark LLaMA2-70b distribute
inference

onnxruntime/python/tools/transformers/models/llama/README.md

@@ -135,7 +135,7 @@ $ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output l
 $ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp16 --precision fp16 --execution_provider cuda --use_gqa
 ```
 
-Note: GroupQueryAttention currently runs on Linux for FP16 CUDA and INT4 CUDA models, and it can provide faster inference than MultiHeadAttention, especially for large sequence lengths (e.g. 1024 or larger). For the best performance, you should pre-allocate the KV cache buffers to have size `(batch_size, num_heads, max_sequence_length, head_size)` so that the past KV and present KV caches share the same memory. You also need to bind them with ONNX Runtime's [IO binding](https://onnxruntime.ai/docs/api/python/api_summary.html#iobinding).
+Note: GroupQueryAttention currently works with the FP16 CUDA and INT4 CUDA models, and it can provide faster inference than MultiHeadAttention, especially for large sequence lengths (e.g. 1024 or larger). For the best performance, you should pre-allocate the KV cache buffers to have size `(batch_size, num_heads, max_sequence_length, head_size)` so that the past KV and present KV caches share the same memory. You also need to bind them with ONNX Runtime's [IO binding](https://onnxruntime.ai/docs/api/python/api_summary.html#iobinding).
 
 Here is an example of how you can bind directly to `torch.tensor` objects:
 ```
@@ -340,6 +340,18 @@ CUDA_VISIBLE_DEVICES=4 python3 -m models.llama.benchmark \
     --device cuda
 ```
 
+9. ONNX Runtime, FP16, convert_to_onnx, LLaMA-2 70B shard to 4 GPUs
+```
+CUDA_VISIBLE_DEVICES=4,5,6,7 bash benchmark_70b_model.sh 4 \
+    --benchmark-type ort-convert-to-onnx \
+    --ort-model-path ./llama2-70b-dis/rank_{}_Llama-2-70b-hf_decoder_merged_model_fp16.onnx \
+    --model-name meta-llama/Llama-2-70b-hf \
+    --precision fp16 \
+    --device cuda \
+    --warmup-runs 5 \
+    --num-runs 100
+```
+
 You can profile a variant by adding the `--profile` flag and providing one batch size and sequence length combination.
 
 ### Benchmark All

onnxruntime/python/tools/transformers/models/llama/dist_settings.py

@@ -2,8 +2,6 @@ import os
 
 import torch.distributed as dist
 
-comm = None
-
 
 def init_dist():
     if "LOCAL_RANK" in os.environ:
@@ -13,10 +11,6 @@ def init_dist():
 
         dist.init_process_group("nccl", init_method="tcp://127.0.0.1:7645", world_size=world_size, rank=rank)
     elif "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ:
-        from mpi4py import MPI
-
-        comm = MPI.COMM_WORLD  # noqa: F841
-
         int(os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", 0))
         rank = int(os.environ.get("OMPI_COMM_WORLD_RANK", 0))
         world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", 1))
@@ -27,15 +21,28 @@ def init_dist():
         pass
 
 
+def _get_comm():
+    try:
+        from mpi4py import MPI
+
+        comm = MPI.COMM_WORLD
+        return comm
+    except ImportError:
+        return None
+
+
 def get_rank():
+    comm = _get_comm()
     return comm.Get_rank() if comm is not None else 0
 
 
 def get_size():
+    comm = _get_comm()
     return comm.Get_size() if comm is not None else 1
 
 
 def barrier():
+    comm = _get_comm()
     if comm is not None:
         comm.Barrier()