Llama example scripts
State: closed | Pull request: no
Labels: name: stale

I was trying to configure a test for the Llama model using Megatron Nemo. While I discovered an example configuration file (https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_llama_config.yaml), I couldn't locate any sample scripts specifically designed for launching training with Llama. In my attempt, I used the Llama config file along with "megatron_gpt_pretraining.py," but unfortunately, it didn't work. Could you provide guidance or share any example scripts tailored for Llama pretraining?

This is the full stack trace, and the issue is `Exception: The provided qkv memory layout is not supported!`
```0: Traceback (most recent call last):
0:   File "/fsx/users/xuteng/NeMo/examples/nlp/language_modeling/megatron_llama_pretraining.py", line 40, in main
0:     trainer.fit(model)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
0:     call._call_and_handle_interrupt(
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
0:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
0:     return function(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
0:     self._run(model, ckpt_path=ckpt_path)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 980, in _run
0:     results = self._run_stage()
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1023, in _run_stage
0:     self.fit_loop.run()
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 202, in run
0:     self.advance()
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 355, in advance
0:     self.epoch_loop.run(self._data_fetcher)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
0:     self.advance(data_fetcher)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 219, in advance
0:     batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 188, in run
0:     self._optimizer_step(kwargs.get("batch_idx", 0), closure)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 266, in _optimizer_step
0:     call._call_lightning_module_hook(
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 145, in _call_lightning_module_hook
0:     output = fn(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1270, in optimizer_step
0:     optimizer.step(closure=optimizer_closure)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 161, in step
0:     step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 257, in optimizer_step
0:     optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 231, in optimizer_step
0:     return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 73, in optimizer_step
0:     return super().optimizer_step(optimizer, model=model, closure=closure, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 116, in optimizer_step
0:     return optimizer.step(closure=closure, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
0:     return wrapped(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/optim/optimizer.py", line 280, in wrapper
0:     out = func(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/apex/contrib/optimizers/distributed_fused_adam.py", line 2071, in step
0:     loss = closure()
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 103, in _wrap_closure
0:     closure_result = closure()
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 142, in __call__
0:     self._result = self.closure(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
0:     return func(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 128, in closure
0:     step_output = self._step_fn()
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 315, in _training_step
0:     training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 293, in _call_strategy_hook
0:     output = fn(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 330, in training_step
0:     return self.model(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/overrides/base.py", line 90, in forward
0:     output = self._forward_module.training_step(*inputs, **kwargs)
0:   File "/fsx/users/xuteng/NeMo/nemo/utils/model_utils.py", line 381, in wrap_training_step
0:     output_dict = wrapped(*args, **kwargs)
0:   File "/fsx/users/xuteng/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 607, in training_step
0:     loss_mean = self.fwd_bwd_step(dataloader_iter, batch_idx, False)
0:   File "/fsx/users/xuteng/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 509, in fwd_bwd_step
0:     losses_reduced_per_micro_batch = fwd_bwd_function(
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/pipeline_parallel/schedules.py", line 327, in forward_backward_no_pipelining
0:     output_tensor = forward_step(
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/pipeline_parallel/schedules.py", line 183, in forward_step
0:     output_tensor, loss_func = forward_step_func(data_iterator, model)
0:   File "/fsx/users/xuteng/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 884, in fwd_output_and_loss_func
0:     output_tensor = model(**forward_args)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/module.py", line 137, in forward
0:     outputs = self.module(*inputs, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/models/gpt/gpt_model.py", line 166, in forward
0:     hidden_states = self.decoder(
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/transformer_block.py", line 302, in forward
0:     hidden_states = self._checkpointed_forward(
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/transformer_block.py", line 204, in _checkpointed_forward
0:     hidden_states, context = tensor_parallel.checkpoint(
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/tensor_parallel/random.py", line 272, in checkpoint
0:     return CheckpointFunction.apply(function, distribute_saved_activations, *args)
0:   File "/opt/conda/lib/python3.10/site-packages/torch/autograd/function.py", line 506, in apply
0:     return super().apply(*args, **kwargs)  # type: ignore[misc]
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/tensor_parallel/random.py", line 211, in forward
0:     outputs = run_function(*args)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/transformer_block.py", line 168, in custom_forward
0:     hidden_states, context = layer(
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/transformer_layer.py", line 153, in forward
0:     attention_output_with_bias = self.self_attention(
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/attention.py", line 263, in forward
0:     core_attn_out = self.core_attention(
0:   File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1504, in _call_impl
0:     return forward_call(*args, **kwargs)
0:   File "/opt/conda/lib/python3.10/site-packages/megatron/core/transformer/custom_layers/transformer_engine.py", line 427, in forward
0:     return super().forward(
0:   File "/opt/conda/lib/python3.10/site-packages/transformer_engine/pytorch/attention.py", line 2085, in forward
0:     qkv_layout = _get_qkv_layout(query_layer, key_layer, value_layer,
0:   File "/opt/conda/lib/python3.10/site-packages/transformer_engine/pytorch/attention.py", line 1195, in _get_qkv_layout
0:     raise Exception("The provided qkv memory layout is not supported!")
0: Exception: The provided qkv memory layout is not supported!```

have you tried training recipes here? https://github.com/NVIDIA/NeMo-Megatron-Launcher/tree/master/launcher_scripts/conf/training/llama

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.