Nemotron lora finetune Error, The provided qkv memory layout is not supported
State: closed | Pull request: no
Labels: name: bug

**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

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!