dim unmatch when doing sft with tensor parallel and sequence parallel and LoRA
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

I was training to run sft based on Mixtral-8x7B-instruct model with tensor parallel size=4 (sequence parallel=True) and LoRA (target modules =[all]). 
It reports that the output dims of original module and the corresponding lora adapter module is not matched so they cannot be added together. 

**Steps/Code to reproduce bug**

I used the recommended docker nvcr.io/nvidia/nemo:24.07 and my scripts is as follow:
```shell
export CUDA_VISIBLE_DEVICES="4,5,6,7"
nemo_ckpt_path=/path/to/Mixtral-8x7B-instruct-v0.1-nemo
train_data='/path/to/train.json'
valid_data='/path/to/valid.json'
output='/data/results/mixtral-8x7b-instruct-sft'

python /opt/NeMo-Aligner/examples/nlp/gpt/train_gpt_sft.py \
   trainer.precision=bf16 \
   trainer.num_nodes=1 \
   trainer.devices=4 \
   trainer.sft.max_steps=-1 \
   trainer.sft.limit_val_batches=40 \
   trainer.sft.val_check_interval=1000 \
   trainer.sft.max_epochs=2 \
   model.tensor_model_parallel_size=4 \
   model.sequence_parallel=True \
   model.use_flash_attention=True \
   model.encoder_seq_length=1024 \
   model.megatron_amp_O2=False \
   model.restore_from_path=$nemo_ckpt_path \
   model.optim.lr=5e-6 \
   model.optim.name=fused_adam\
   model.answer_only_loss=True \
   model.data.num_workers=0 \
   model.data.train_ds.micro_batch_size=1 \
   model.data.train_ds.global_batch_size=128 \
   model.data.train_ds.file_path=$train_data \
   model.data.validation_ds.micro_batch_size=1 \
   model.data.validation_ds.global_batch_size=128 \
   model.data.validation_ds.file_path=$valid_data \
   model.peft.peft_scheme=lora \
   model.peft.lora_tuning.target_modules=[attention_dense,mlp_fc1,mlp_fc2] \
   model.peft.lora_tuning.adapter_dim=128 \
   exp_manager.create_wandb_logger=False \
   exp_manager.explicit_log_dir=$output \
   exp_manager.wandb_logger_kwargs.project=sft_run \
   exp_manager.wandb_logger_kwargs.name=ruozhiba_sft_run \
   exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True \
   exp_manager.resume_if_exists=True \
   exp_manager.resume_ignore_no_checkpoint=True \
   exp_manager.create_checkpoint_callback=True \
   exp_manager.checkpoint_callback_params.monitor=validation_loss

```

And it runs into error like:
```shell
  ...
  File "/opt/NeMo-Aligner/nemo_aligner/algorithms/supervised.py", line 145, in train_single_step
    loss_mean, metrics = self.model.get_loss_and_metrics(batch=batch, forward_only=False)
  File "/opt/NeMo-Aligner/nemo_aligner/models/nlp/gpt/gpt_sft_model.py", line 93, in get_loss_and_metrics
    losses_reduced = fwd_bwd_function(
  File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 439, in forward_backward_no_pipelining
    output_tensor, num_tokens = forward_step(
  File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 264, in forward_step
    output_tensor, loss_func = forward_step_func(data_iterator, model)
  File "/opt/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1273, in fwd_output_and_loss_func
    output_tensor = model(**forward_args)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/megatron-lm/megatron/core/models/gpt/gpt_model.py", line 191, in forward
    hidden_states = self.decoder(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/megatron-lm/megatron/core/transformer/transformer_block.py", line 411, in forward
    hidden_states, context = layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/megatron-lm/megatron/core/transformer/transformer_layer.py", line 178, in forward
    attention_output_with_bias = self.self_attention(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/NeMo/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py", line 202, in forward
    query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)
  File "/opt/NeMo/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py", line 117, in get_query_key_value_tensors
    mixed_qkv = mixed_qkv + lora_mixed_qkv
RuntimeError: The size of tensor a (400) must match the size of tensor b (1600) at non-singleton dimension 0

```
And I try to fix this by modify the following two .py files:
/opt/NeMo/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py


```patch
--- mcore_mixins.py     2024-08-28 02:51:12.000000000 +0000
+++ /opt/NeMo/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py     2024-08-28 02:51:13.952291250 +0000
@@ -379,7 +379,6 @@
             if lora_fc2_adapter and self.adapter_cfg[AdapterName.LORA_4HtoH_ADAPTER]['enabled']:
                 lora_output = lora_fc2_adapter(intermediate_parallel)
             elif lora_moe_fc2_adapter and self.adapter_cfg[AdapterName.LORA_MOE_4HtoH_ADAPTER]['enabled']:
-                lora_moe_fc2_adapter.expert_adapters[expert_idx].input_is_parallel=False
                 lora_output = lora_moe_fc2_adapter(intermediate_parallel, expert_idx)

             output = output + lora_output
```


/opt/NeMo/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py
```patch
--- parallel_adapters.py        2024-08-28 02:52:03.000000000 +0000
+++ /opt/NeMo/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py        2024-08-28 02:52:04.712871536 +0000
@@ -291,7 +291,7 @@

         if self.norm_position == 'pre':
             x = self.layer_norm(x)
-        if self._sequence_parallel and not self.input_is_parallel and self.norm_position=='pre':
+        if self._sequence_parallel and not self.input_is_parallel:
             # for attention_qkv and linear_fc1
             # layernorm before lora is impacted by sequence parallel,
             # hence seq dim need to be gathered right before lora linear layers

```
After the modifications, I can run sft with LoRA and tensor and sequence parallel, but I am not sure it runs correctly. Hope you guys can provide elegant solutions for it.


**Expected behavior**

LoRA can be used with tensor and sequence parallel.

**Environment overview (please complete the following information)**

nvcr.io/nvidia/nemo:24.07
docker run
```shell
HERE=$(pwd -P) 
user=`whoami`
uid=`id -u`
gid=`id -g`

docker run \
  -v /dev/shm:/dev/shm\
  -v /data:/data \
  -e USER=$user -e UID=$uid -e GID=$gid \
  -v $HERE:/home/dir/ \
  -w /home/dir/ \
  --security-opt \
  seccomp=unconfined \
  -it \
  --rm \
  --network=host \
  --gpus all \
  nvcr.io/nvidia/nemo:24.07

```

**Environment details**

I used the default environment of the nemo docker




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.

Hello, I have also encountered this bug. I am referring to the documentation here: [NeMo Framework PEFT with Llama2, Mixtral-8x7B, and Nemotron 4 340B](https://docs.nvidia.com/nemo-framework/user-guide/latest/playbooks/nemoframeworkpeft.html). 

When training Llama2 with Nemo, an error occurs if the following parameters are set: `TP=2, sequence_parallel=True, and model.peft.peft_scheme="lora"`. 

The error is a `RuntimeError: "The size of tensor A (2256) must match the size of tensor B (4512) at non-singleton dimension 0."`

Hi, 

can you pass `+model.tp_comm_overlap_disable_qkv=True` to CLI arguments and report back if that fixes the issue?

> Hi,
> 
> can you pass `+model.tp_comm_overlap_disable_qkv=True` to CLI arguments and report back if that fixes the issue?

I added `+model. tp_comm_overlap-disable_qkv=True`, but it still doesn't work

Under `TP2PP1DP1` participant

When I set `model. sequence paramallel=False` and `model. tp_comm_overlap-disable_qkv=False` and `Model. sequence=False` and `model. tp_comm_overlap_isable_qkv=True`, it can run normally,

But when I set them separately: `Model. sequence=True` and `model. tp_comm_overlap_isable_qkv=False` and `Model. sequence=True` and `model. tp_comm_overlap_isable_qkv=True`. Error will be reported: `AttributeError: 'TELayerNormColumnParallelLinear' object has no attribute 'ub_overlap.ag'. `

This is the log: 
- Work smoothly:
  - SP=False disable_qkv=False: [A800_tune_peft_SP0overlap0.log](https://github.com/user-attachments/files/17488768/A800_tune_peft_SP0overlap0.log)
  - SP=False disable_qkv=True: [A800_tune_peft_SP0overlap1.log](https://github.com/user-attachments/files/17488732/A800_tune_peft_SP0overlap1.log)
- Not working: 
  - SP=True disable_qkv=False: [A800_tune_peft_SP1overlap0.log](https://github.com/user-attachments/files/17488733/A800_tune_peft_SP1overlap0.log)
  - SP=True disable_qkv=True: [A800_tune_peft_SP1overlap1.log](https://github.com/user-attachments/files/17488731/A800_tune_peft_SP1overlap1.log)

I have reviewed the source code of Nemo and Megatron LM,

In `/Megatron-LM-core_r0.7.0.beta/megatron/core/transformer/custom_1ayers/transformar_engine. py`

```python
if tp_comm_buffer_name == 'qkv' and self.config.tp_comm_disable_qkv:
extra_kwargs["ub_overlap_ag"] = False
extra_kwargs["ub_overlap_rs_dgrad"] = False
```
Could it be this part of the content that caused the error?

My environment version is
- NeMo-2.0.0.rc0.beta
- Megatron-LM-core_r0.7.0.beta
- TransformerEngine-1.6rc2

can you also try with https://github.com/NVIDIA/NeMo/pull/10920? thanks

> can you also try with [#10920](https://github.com/NVIDIA/NeMo/pull/10920)? thanks

succeed! Thank you very much for your help!

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.