NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Inconsistent Log Prob Calculation in MCore When Using Chunked vs Padded Inputs

Open
#1,809 1 comment 0 reactions 0 assignees View on GitHub
module: rl question
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

### Summary
We are using partial rollout for LLM + RL. When computing `old_logprob` for the first chunk (`chunk1`) using MCore, we observe inconsistent results:

- `chunk1 + padding` (padded to full context length), or
- `chunk1 + chunk2` (full sequence without padding).

Using the same model (Qwen3-1.7B, dense), same weights, and same position IDs for `chunk1`, the computed log probabilities for `chunk1` differ significantly — especially on low-probability tokens and even starting from the first token of the response.

This issue **does not occur** in Transformer native forward passes - chunked and padded runs produce identical logits for `chunk1`.

---

### Observed Issues

#### 1. MCore Self-Consistency (Single Chunk)

- When generating in **one chunk**, recomputing log probs twice yields identical results (with `temperature=1` and `0.6`) — so MCore is deterministic in isolation.

#### 2. Chunk + Padding vs Chunk + Chunk Inconsistency (Main Issue)

All experiments use **Qwen3-1.7B (dense)**.

- **Exp 1**:
- `output_weight`, `rotary_pos_emb` (first 10360), and `decoder_input` (first 2590 tokens) are identical.
- But **hidden states diverge** in the first half — resembling precision/numerical instability, not logical error.

in: megatron/core/models/gpt/gpt_model.py
Image

chunk 1+ chunk 2:

Image

chunk 1 + padding:

Image

---

### Expected Behavior

Log probabilities for `chunk1` should be **identical** regardless of whether it is:

- Run standalone + padding, or
- Run as prefix to `chunk2`.

This is fundamental for RLHF/RLAIF training where `old_logprob` must remain consistent across rollouts and re-computations.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.