NVIDIA / NVIDIA/Megatron-LM

[BUG] grad lost when sharing embeddings if decoder_pp_size = 1

Open
#1,548 3 comments 0 reactions 0 assignees View on GitHub
bug community-request waiting-on-customer
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.