Nemo run pretrain raised Exception: No dot product attention support for the provided inputs
State: closed | Pull request: no
Labels: name: bug

name: 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


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.