[BREAKING][rollout, trainer, algo] feat: comprehensive rollout importance sampling implementation

ROLLOUT_IS_MIGRATION.md

@@ -0,0 +1,320 @@
+# Rollout Importance Sampling - Migration Guide
+
+This document provides a comprehensive overview of the Rollout Importance Sampling (IS) implementation merged from aiic_verl into verl.
+
+## References
+
+- **When Speed Kills Stability**: https://yingru.notion.site/When-Speed-Kills-Stability-271211a558b7808d8b12d403fd15edda
+- **Off-policy RL**: https://fengyao.notion.site/off-policy-rl
+
+## Overview
+
+Rollout Importance Sampling corrects for distribution mismatch between:
+- **Rollout policy**: e.g., vLLM with BFloat16
+- **Training policy**: e.g., FSDP with FP32
+
+This mismatch can lead to biased gradient estimates and unstable training. Rollout IS applies importance sampling weights to correct these biases.
+
+## What Changed
+
+### **Removed (Old Implementation)**
+
+```yaml
+# Old TIS configuration (REMOVED)
+actor:
+  tis_imp_ratio_cap: 2.0  # ❌ No longer supported
+```
+
+The old implementation:
+- Only supported token-level truncate mode
+- Had no metrics tracking
+- Lacked numerical stability safeguards
+- No configurability for different scenarios
+
+### **Added (New Implementation)**
+
+```yaml
+# New Rollout IS configuration
+algorithm:
+  rollout_is: true
+  rollout_is_threshold: 2.0
+  rollout_is_threshold_lower: null  # Auto-reciprocal
+  rollout_is_level: token
+  rollout_is_mode: truncate
+  rollout_is_veto_threshold: 1e-4
+
+# REQUIRED: Enable log prob calculation
+actor_rollout_ref:
+  rollout:
+    calculate_log_probs: true
+```
+
+The new implementation:
+- ✅ Three aggregation levels: token, sequence, geometric
+- ✅ Two bounding modes: truncate, clip
+- ✅ Dual threshold support (upper/lower)
+- ✅ Veto mechanism for catastrophic outliers
+- ✅ 30+ comprehensive metrics
+- ✅ Log-space computation for numerical stability
+- ✅ Memory-efficient implementation
+
+## Files Modified
+
+### **Core Implementation**
+
+1. **NEW**: [verl/trainer/ppo/mismatch_helper.py](verl/trainer/ppo/mismatch_helper.py)
+   - Contains `compute_rollout_importance_weights()` - main function
+   - Contains `compute_is_metrics()` - comprehensive metrics
+
+2. **MODIFIED**: [verl/trainer/ppo/core_algos.py](verl/trainer/ppo/core_algos.py#L962-991)
+   - Replaced old TIS implementation (lines 962-967)
+   - Added new rollout IS with metrics support
+
+3. **MODIFIED**: [verl/workers/actor/dp_actor.py](verl/workers/actor/dp_actor.py)
+   - Updated to use `rollout_is_threshold` instead of `tis_imp_ratio_cap`
+   - Collects and logs all rollout IS metrics
+
+### **Configuration Files**
+
+4. **MODIFIED**: [verl/trainer/config/algorithm.py](verl/trainer/config/algorithm.py#L95-100)
+   - Added 6 new rollout IS parameters to `AlgoConfig`
+
+5. **MODIFIED**: [verl/workers/config/actor.py](verl/workers/config/actor.py#L110-115)
+   - Added 6 new rollout IS parameters to `ActorConfig`
+
+6. **MODIFIED**: [verl/trainer/config/actor/actor.yaml](verl/trainer/config/actor/actor.yaml#L77-89)
+   - Added rollout IS configuration section
+
+7. **MODIFIED**: [verl/trainer/config/ppo_trainer.yaml](verl/trainer/config/ppo_trainer.yaml#L116-133)
+   - Added rollout IS to algorithm config
+
+### **Documentation**
+
+8. **MODIFIED**: [docs/examples/config.rst](docs/examples/config.rst)
+   - Updated actor config with rollout IS parameters
+   - Updated algorithm config with rollout IS parameters
+   - Added detailed parameter descriptions
+
+### **Example Scripts**
+
+9. **MODIFIED**: [recipe/dapo/run_dapo_qwen2.5_32b_tis.sh](recipe/dapo/run_dapo_qwen2.5_32b_tis.sh)
+   - Updated from `tis_imp_ratio_cap` to rollout IS parameters
+   - Added comprehensive comments
+
+10. **NEW**: [examples/rollout_importance_sampling/README.md](examples/rollout_importance_sampling/README.md)
+    - Comprehensive guide with usage patterns
+    - Troubleshooting section
+    - Performance considerations
+
+11. **NEW**: [examples/rollout_importance_sampling/run_with_rollout_is.sh](examples/rollout_importance_sampling/run_with_rollout_is.sh)
+    - Basic example with token-level truncate
+
+### **Tests**
+
+12. **NEW**: [test_rollout_is.py](test_rollout_is.py)
+    - Unit tests for rollout IS functionality
+
+13. **NEW**: [tests/trainer/ppo/test_rollout_is_integration.py](tests/trainer/ppo/test_rollout_is_integration.py)
+    - Integration tests with PPO
+
+## Configuration Parameters
+
+### `algorithm.rollout_is` (bool)
+Enable/disable IS correction. Default: `False`
+
+### `algorithm.rollout_is_threshold` (float or null)
+Upper threshold for IS weights. Set to `null` to disable IS completely.
+- Token level: 1.5 - 5.0
+- Sequence level: 2.0 - 10.0
+- Geometric level: 1.0002 - 1.001
+
+### `algorithm.rollout_is_threshold_lower` (float or null)
+Lower threshold for IS weights. If `null`, defaults to 1/upper (reciprocal).
+
+### `algorithm.rollout_is_level` (str)
+Aggregation level for IS weights:
+- `"token"`: Per-token ratios (biased)
+- `"sequence"`: Product of ratios (unbiased)
+- `"geometric"`: Geometric mean (experimental)
+
+### `algorithm.rollout_is_mode` (str)
+Bounding mode:
+- `"truncate"`: Cap weights at upper threshold only
+- `"clip"`: Zero out weights outside [lower, upper]
+
+### `algorithm.rollout_is_veto_threshold` (float)
+Per-token veto threshold. If any token ratio < this, entire sequence is rejected.
+Default: `1e-4` (ratio 10,000x off)
+
+## Migration Steps
+
+### Step 1: Update Your Configuration
+
+**Before (Old):**
+```yaml
+actor_rollout_ref:
+  actor:
+    tis_imp_ratio_cap: 2.0
+  rollout:
+    calculate_log_probs: true
+```
+
+**After (New):**
+```yaml
+algorithm:
+  rollout_is: true
+  rollout_is_threshold: 2.0
+  rollout_is_level: token
+  rollout_is_mode: truncate
+
+actor_rollout_ref:
+  rollout:
+    calculate_log_probs: true  # Still required!
+```
+
+### Step 2: Monitor New Metrics
+
+Add monitoring for these key metrics (all prefixed with `actor/`):
+
+**Health Indicators:**
+- `rollout_is_mean`: Mean IS weight across sequences
+- `rollout_is_eff_sample_size`: Effective sample size after weighting
+- `rollout_is_veto_fraction`: Fraction of sequences vetoed
+
+**Distribution Metrics:**
+- `rollout_is_max`, `rollout_is_min`
+- `rollout_is_std`
+- `rollout_is_p50`, `rollout_is_p95`, `rollout_is_p99`
+
+**Diagnostic Metrics:**
+- `rollout_is_ratio_fraction_high`
+- `rollout_is_ratio_fraction_low`
+- `rollout_is_catastrophic_token_fraction`
+
+**Mismatch Metrics (Training vs Rollout Policy):**
+- `mismatch_training_ppl`: Perplexity of training policy
+- `mismatch_rollout_ppl`: Perplexity of rollout policy
+- `mismatch_ppl_ratio`: Ratio of training/rollout PPL
+- `mismatch_kl`: KL divergence KL(π_rollout || π_training)
+- `mismatch_k3_kl`: K3 KL estimator
+- `mismatch_log_ppl_diff`: Log perplexity difference
+
+### Step 3: Test Your Training
+
+Start with the basic token-level truncate configuration:
+```bash
+bash examples/rollout_importance_sampling/run_with_rollout_is.sh
+```
+
+Monitor metrics for 1-2 epochs before adjusting parameters.
+
+## Configuration Examples
+
+### Example 1: Token-level with Truncate
+```yaml
+algorithm:
+  rollout_is: true
+  rollout_is_threshold: 2.0
+  rollout_is_level: token
+  rollout_is_mode: truncate
+```
+
+### Example 2: Geometric Mean with Clip
+```yaml
+algorithm:
+  rollout_is: true
+  rollout_is_threshold: 1.0002
+  rollout_is_threshold_lower: 0.9998
+  rollout_is_level: geometric
+  rollout_is_mode: clip
+```
+
+### Example 3: Wider Threshold with Clip
+```yaml
+algorithm:
+  rollout_is: true
+  rollout_is_threshold: 3.0
+  rollout_is_threshold_lower: 0.33
+  rollout_is_level: token
+  rollout_is_mode: clip
+  rollout_is_veto_threshold: 1e-5
+```
+
+### Example 4: Asymmetric Thresholds
+```yaml
+algorithm:
+  rollout_is: true
+  rollout_is_threshold: 5.0
+  rollout_is_threshold_lower: 0.8
+  rollout_is_level: token
+  rollout_is_mode: clip
+```
+
+## Troubleshooting
+
+### Issue: High variance in IS weights
+**Symptoms:** `rollout_is_std` > 1.0, `rollout_is_eff_sample_size` < 0.3
+
+**Solutions:**
+1. Switch from `sequence` to `geometric` level
+2. Tighten thresholds
+3. Verify rollout and training aren't too different
+
+### Issue: Too many sequences vetoed
+**Symptoms:** `rollout_is_veto_fraction` > 0.1
+
+**Solutions:**
+1. Relax veto threshold: `rollout_is_veto_threshold: 1e-3`
+2. Check for numerical issues in log prob computation
+3. Verify policies aren't completely different
+
+### Issue: Mean IS weight far from 1.0
+**Symptoms:** `rollout_is_mean` < 0.5 or > 2.0
+
+**Solutions:**
+1. Verify `calculate_log_probs=True` is set
+2. Check rollout_log_probs are correctly passed
+3. Check for systematic bias
+
+## Performance Impact
+
+- **Memory overhead**: ~1% of model memory
+- **Computational overhead**: 1-3% depending on level
+- **Training stability**: Significantly improved when mismatch exists
+
+## Backward Compatibility
+
+**The old `tis_imp_ratio_cap` parameter is completely removed.** There is no backward compatibility mode.
+
+All scripts and configurations must be updated to use the new rollout IS parameters.
+
+## Testing
+
+Run the test suite to verify everything works:
+
+```bash
+# Basic unit tests
+python test_rollout_is.py
+
+# Integration tests (if pytest is available)
+pytest tests/trainer/ppo/test_rollout_is_integration.py -v
+```
+
+Expected output: All tests pass ✓
+
+## Additional Resources
+
+- **Implementation**: [verl/trainer/ppo/mismatch_helper.py](verl/trainer/ppo/mismatch_helper.py)
+- **Examples**: [examples/rollout_importance_sampling/](examples/rollout_importance_sampling/)
+- **DAPO Example**: [recipe/dapo/run_dapo_qwen2.5_32b_tis.sh](recipe/dapo/run_dapo_qwen2.5_32b_tis.sh)
+
+## Summary
+
+The new Rollout Importance Sampling implementation provides:
+- ✅ More robust handling of distribution mismatch
+- ✅ Better numerical stability
+- ✅ Comprehensive metrics for monitoring
+- ✅ Flexibility for different scenarios
+- ✅ Memory-efficient computation
+
+Migration is straightforward: replace `tis_imp_ratio_cap` with the new `rollout_is_*` parameters in the `algorithm` config section.

docs/examples/config.rst

@@ -118,7 +118,13 @@ Actor/Rollout/Reference Policy
       clip_ratio: 0.2
       entropy_coeff: 0.0
       use_kl_loss: False # True for GRPO
-      tis_imp_ratio_cap: -1 # set to positive values for Truncated Importance Sampling (requires setting `rollout.calculate_log_probs` as True)
+      # Rollout Importance Sampling (corrects distribution mismatch between rollout and training)
+      rollout_is: False # Enable IS correction
+      rollout_is_threshold: null # Upper threshold for IS weights (null to disable)
+      rollout_is_threshold_lower: null # Lower threshold (null = auto 1/upper)
+      rollout_is_level: token # Aggregation: token/sequence/geometric
+      rollout_is_mode: truncate # Bounding: truncate/clip
+      rollout_is_veto_threshold: 1e-4 # Catastrophic outlier threshold
       use_torch_compile: True # False to disable torch compile
       kl_loss_coef: 0.001 # for grpo
       kl_loss_type: low_var_kl # for grpo
@@ -498,6 +504,13 @@ Algorithm
        kl_coef: 0.005
        horizon: 10000
        target_kl: 0.1
+     # Rollout Importance Sampling
+     rollout_is: False
+     rollout_is_threshold: null
+     rollout_is_threshold_lower: null
+     rollout_is_level: token
+     rollout_is_mode: truncate
+     rollout_is_veto_threshold: 1e-4
 
 - ``gamma``: discount factor
 - ``lam``: Trade-off between bias and variance in the GAE estimator
@@ -510,6 +523,13 @@ Algorithm
   - ``kl_coef``: The (initial) coefficient of in-reward kl_penalty. Default is 0.001.
   - ``type``: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController.
   - ``horizon`` and ``target_kl``: See source code of AdaptiveKLController for details.
+- ``rollout_is``: Whether to enable rollout importance sampling correction. Default is False.
+- ``rollout_is_threshold``: Upper threshold for IS weights. Set to ``null`` to disable IS completely.
+- ``rollout_is_threshold_lower``: Lower threshold for IS weights. If ``null``, defaults to reciprocal of upper (1/upper).
+- ``rollout_is_level``: Aggregation level: ``token`` (biased), ``sequence`` (unbiased), or ``geometric`` (experimental).
+- ``rollout_is_mode``: Bounding mode: ``truncate`` (cap upper only) or ``clip`` (zero outside bounds).
+- ``rollout_is_veto_threshold``: Per-token veto threshold for catastrophic outliers. Default is 1e-4.
+  Note: Rollout IS requires setting ``actor_rollout_ref.rollout.calculate_log_probs=True``.
 
 Trainer
 ~~~~~~~