qkv format error in GPT Eval when bs=1 and using fused rope kernel

[cuichenx]

Describe the bug

Opening a github issue for public tracking purposes. (internal #4489257)

megatron_gpt_eval.py throws the following error:

...
    core_attn_out = super().forward(
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/attention.py", line 2099, in forward
    qkv_layout, query_layer, key_layer, value_layer = _get_qkv_layout(
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/attention.py", line 1209, in _get_qkv_layout
    raise Exception("The provided qkv memory layout is not supported!")
Exception: The provided qkv memory layout is not supported!
This only happens when all of the following are true:

micro batch size = 1 (mbs = 2 works)
apply_rope_fusion = True
rope fusion is available (i.e. container has latest Apex)
running an inference workload – megatron_gpt_eval.py, or megatron_gpt_generate.py, or validation loop during training. (training with mbs = 1 and mbs = 2 both work)

Steps/Code to reproduce bug

run inference with any gpt-style model (e.g. llama2)

torchrun --nproc_per_node=1 /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
    gpt_model_file=<path to model> \
    inference.greedy=True \
    inference.add_BOS=True \
    trainer.devices=1 \
    trainer.num_nodes=1 \
    tensor_model_parallel_size=1 \
    pipeline_model_parallel_size=1 \
    prompts=["deep learning is"]

Environment overview (please complete the following information)

nvcr.io/nvidia/nemo:24.01.framework

[github-actions]

This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[cuichenx]

fixed in nemo 24.03 release / megatron core 0.6.0 release