I followed NVIDIA docs hub, and I got an error, `Exception: The provided qkv memory layout is not supported!`
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. This issue was closed because it has been inactive for 7 days since being marked as stale.
**Describe the bug**

I'm trying to train with the nemo model following this [guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/playbooks/llama2sft.html).
I've converted llama2-7b-hf to nemo model.

When I execute training script, I got this message `Exception: The provided qkv memory layout is not supported!`

**Steps/Code to reproduce bug**

I followed the guide from the beginning.
In the guide, followed [Step 4: Run the SFT finetuning script](https://docs.nvidia.com/nemo-framework/user-guide/latest/playbooks/llama2sft.html#step-4-run-the-sft-finetuning-script), only changing `megatron_amp_O2` from `True` to `False`.
(Without changing, I got FileNotFoundError regarding `.../model_weights/model.module. ... / shard_0_32.pt`)
And I got an error as below.
```
  File "/home/projects/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_sft.py", line 226, in main
    trainer.fit(model)
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 980, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 1021, in _run_stage
    self._run_sanity_check()
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 1050, in _run_sanity_check
    val_loop.run()
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 376, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py", line 293, in _call_strategy_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/strategies/ddp.py", line 338, in validation_step
    return self.model(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1519, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1355, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/pytorch_lightning/overrides/base.py", line 102, in forward
    return self._forward_module.validation_step(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/projects/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 387, in validation_step
    return self.inference_step(dataloader_iter, batch_idx, 'validation', dataloader_idx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/projects/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 403, in inference_step
    loss = super().validation_step(itertools.chain([batch]), batch_idx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/projects/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 960, in validation_step
    loss = self.fwd_bwd_step(dataloader_iter, batch_idx, True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/projects/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 347, in fwd_bwd_step
    losses_reduced_per_micro_batch = fwd_bwd_function(
                                     ^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/pipeline_parallel/schedules.py", line 329, in forward_backward_no_pipelining
    output_tensor = forward_step(
                    ^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/pipeline_parallel/schedules.py", line 185, in forward_step
    output_tensor, loss_func = forward_step_func(data_iterator, model)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/projects/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 860, in fwd_output_and_loss_func
    output_tensor = model(**forward_args)
                    ^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/models/gpt/gpt_model.py", line 172, in forward
    hidden_states = self.decoder(
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/transformer/transformer_block.py", line 250, in forward
    hidden_states = layer(
                    ^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/transformer/transformer_layer.py", line 124, in forward
    attention_output_with_bias = self.self_attention(
                                 ^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/transformer/attention.py", line 233, in forward
    core_attn_out = self._checkpointed_attention_forward(query, key, value, attention_mask)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/transformer/attention.py", line 81, in _checkpointed_attention_forward
    hidden_states = tensor_parallel.checkpoint(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/tensor_parallel/random.py", line 249, in checkpoint
    return CheckpointFunction.apply(function, distribute_saved_activations, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/tensor_parallel/random.py", line 191, in forward
    outputs = run_function(*args)
              ^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/megatron/core/transformer/attention.py", line 78, in custom_forward
    output_ = self.dot_product_attention(query, key, value, attention_mask)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/transformer_engine/pytorch/attention.py", line 2085, in forward
    qkv_layout = _get_qkv_layout(query_layer, key_layer, value_layer,
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mambaforge/envs/nemo/lib/python3.11/site-packages/transformer_engine/pytorch/attention.py", line 1195, in _get_qkv_layout
    raise Exception("The provided qkv memory layout is not supported!")
Exception: The provided qkv memory layout is not supported!

```


**Expected behavior**

Train a model properly.

**Environment overview (please complete the following information)**

 - Environment location: Bare-metal
 - Method of NeMo install: 
 ```
$ pip install torch==2.1.0 torchvision --index-url https://download.pytorch.org/whl/cu118
$ git clone https://github.com/NVIDIA/apex
$ pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./apex
$ pip install -r requirements.txt  # except torch
$ pip install -r requirements_lightning 
$ pip install -r requirements_common # except youtokentome, run manually by conda install youtokentome
$ pip install -r requirements_nlp.txt
 ```

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version: ubuntu 20.04
- PyTorch version: 2.1.0
- Python version: 3.11

**Additional context**
 
GPU: A100 1node 8gpus
