[QUESTION] Inconsistent Log Prob Calculation in MCore When Using Chunked vs Padded Inputs
- 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
chunk 1+ chunk 2:
chunk 1 + padding:
---
### 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
Assessment
This issue has not been assessed yet.