Refactoring training inference importance sampling with seqeunce/geometry level

[zhaochenyang20]

Thanks so much to the contribution of this paper When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch

This PR refactors the Importance Sampling (IS) functionality, replacing the legacy --use-tis parameter with a more flexible --use-train-infer-is parameter system. We introduce multiple aggregation levels and processing modes to handle training-inference mismatch problems better.

1. Parameter System Refactoring

• Removed legacy parameters: Deleted --use-tis, --tis-clip, --tis-clip-low parameters
• New parameters:
  • --use-train-infer-is: Enable training-inference importance sampling
  • --train-infer-is-level: Aggregation level (token/sequence/geometric)
  • --train-infer-is-mode: Processing mode (truncate/mask/clip)
  • --train-infer-is-lower-bound/--train-infer-is-upper-bound: Weight bounds
  • --train-infer-is-veto-threshold: Catastrophic token threshold

2. Aggregation Levels

Token Level (default):

• Computes importance weights independently for each token
• Formula: w_i = exp(log π_train(x_i) - log π_rollout(x_i))
• Characteristics: Biased but computationally simple, suitable for most scenarios

Sequence Level:

• Uses the product of all token weights as the sequence weight
• Formula: w_seq = exp(Σ(log π_train(x_i) - log π_rollout(x_i)))
• Characteristics: Unbiased but high variance, suitable for sequence-level optimization

Geometric Level:

• Uses geometric mean to compute sequence weights
• Formula: w_seq = exp(mean(log π_train(x_i) - log π_rollout(x_i)))
• Characteristics: Biased but low variance, balances bias and variance

3. Processing Modes

Truncate Mode (TIS):

• Clips weights exceeding the upper bound to the upper bound
• Maintains original TIS behavior, suitable for variance control

Mask Mode (MIS):

• Sets weights outside [lower, upper] range to zero
• More aggressive filtering strategy, suitable for handling extreme mismatches

Clip Mode (CIS):

• Constrains weights within [lower, upper] range
• Balanced truncation strategy

4. Others

• Catastrophic token detection: Detects and filters sequences containing catastrophic tokens via --train-infer-is-veto-threshold
• monitoring metrics: Added training/inference perplexity, KL divergence, K3 KL estimator, and more

5. Usage

# Using geometric mean + mask mode
--use-train-infer-is \
--train-infer-is-level geometric \
--train-infer-is-mode mask \
--train-infer-is-lower-bound 0.5 \
--train-infer-is-upper-bound 2.0 \
--train-infer-is-veto-threshold 1e-3

[zhaochenyang20]

/gemini review

[yitianlian]

I think it might have a problem when cp>1.

[yitianlian]

I think it might have a problem when cp>1. Because the loss mask is not split by cp_size, while logp is split by cp_size. You can reuse the implementation of sum_of_sample_mean in cp_utils.