swiglu activation doesn't support sequence_parallel
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

When changing activation to `swiglu` and enable `sequence_parallel`, the training crashes.

**Steps/Code to reproduce bug**

1. Checkout out repo tag: `git checkout v1.19.1`
2. Config modifications:
  ```diff
  --- a/examples/nlp/language_modeling/conf/megatron_gpt_config.yaml
  +++ b/examples/nlp/language_modeling/conf/megatron_gpt_config.yaml
  @@ -51,7 +51,7 @@ model:
     virtual_pipeline_model_parallel_size: null # interleaved pipeline
  
     # model architecture
  -  encoder_seq_length: 512
  +  encoder_seq_length: 8192
     max_position_embeddings: ${.encoder_seq_length}
     num_layers: 12
     hidden_size: 768
  @@ -71,8 +71,8 @@ model:
     pre_process: True # add embedding
     post_process: True # add pooler
     persist_layer_norm: True # Use of persistent fused layer norm kernel.
  -  bias: True # Whether to use bias terms in all weight matrices.
  -  activation: 'gelu' # Options ['gelu', 'geglu', 'swiglu', 'reglu', 'squared-relu', 'fast-geglu', 'fast-swiglu', 'fast-reglu']
  +  bias: False # Whether to use bias terms in all weight matrices.
  +  activation: 'swiglu' # Options ['gelu', 'geglu', 'swiglu', 'reglu', 'squared-relu', 'fast-geglu', 'fast-swiglu', 'fast-reglu']
     headscale: False # Whether to learn extra parameters that scale the output of the each self-attention head.
     transformer_block_type: 'pre_ln' # Options ['pre_ln', 'post_ln', 'normformer']
     openai_gelu: False # Use OpenAI's GELU instead of the default GeLU
  @@ -88,8 +88,8 @@ model:
       library: 'megatron'
       type: 'GPT2BPETokenizer'
       model: null
  -    vocab_file: null
  -    merge_file: null
  +    vocab_file: /data/meta/gpt2/vocab.json
  +    merge_file: /data/meta/gpt2/merges.txt
       delimiter: null # only used for tabular tokenizer
       sentencepiece_legacy: False # Legacy=True allows you to add special tokens to sentencepiece tokenizers.
  
  @@ -107,8 +107,8 @@ model:
     # Fusion
     grad_div_ar_fusion: True # Fuse grad division into torch.distributed.all_reduce. Only used with O2 and no pipeline parallelism..
     gradient_accumulation_fusion: False # Fuse weight gradient accumulation to GEMMs. Only used with pipeline parallelism and O2.
  -  bias_activation_fusion: True # Use a kernel that fuses the bias addition from weight matrices with the subsequent activation function.
  -  bias_dropout_add_fusion: True # Use a kernel that fuses the bias addition, dropout and residual connection addition.
  +  bias_activation_fusion: False # Use a kernel that fuses the bias addition from weight matrices with the subsequent activation function.
  +  bias_dropout_add_fusion: False # Use a kernel that fuses the bias addition, dropout and residual connection addition.
     masked_softmax_fusion: True # Use a kernel that fuses the attention softmax with it's mask.
     get_attention_mask_from_fusion: True # When using fused softmax it will create the attention mask so we won't copy it to the pipeline stages.
  
  @@ -187,7 +187,9 @@ model:
       # Dictionary: can override from CLI "model.data.data_prefix"={"train":[1.0, /path/to/data], "validation":/path/to/data, "test":/path/to/test}
       # Or see example below:
       # "model.data.data_prefix: {train:[1.0,/path/to/data], validation:[/path/to/data], test:[/path/to/test]}"
  -    data_prefix: ???
  +    data_prefix:
  +    - 1.0
  +    - /data/pretrain/owt2_clean_v2/hfbpe_gpt_training_data_text_document
       index_mapping_dir: null # path to save index mapping .npy files, by default will save in the same location as data_prefix
       data_impl: mmap
       splits_string: 900,50,50
  ```

3. Run command: 
```
python examples/nlp/language_modeling/megatron_gpt_pretraining.py --config-path=/workspace/project/NeMo-upstream/examples/nlp/language_modeling/conf --config-name=megatron_gpt_config trainer.devices=4 trainer.num_nodes=1 model.micro_batch_size=1 model.global_batch_size=8 model.tensor_model_parallel_size=4 model.pipeline_model_parallel_size=1 model.sequence_parallel=True
```
4. The error trackback:
```
Traceback (most recent call last):
  File "/workspace/project/NeMo-upstream/examples/nlp/language_modeling/megatron_gpt_pretraining.py", line 86, in main
    trainer.fit(model)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit
    call._call_and_handle_interrupt(
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 1112, in _run
    results = self._run_stage()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 1191, in _run_stage
    self._run_train()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 1214, in _run_train
    self.fit_loop.run()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/fit_loop.py", line 267, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 213, in advance
    batch_output = self.batch_loop.run(kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(optimizers, kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 202, in advance
    result = self._run_optimization(kwargs, self._optimizers[self.optim_progress.optimizer_position])
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 249, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, kwargs.get("batch_idx", 0), closure)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 370, in _optimizer_step
    self.trainer._call_lightning_module_hook(
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 1356, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/core/module.py", line 1742, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/core/optimizer.py", line 169, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/strategies/ddp.py", line 280, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/strategies/strategy.py", line 234, in optimizer_step
    return self.precision_plugin.optimizer_step(
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/plugins/precision/native_amp.py", line 75, in optimizer_step
    closure_result = closure()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 149, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 135, in closure
    step_output = self._step_fn()
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 419, in _training_step
    training_step_output = self.trainer._call_strategy_hook("training_step", *kwargs.values())
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/trainer/trainer.py", line 1494, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/strategies/ddp.py", line 351, in training_step
    return self.model(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1533, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/distributed.py", line 1137, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/distributed.py", line 1094, in _run_ddp_forward
    return self.module(*inputs, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1533, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/overrides/base.py", line 98, in forward
    output = self._forward_module.training_step(*inputs, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/nemo/utils/model_utils.py", line 380, in wrap_training_step
    output_dict = wrapped(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 469, in training_step
    self.allreduce_sequence_parallel_gradients()
  File "/usr/local/lib/python3.8/dist-packages/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 571, in allreduce_sequence_parallel_gradients
    coalesced = torch._utils._flatten_dense_tensors(grads)
  File "/usr/local/lib/python3.8/dist-packages/torch/_utils.py", line 459, in _flatten_dense_tensors
    return torch._C._nn.flatten_dense_tensors(tensors)
RuntimeError: torch.cat(): expected a non-empty list of Tensors
```
**Expected behavior**

The training can proceed normally

**Environment overview (please complete the following information)**

 - Docker images: `nvcr.io/nvidia/nemo:23.04`
 - Start docker container: `sudo nvidia-docker run -it --shm-size 64g --network host --hostname nemo-2304 --name nemo-2304 -v /path/on/host:/path/on/container/ nvcr.io/nvidia/nemo:23.04 bash`

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version

**Additional context**

GPU: 4 V100