AssertionError: (RMSNorm) is not supported in FusedLayerNorm
State: closed | Pull request: no
Labels: name: stale

We are using the nemo:24.05 container.

We migrated from GPT-Neox. We are trying to continually pretrain a Llama 2 7B model using Nemo. We made it work, but, for some reason, it needs 8x more GPUs for the same configuration when compared to GPT-Neox.

Some people told us that flash attention does not work when setting `transformer_engine=True`, and we think this is the issue. However, when we set `transformer_engine=False` it gives the following error:

Traceback (most recent call last):
  File "/opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py", line 40, in main
    model = MegatronGPTModel(cfg.model, trainer)
  File "/opt/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 342, in __init__
    self.model = build_model(
  File "/opt/NeMo/nemo/collections/nlp/modules/common/megatron/build_model.py", line 90, in build_model
    model = model_provider_func(
  File "/opt/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 413, in model_provider_func
    model = MCoreGPTModel(
  File "/opt/megatron-lm/megatron/core/models/gpt/gpt_model.py", line 95, in __init__
    self.decoder = TransformerBlock(
  File "/opt/megatron-lm/megatron/core/transformer/transformer_block.py", line 143, in __init__
    self._build_layers()
  File "/opt/megatron-lm/megatron/core/transformer/transformer_block.py", line 158, in _build_layers
    [
  File "/opt/megatron-lm/megatron/core/transformer/transformer_block.py", line 159, in <listcomp>
    build_layer(layer_spec, i + 1)
  File "/opt/megatron-lm/megatron/core/transformer/transformer_block.py", line 154, in build_layer
    return build_module(layer_spec, config=self.config, layer_number=layer_number,)
  File "/opt/megatron-lm/megatron/core/transformer/spec_utils.py", line 107, in build_module
    raise type(e)(f"{str(e)} when instantiating {module.__name__}").with_traceback(
  File "/opt/megatron-lm/megatron/core/transformer/spec_utils.py", line 99, in build_module
    return module(
  File "/opt/megatron-lm/megatron/core/transformer/transformer_layer.py", line 76, in __init__
    self.input_layernorm = build_module(
  File "/opt/megatron-lm/megatron/core/transformer/spec_utils.py", line 107, in build_module
    raise type(e)(f"{str(e)} when instantiating {module.__name__}").with_traceback(
  File "/opt/megatron-lm/megatron/core/transformer/spec_utils.py", line 99, in build_module
    return module(
  File "/opt/megatron-lm/megatron/core/fusions/fused_layer_norm.py", line 68, in __init__
    self.config.normalization == "LayerNorm"
AssertionError: (RMSNorm) is not supported in FusedLayerNorm when instantiating FusedLayerNorm when instantiating TransformerLayer


Apparently, if we do not use transformer engine then only LayerNorm is supported, but Llama uses RMSNorm. 

How can we fix this? Is it true that flash attention does not work when enabling transformer engine? Why does it need 8x more GPUs with respect to GPT-Neox?

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.

can you solved it?

@yuanzhiyong1999 No. As of now, we think that flash attention is working properly even when `transformer_engine=True` (with this setting, flash attention is always uses, that is why the `use_flash_attention` parameter does nothing).

We think that the 8x more GPU issue cannot be solved from our end. We think it is a consequence of using Nemo and that we have to live with it (probably because it's not well optimized or something). In newer containers we always use `transformer_engine=True` now.

thank you for your reply，I want to know which framework is more efficient in training, gpt-neox or nemo?

---- Replied Message ----
| From | ***@***.***> |
| Date | 09/25/2024 19:25 |
| To | NVIDIA/NeMo ***@***.***> |
| Cc | Zhiyong ***@***.***>,
Mention ***@***.***> |
| Subject | Re: [NVIDIA/NeMo] AssertionError: (RMSNorm) is not supported in FusedLayerNorm (Issue #9790) |

@yuanzhiyong1999 No. As of now, we think that flash attention is working properly even when transformer_engine=True (with this setting, flash attention is always uses, that is why the use_flash_attention parameter does nothing).

We think that the 8x more GPU issue cannot be solved from our end. We think it is a consequence of using Nemo and that we have to live with it (probably because it's not well optimized or something). In newer containers we always use transformer_engine=True now.

—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you were mentioned.Message ID: ***@***.***>

@yuanzhiyong1999 From what we have experienced GPT-NeoX is more efficient. But it has some downsides. For instance, I don't think it supports newer models like Llama 3.1 (although it may in the future). 

thank you very much！

---- Replied Message ----
| From | ***@***.***> |
| Date | 09/25/2024 21:54 |
| To | NVIDIA/NeMo ***@***.***> |
| Cc | Zhiyong ***@***.***>,
Mention ***@***.***> |
| Subject | Re: [NVIDIA/NeMo] AssertionError: (RMSNorm) is not supported in FusedLayerNorm (Issue #9790) |

@yuanzhiyong1999 From what we have experienced GPT-NeoX is more efficient. But it has some downsides. For instance, I don't think it supports newer models like Llama 3.1 (although it may in the future).

—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you were mentioned.Message ID: ***@***.***>