Nemo run pretrain raised Exception: No dot product attention support for the provided inputs
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**

When trying to run `nemo` with `gemma_2b`, `Exception: No dot product attention support for the provided inputs!` was returned.
```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/local/bin/nemo", line 8, in <module>
[rank0]:     sys.exit(app())
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/main.py", line 326, in __call__
[rank0]:     raise e
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/main.py", line 309, in __call__
[rank0]:     return get_command(self)(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1157, in __call__
[rank0]:     return self.main(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/core.py", line 723, in main
[rank0]:     return _main(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/core.py", line 193, in _main
[rank0]:     rv = self.invoke(ctx)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1688, in invoke
[rank0]:     return _process_result(sub_ctx.command.invoke(sub_ctx))
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1688, in invoke
[rank0]:     return _process_result(sub_ctx.command.invoke(sub_ctx))
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1434, in invoke
[rank0]:     return ctx.invoke(self.callback, **ctx.params)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 783, in invoke
[rank0]:     return __callback(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/main.py", line 692, in wrapper
[rank0]:     return callback(**use_params)
[rank0]:   File "/opt/NeMo-Run/src/nemo_run/cli/api.py", line 793, in command
[rank0]:     self.cli_execute(fn, ctx.args, type)
[rank0]:   File "/opt/NeMo-Run/src/nemo_run/cli/api.py", line 845, in cli_execute
[rank0]:     self._execute_task(fn, filtered_args)
[rank0]:   File "/opt/NeMo-Run/src/nemo_run/cli/api.py", line 895, in _execute_task
[rank0]:     run_task()
[rank0]:   File "/opt/NeMo-Run/src/nemo_run/cli/api.py", line 874, in run_task
[rank0]:     run.run(
[rank0]:   File "/opt/NeMo-Run/src/nemo_run/run/api.py", line 65, in run
[rank0]:     direct_run_fn(fn_or_script, dryrun=dryrun)
[rank0]:   File "/opt/NeMo-Run/src/nemo_run/run/task.py", line 77, in direct_run_fn
[rank0]:     built_fn()
[rank0]:   File "/opt/NeMo/nemo/collections/llm/api.py", line 134, in pretrain
[rank0]:     return train(
[rank0]:   File "/opt/NeMo/nemo/collections/llm/api.py", line 92, in train
[rank0]:     trainer.fit(model, data)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 538, in fit
[rank0]:     call._call_and_handle_interrupt(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 574, in _fit_impl
[rank0]:     self._run(model, ckpt_path=ckpt_path)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 981, in _run
[rank0]:     results = self._run_stage()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1023, in _run_stage
[rank0]:     self._run_sanity_check()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1052, in _run_sanity_check
[rank0]:     val_loop.run()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/utilities.py", line 178, in _decorator
[rank0]:     return loop_run(self, *args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 135, in run
[rank0]:     self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 396, in _evaluation_step
[rank0]:     output = call._call_strategy_hook(trainer, hook_name, *step_args)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 319, in _call_strategy_hook
[rank0]:     output = fn(*args, **kwargs)
[rank0]:   File "/opt/NeMo/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 571, in validation_step
[rank0]:     out = self.model.validation_step(dataloader_iter, *args, **kwargs)
[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 342, in validation_step
[rank0]:     return self._step(
[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 429, in _step
[rank0]:     return self.forward(
[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 279, in forward
[rank0]:     microbatch_outputs = step()
[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 1126, in __call__
[rank0]:     return self.forward_backward_func(
[rank0]:   File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 450, in forward_backward_no_pipelining
[rank0]:     output_tensor, num_tokens = forward_step(
[rank0]:   File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 273, in forward_step
[rank0]:     output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 492, in wrapped_forward_step_func
[rank0]:     output_tensor = _forward_step(model, batch)
[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 742, in wrapped
[rank0]:     return attr(*args)
[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 311, in validation_step
[rank0]:     return self.forward_step(batch)
[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 302, in forward_step
[rank0]:     return self.config.forward_step_fn(self, batch)
[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 101, in gpt_forward_step
[rank0]:     return model(**forward_args)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1561, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 286, in forward
[rank0]:     output_tensor = self.module(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1561, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/distributed/distributed_data_parallel.py", line 305, in forward
[rank0]:     return self.module(*inputs, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1602, in _call_impl
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/transformer/module.py", line 178, in forward
[rank0]:     outputs = self.module(*inputs, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1602, in _call_impl
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/models/gpt/gpt_model.py", line 226, in forward
[rank0]:     hidden_states = self.decoder(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1602, in _call_impl
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/transformer/transformer_block.py", line 493, in forward
[rank0]:     hidden_states, context = layer(
[rank0]:   File "/opt/megatron-lm/megatron/core/transformer/transformer_layer.py", line 390, in __call__
[rank0]:     return super(MegatronModule, self).__call__(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1602, in _call_impl
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/transformer/transformer_layer.py", line 294, in forward
[rank0]:     attention_output_with_bias = self.self_attention(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1602, in _call_impl
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/transformer/attention.py", line 314, in forward
[rank0]:     core_attn_out = self.core_attention(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1602, in _call_impl
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/extensions/transformer_engine.py", line 699, in forward
[rank0]:     core_attn_out = super().forward(
[rank0]:   File "/opt/TransformerEngine/transformer_engine/pytorch/attention.py", line 6984, in forward
[rank0]:     raise Exception("No dot product attention support for the provided inputs!")
[rank0]: Exception: No dot product attention support for the provided inputs!
```

**Steps/Code to reproduce bug**

```
nemo llm pretrain -f gemma_2b trainer.max_steps=10
```


**Expected behavior**

Exception should not pop up or abort.

**Environment overview (please complete the following information)**
```
sudo docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it nvcr.io/nvidia/nemo:dev
```
on GCP VM
**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version: Ubuntu Ubuntu 22.04.4 LTS
- PyTorch version: 2.4.0a0+3bcc3cddb5.nv24.07
- Python version: 3.10.12

**Additional context**
 GPU model: NVIDIA L4 * 8
