[BUG] grad lost when sharing embeddings if decoder_pp_size = 1
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**
When sharing embeddings on same pipeline stage, `zero_out_wgrad` will be set to the weight of embeddings to make sure grad accumulation into `main_grad` is correct, as shown in [language_module.py#L127](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/common/language_module/language_module.py#L127) and [layers.py#L526](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/layers.py#L526)
However, if model has encoder and **decoder_pp_size = 1** , function `get_pipeline_model_parallel_world_size` cannot return 1, while decoder needs `zero_out_wgrad`.
So, decoder will loss head embeddings grad in the hook that copy grad to `main_grad`, as shown in [distributed_data_parallel.py#L493](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/distributed_data_parallel.py#L493)
**To Reproduce**
Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.
**Expected behavior**
A clear and concise description of what you expected to happen.
**Stack trace/logs**
If applicable, add the stack trace or logs from the time of the error.
**Environment (please complete the following information):**
- Megatron-LM commit ID
- PyTorch version
- CUDA version
- NCCL version
**Proposed fix**
```
if parallel_state.get_pipeline_model_parallel_world_size() - parallel_state.get_pipeline_model_parallel_decoder_start() == 1:
self.shared_embedding_or_output_weight().zero_out_wgrad = True
```
**Additional context**
Add any other context about the problem here.
Contributor guide
Research direction
Start with the linked sections of language_module.py, layers.py, and distributed_data_parallel.py, then trace the shared embedding path when an encoder is present and decoder_pp_size is 1. Reproduce the configuration described in the issue and verify that the decoder head embedding gradient is preserved in main_grad; the proposed pipeline-size condition indicates the expected behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100