NVIDIA / NVIDIA/Megatron-LM

Feature Request: DSA/DSv4 Indexer Replay for RL

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

Description

## Summary

Add replay support for DSA / DeepSeek-V4 hybrid-attention indexer top-k decisions during RL training.

Megatron already has MoE router replay support for R3-style RL. We should add something similar for sparse-attention indexer decisions in DSA / DSv4 hybrid attention.

## Motivation

For RL, rollout inference and train-side recomputation may run through different kernels, precisions, batching layouts, CUDA graph paths, or packed-sequence layouts. In DSA/DSv4 sparse attention, the indexer makes a discrete top-k selection of which KV positions each query token attends to.

Small numerical differences can cause rollout inference and training to choose different indexer top-k positions for the same tokens. That means train-side logprobs may be computed through a different sparse-attention graph than the one used during rollout, creating train/rollout mismatch.

This is the same instability that MoE R3 addresses for expert routing, but applied to sparse-attention indexer selections.

## Proposed Solution

Add an opt-in replay path, e.g. `dsa_enable_indexer_replay`, that supports record mode, replay mode, and RL frameworks to carry recorded indexer metadata with rollout data.

This should also be tested with downstream RL frameworks like verl.

## Related Work
MoE Router Replay / R3 issue: https://github.com/NVIDIA/Megatron-LM/issues/4168
MCore RouterReplay PR: https://github.com/NVIDIA/Megatron-LM/pull/2693
DSA support: https://github.com/NVIDIA/Megatron-LM/pull/2440

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.