[ENHANCEMENT] Shared-prefix training with context parallelism for exact agent trajectories
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Is your feature request related to a problem? Please describe.**
Agentic RL episodes contain many exact model calls whose contexts share prefixes. Retokenization, subagents, and different executed expert routes mean that concatenating conversation text is not a faithful training representation. Independently packing each captured call is correct, but repeatedly computes identical prefixes.
We would like an opt-in training representation that shares compatible physical prefixes while preserving each sampled action's conditioning context and loss contribution, including with context parallelism. This issue precedes a draft implementation for design discussion; it is not a request to merge an unreviewed numerical change.
**Describe the solution you'd like**
- Extend packed-sequence metadata with DFS-linearized segments, parent/depth information, CP-local query runs, and sampled-edge projection indices.
- Keep exact token/routing/provenance compatibility checks and bounded fragment planning in the training framework, outside MCore.
- Compute attention for each query run against only its ancestor path. Gather CP-sharded K/V with autograd support; keep queries local. Checkpoint per-run attention to bound saved intermediates.
- Restore differentiable Mamba convolution and recurrent state at branch points. Ordinary packed execution remains unchanged.
- Project only sampled-edge source hidden states through the vocabulary head, retaining one loss owner per sampled action.
- Reject unsupported attention modifiers, inference contexts, custom output processors, and incompatible MTP execution rather than silently ignoring them.
CP>1 is necessary for long-context training. The proposed implementation gathers full K/V within the CP group, so communication and temporary storage scale with physical fragment length; it is not a claim of constant-memory or zero-communication tree attention. The trainer bounds fragments and may repeat ancestors across fragments to honor the memory bound.
**Describe alternatives you've considered**
1. Independent exact-call sequences: the correctness reference, but duplicates every shared prefix.
2. Maximal compatible root-to-leaf paths: removes strict-prefix calls, but still duplicates common trunks across branches.
3. Concatenating rendered conversation text: changes tokenization and conditioning, so it is not an equivalent optimization.
4. Restricting training to CP1: insufficient for our long-context workloads.
**Additional context**
In one integrated 16-episode training gate, naive captured calls contained 40,398,130 input/output tokens. Compatible physical trees contained 3,744,651 tokens before fragment duplication and 3,919,929 after bounded fragmentation. All 987,024 sampled output tokens retained exactly one loss owner. These are token-count measurements, not an equivalent wall-clock or FLOP speedup claim.
Attention CP1/CP2 and TP2 edge-projection parity checks passed. BF16 Mamba parameter-gradient parity does **not** pass all existing test thresholds: a focused control also reproduced a discrepancy between unmodified upstream fused and unfused implementations. FP32 comparisons were close (maximum reported relative L2 approximately 1.21e-5). This separates a pre-existing precision concern from tree topology, but does not establish BF16 equivalence. No tolerances have been relaxed, and a precision fix is intentionally outside this feature stack.
Independent prerequisite: optional MTP-loss bypass in #7384. Gradient preservation for the async RL consumer is separate in #7383. The forthcoming shared-prefix draft will carry the model-side parity tests and explicit limitations for review.
Contributor guide
Research direction
Start by reviewing the forthcoming shared-prefix draft and the existing attention CP1/CP2 and TP2 edge-projection parity checks referenced in the issue. Use #7384 and #7383 to understand prerequisite and separate work, then verify exact loss ownership, CP behavior, supported-operation rejection, and the stated BF16 limitation without relaxing tolerances.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100