Nemotron lora finetune Error, The provided qkv memory layout is not supported
same here 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. can you guide about this issue? @ericharper  thank you in advance!
**Describe the bug**
Hi,
I'm trying to finetune Nemotron-3 Model using nvcr.io/nvidia/nemo:23.10 container with this [lora.ipynb notebook](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/lora.ipynb) and this exception occurs:

raise Exception("The provided qkv memory layout is not supported!")

**The full trace is like this:**

```
Traceback (most recent call last):
  File "/workspace/nemo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_tuning.py", line 77, in main
    trainer.fit(model)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/usr/local/lib/python3.10/dist-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 "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 980, in _run
    results = self._run_stage()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1021, in _run_stage
    self._run_sanity_check()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1050, in _run_sanity_check
    val_loop.run()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 376, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 293, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/ddp.py", line 338, in validation_step
    return self.model(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py", line 1521, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py", line 1357, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/overrides/base.py", line 102, in forward
    return self._forward_module.validation_step(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/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 "/usr/local/lib/python3.10/dist-packages/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 "/usr/local/lib/python3.10/dist-packages/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 963, in validation_step
    loss = self.fwd_bwd_step(dataloader_iter, batch_idx, True)
  File "/usr/local/lib/python3.10/dist-packages/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 "/usr/local/lib/python3.10/dist-packages/megatron/core/pipeline_parallel/schedules.py", line 327, in forward_backward_no_pipelining
    output_tensor = forward_step(
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/pipeline_parallel/schedules.py", line 183, in forward_step
    output_tensor, loss_func = forward_step_func(data_iterator, model)
  File "/usr/local/lib/python3.10/dist-packages/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 863, in fwd_output_and_loss_func
    output_tensor = model(**forward_args)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/models/gpt/gpt_model.py", line 166, in forward
    hidden_states = self.decoder(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/transformer/transformer_block.py", line 311, in forward
    hidden_states, context = layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/transformer/transformer_layer.py", line 153, in forward
    attention_output_with_bias = self.self_attention(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/transformer/attention.py", line 263, in forward
    core_attn_out = self.core_attention(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/transformer/custom_layers/transformer_engine.py", line 427, in forward
    return super().forward(
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/attention.py", line 1235, in forward
    qkv_layout = _get_qkv_layout(query_layer, key_layer, value_layer,
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/attention.py", line 464, in _get_qkv_layout
    raise Exception("The provided qkv memory layout is not supported!")
Exception: The provided qkv memory layout is not supported!
```


**and my training config file is this:**


```
  name: megatron_gpt_peft_${model.peft.peft_scheme}_tuning
    trainer:
      devices: 1
      accelerator: gpu
      num_nodes: 1
      precision: '32'
      logger: false
      enable_checkpointing: false
      use_distributed_sampler: false
      max_epochs: 3
      max_steps: 20000
      log_every_n_steps: 10
      val_check_interval: 200
      gradient_clip_val: 1.0
    exp_manager:
      explicit_log_dir: null
      exp_dir: null
      name: ${name}
      create_wandb_logger: false
      wandb_logger_kwargs:
        project: null
        name: null
      resume_if_exists: true
      resume_ignore_no_checkpoint: true
      create_checkpoint_callback: true
      checkpoint_callback_params:
        monitor: validation_${model.data.validation_ds.metric.name}
        save_top_k: 1
        mode: min
        save_nemo_on_train_end: true
        filename: ${name}--{${exp_manager.checkpoint_callback_params.monitor}:.3f}-{step}-{consumed_samples}
        model_parallel_size: ${model.tensor_model_parallel_size}
        always_save_nemo: false
        save_best_model: true
      create_early_stopping_callback: true
      early_stopping_callback_params:
        monitor: val_loss
        mode: min
        min_delta: 0.001
        patience: 10
        verbose: true
        strict: false
    model:
      seed: 1234
      tensor_model_parallel_size: 1
      pipeline_model_parallel_size: 1
      global_batch_size: 4
      micro_batch_size: 1
      restore_from_path: /project/models/Nemotron-3-8B-Chat-4k-RLHF.nemo
      resume_from_checkpoint: null
      save_nemo_on_validation_end: false
      sync_batch_comm: false
      megatron_amp_O2: false
      sequence_parallel: false
      activations_checkpoint_granularity: null
      activations_checkpoint_method: null
      activations_checkpoint_num_layers: null
      activations_checkpoint_layers_per_pipeline: null
      answer_only_loss: true
      gradient_as_bucket_view: false
      hidden_dropout: 0.0
      attention_dropout: 0.0
      ffn_dropout: 0.0
      peft:
        peft_scheme: ptuning
        restore_from_path: null
        adapter_tuning:
          type: parallel_adapter
          adapter_dim: 32
          adapter_dropout: 0.0
          norm_position: pre
          column_init_method: xavier
          row_init_method: zero
          norm_type: mixedfusedlayernorm
          layer_selection: null
          weight_tying: false
          position_embedding_strategy: null
        lora_tuning:
          adapter_dim: 32
          adapter_dropout: 0.0
          column_init_method: xavier
          row_init_method: zero
          layer_selection: null
          weight_tying: false
          position_embedding_strategy: null
        p_tuning:
          virtual_tokens: 15
          bottleneck_dim: 1024
          embedding_dim: 1024
          init_std: 0.023
        ia3_tuning:
          layer_selection: null
      data:
        train_ds:
          file_names:
          - /project/data/therapies.jsonl
          global_batch_size: ${model.global_batch_size}
          micro_batch_size: ${model.micro_batch_size}
          shuffle: true
          num_workers: 0
          memmap_workers: 2
          pin_memory: true
          max_seq_length: 2048
          min_seq_length: 1
          drop_last: true
          concat_sampling_probabilities:
          - 1.0
          label_key: output
          add_eos: true
          add_sep: false
          add_bos: false
          truncation_field: input
          index_mapping_dir: null
          prompt_template: '{input}
    
    
            {output}'
          truncation_method: right
        validation_ds:
          file_names:
          - /project/data/therapies.jsonl
          names: null
          global_batch_size: ${model.global_batch_size}
          micro_batch_size: ${model.micro_batch_size}
          shuffle: false
          num_workers: 0
          memmap_workers: ${model.data.train_ds.memmap_workers}
          pin_memory: true
          max_seq_length: 2048
          min_seq_length: 1
          drop_last: false
          label_key: ${model.data.train_ds.label_key}
          add_eos: ${model.data.train_ds.add_eos}
          add_sep: ${model.data.train_ds.add_sep}
          add_bos: ${model.data.train_ds.add_bos}
          write_predictions_to_file: false
          output_file_path_prefix: null
          truncation_field: ${model.data.train_ds.truncation_field}
          index_mapping_dir: null
          prompt_template: ${model.data.train_ds.prompt_template}
          tokens_to_generate: 32
          truncation_method: right
          metric:
            name: loss
            average: null
            num_classes: null
        test_ds:
          file_names: null
          names: null
          global_batch_size: ${model.global_batch_size}
          micro_batch_size: ${model.micro_batch_size}
          shuffle: false
          num_workers: 0
          memmap_workers: ${model.data.train_ds.memmap_workers}
          pin_memory: true
          max_seq_length: 2048
          min_seq_length: 1
          drop_last: false
          label_key: ${model.data.train_ds.label_key}
          add_eos: ${model.data.train_ds.add_eos}
          add_sep: ${model.data.train_ds.add_sep}
          add_bos: ${model.data.train_ds.add_bos}
          write_predictions_to_file: false
          output_file_path_prefix: null
          truncation_field: ${model.data.train_ds.truncation_field}
          index_mapping_dir: null
          prompt_template: ${model.data.train_ds.prompt_template}
          tokens_to_generate: 32
          truncation_method: right
          metric:
            name: loss
            average: null
            num_classes: null
      optim:
        name: fused_adam
        lr: 0.0001
        weight_decay: 0.01
        betas:
        - 0.9
        - 0.98
        sched:
          name: CosineAnnealing
          warmup_steps: 50
          min_lr: 0.0
          constant_steps: 0
          monitor: val_loss
          reduce_on_plateau: false
```
**Expected behavior**

finetuning the model properly.

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

 - Environment location: Docker - Nemo:23.10 container


**Additional context**
GPU: A100 - 80GB