Update main readme to the point we can start the inference benchmark with ./nki-llama inference benchmark

[crizCraig]

No description provided.

[crizCraig]

@arm-diaz Hello from AGI House! We are preparing for a trainium competition here and I was running into the following on a trn1.2xlarge (as I don't have vCPU quota for a The trn1.32xlarge)

./nki-llama inference benchmark
📝 Logging to: /home/ubuntu/nki-llama/logs/nki-llama_20250628_001818.log
Running NKI benchmark evaluation...
💡 Running benchmark in evaluate_all mode
   This includes model compilation with NKI optimizations (10-30 min on first run).
   The compiled model will be cached for future use.
   Auto cache recovery is enabled by default.
   Using tmux is strongly recommended!
   Running: /home/ubuntu/nki-llama/src/inference/scripts/run-nki-benchmark.sh --mode evaluate_all
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
NKI Benchmark Evaluation
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Mode:              evaluate_all
Model:             llama-3-8b
Model Path:        /home/ubuntu/models/llama-3-8b
Compiled Path:     /home/ubuntu/traced_model/llama-3-8b
NKI Enabled:       true
Sequence Length:   512
TP Degree:         2
Auto Clear Cache:  true

✓ NKI optimizations enabled
📊 Benchmark logs will be saved to:
   /home/ubuntu/nki-llama/logs/benchmarks/20250628_001818

📊 Running benchmark in evaluate_all mode...
This mode evaluates all model configurations with NKI optimizations.
Executing evaluate_all benchmark...
python main.py --mode evaluate_all --model-path /home/ubuntu/models/llama-3-8b --compiled-model-path /home/ubuntu/traced_model/llama-3-8b --seq-len 512 --tp-degree 2 --enable-nki
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/parallel_layers/layers.py:16: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from .mappings import (
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/parallel_layers/layers.py:16: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from .mappings import (
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/parallel_layers/layers.py:16: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from .mappings import (
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/modules/moe/blockwise.py:79: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  component, error = import_nki(config)
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/modules/moe/blockwise.py:79: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  component, error = import_nki(config)
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/modules/moe/blockwise.py:79: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  component, error = import_nki(config)
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/modules/moe/blockwise.py:79: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  component, error = import_nki(config)
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/modules/moe/blockwise.py:79: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  component, error = import_nki(config)
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/models/dbrx/modeling_dbrx.py:38: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/models/dbrx/modeling_dbrx.py:38: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/models/dbrx/modeling_dbrx.py:38: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/utils/constants.py:1: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.models.dbrx.modeling_dbrx import NeuronDbrxForCausalLM
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/utils/constants.py:3: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from neuronx_distributed_inference.models.mixtral.modeling_mixtral import NeuronMixtralForCausalLM
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/models/mllama/modeling_mllama.py:72: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead.
  from .modeling_mllama_vision import NeuronMllamaVisionModel  # noqa: E402
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed_inference/utils/accuracy.py:30: UserWarning: Intel extension for pytorch not found. For faster CPU references install `intel-extension-for-pytorch`.
  warnings.warn(
root: NeuronConfig init: Unexpected keyword arguments: {'mode': 'evaluate_all', 'enable_nki': True, 'base_latency': 526.15, 'base_throughput': 134.61, 'model_path': '/home/ubuntu/models/llama-3-8b', 'compiled_model_path': '/home/ubuntu/traced_model/llama-3-8b', 'benchmark': False, 'divergence_difference_tol': 0.15, 'tol_map': '{None: (1e-5, 0.05), 1000: (1e-5, 0.03), 50: (1e-5, 0.03), 5: (1e-5, 0.03)}', 'prompts': ['I believe the meaning of life is'], 'top_k': 1, 'top_p': 1.0, 'temperature': 1.0, 'do_sample': True, 'dynamic': False, 'pad_token_id': 2, 'on_device_sampling': False, 'quantized_kernel_lower_bound': 1200.0}
Neuron: Saving the neuron_config to /home/ubuntu/traced_model/llama-3-8b/
Neuron: Generating HLOs for the following models: ['context_encoding_model', 'token_generation_model']
Loading configs...

Compiling and saving model...
[2025-06-28 00:18:42.744: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializing tensor model parallel with size 2
[2025-06-28 00:18:42.745: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1
[2025-06-28 00:18:42.745: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with size 1
[2025-06-28 00:18:42.745: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1
[2025-06-28 00:18:42.745: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializing world size to 2
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=<PG_Group_Logic.LOGIC1: (<function ascending_ring_PG_group at 0x79384739e440>, 'Ascending Ring PG Group')>
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0, 1]]
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0], [1]]
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0], [1]]
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0], [1]]
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0], [1]]
[2025-06-28 00:18:42.746: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0], [1]]
Neuron: Generating 3 hlos for key: context_encoding_model
Neuron: Started loading module context_encoding_model
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: TP degree (2) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!
Neuron: Finished loading module context_encoding_model in 0.05821490287780762 seconds
Neuron: generating HLO: context_encoding_model, input example shape = torch.Size([1, 128])
/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/neuronx_distributed/parallel_layers/layers.py:485: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with torch.cuda.amp.autocast(enabled=False):
/home/ubuntu/nki-llama/src/inference/scripts/run-nki-benchmark.sh: line 249:  5223 Killed                  $CMD 2>&1
      5224 Done                    | tee "${BENCHMARK_LOG_DIR}/benchmark.log"
      5225 Done                    | tee "$TEMP_LOG"
✗ evaluate_all benchmark failed!

I've reduced the resources in my .env to the following in case it was a memory error

MODEL_ID=meta-llama/Meta-Llama-3-8B
MODEL_NAME=llama-3-8b

# Inference Configuration
INFERENCE_PORT=8080
MAX_MODEL_LEN=512 # used by vllm- ensure it is the same as seq len
SEQ_LEN=512 #used by main.py

MAX_NUM_SEQS=2  # Reduced due to limited cores
TENSOR_PARALLEL_SIZE=2 # Match available NeuronCores

# Dataset Configuration
DATASET_NAME=databricks/databricks-dolly-15k

# Neuron Configuration
# Set according to the cores available on your instance
NEURON_RT_NUM_CORES=2

# Jupyter Configuration
JUPYTER_PORT=8888

Would very much appreciate any help.

Best,
Craig

[crizCraig]

I asked the above as there are no issues on this repo. However, this PR could also be merged.

[EmilyWebber]

Hi Craig, this project only supports higher TP degrees at the moment. It is not intended to work with TP=2 given the current setup.

[hezyin]

@EmilyWebber Hi Emily, One logistical wrinkle we’re seeing: a single trn1.32xlarge requires the user’s account to have a 128-vCPU quota. Submitting a Service Quotas request through the standard support channel often takes several days to approve. That delay is manageable for long-running online entrants, but it could be a real hurdle for new participants and especially for our one-day in-person hackathons, where teams need to spin up a node immediately.

Would love your thoughts on how we might streamline this (pre-approved event accounts, a smaller default instance, or an expedited quota path). Let me know what you think.

cc @crizCraig