Training-generation mismatch
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Problem
Colocated Megatron GRPO can report non-zero `gen_kl_error` even when train and gen share weights: train/scoring and generation take different kernel paths (MoE combine order, log-softmax, batch-variant cuBLAS, stale KV after refit, Mamba prefill/decode drift, etc.). That breaks true on-policy RL.
**Success criterion:** `gen_kl_error → 0` on certified beds (step-0 and multi-step), with documented recipes and a fast cert harness where possible.
This issue tracks **three complementary workstreams**. They share the same metric but use different engine/layout assumptions — they are **not** duplicates and should not be merged into one PR.
---
## Workstreams
### 1. Colocated `transformer_engine` — [PR #3208](https://github.com/NVIDIA-NeMo/RL/pull/3208)
**Owner:** @YigongQin
**Status:** Open (community-request)
Single yaml switch: `policy.megatron_cfg.zero_train_gen_mismatch: true`
| Aspect | Choice |
|--------|--------|
| Layout | Colocated train+gen (1-node) |
| Engine | `transformer_engine` |
| Patches | NeMo-RL runtime only (`zero_train_gen_kl_patches/`) — **no Megatron-LM fork** |
| KV / graphs | `kv_cache_management_mode: recompute`, CUDA graphs off (correctness first) |
| Models | Qwen2.5-1.5B, Qwen3-30B-A3B, Nemotron-3-Nano-30B-A3B (Mamba) |
| Precision | BF16, MXFP8 |
**Deliverables**
- [ ] Merge #3208 with BF16 zero-KL recipes
- [ ] Unit tests under `tests/unit/models/generation/megatron/zero_train_gen_kl_patches/`
- [ ] Repro launcher: `research/megatron-inference-true-on-policy/run_zero_kl_precision.sh`
**Follow-up (separate PRs, not blocking merge)**
- Re-enable CUDA graphs + `persist` KV once upstream post-refit invalidation exists
- Port shared infra from #3531 (see Shared items below)
---
### 2. Dedicated-node `inference_optimized` (EP8/TP1) — [PR #3531](https://github.com/NVIDIA-NeMo/RL/pull/3531)
**Owner:** @utkarsh530
**Status:** Open (draft)
Certified fast path: zero `gen_kl` with ~1.02× throughput tax vs non-det, CUDA graphs on.
| Aspect | Choice |
|--------|--------|
| Layout | Dedicated gen node (`colocated.enabled=false`) |
| Engine | `inference_optimized` + [Megatron-LM fork](https://github.com/utkarsh530/Megatron-LM/tree/det-inference-ep8tp1-certified) |
| Parallelism | Gen EP8/TP1 |
| Router replay | **Off** (bitwise logits ⇒ same routes) |
| Model | Qwen3-30B-A3B (cert bed) |
**Deliverables**
- [ ] Merge #3531 certified recipe + `examples/gen_kl_harness.py` (~8 min determinism gate)
- [ ] Colocated config fail-fast in `megatron_generation.py`
- [ ] Refit param-gather sync (`NRL_REFIT_PARAM_SYNC`) for overlap_param_gather safety
- [ ] Design doc: deterministic inference EP8/TP1
---
### 3. MXFP8 zero-KL (future)
**Owner:** TBD (likely @YigongQin after #3208 lands)
**Status:** Not started / blocked on BF16 cert
Extend zero-KL to MXFP8 training + colocated Megatron gen on Blackwell.
**Scope**
- Certify `gen_kl_error → 0` with `fp8_cfg.fp8_recipe: mxfp8` on Qwen3-30B-A3B (and Nano if applicable)
- Align train/gen FP8 scaling paths so refit does not introduce logprob drift
- Recipe pattern already sketched in #3208:
```yaml
policy:
megatron_cfg:
zero_train_gen_mismatch: true
fp8_cfg:
enabled: true
fp8_recipe: mxfp8
optimizer:
use_precision_aware_optimizer: false # Qwen3-30B
Contributor guide
Research direction
First review PRs #3208 and #3531 to identify which of the three separate workstreams is available, then read the referenced zero_train_gen_kl_patches tests and relevant generation entry points. Run the applicable certification recipe or gen_kl_harness.py; done means gen_kl_error approaches zero on the stated step-0 and multi-step certified beds, with the recipe and tests documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning, performance, testing
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100