Add GMPO to fused linear GRPO loss

[chinoll]

Summary

This PR adds native Geometric-Mean Policy Optimization (GMPO) support to LigerFusedLinearGRPOLoss through loss_type="gmpo".

GMPO clips each token's log importance ratio before aggregating the valid tokens with a sequence-level geometric mean. Integrating the objective into Liger's existing fused linear and chunked-loss path avoids materializing the full [batch, sequence, vocabulary] logits tensor while preserving the GMPO objective and gradients.

The official GMPO implementation is available at callsys/GMPO.

Details

Example:

loss_fn = LigerFusedLinearGRPOLoss(
    loss_type="gmpo",
    importance_sampling_level="token",
    epsilon_low=0.4,
    epsilon_high=0.4,
)

The implementation:

• applies advantage-dependent, one-sided clipping directly in log-ratio space;
• computes exp(masked_mean(clipped_log_ratio)) for each completion;
• reuses the existing per-sample reduction, K3 KL path, clipping metric, chunking, and fused linear backward infrastructure;
• supports optional KL regularization and bias-corrected KL;
• rejects sequence-level importance sampling because GMPO clips at token level before sequence aggregation;
• intentionally does not apply vllm_is_ratio or delta to the GMPO policy objective; and
• documents that epsilon_low and epsilon_high are GMPO log-space bounds [-epsilon_low, +epsilon_high], rather than the ratio-space bounds used by PPO-style losses.

The test reference implements GMPO independently with full PyTorch operations. Tests cover:

• FP32 and BF16;
• outer chunk_size values 1, 2, and 5;
• asymmetric clipping bounds and positive, negative, and zero advantages;
• beta=0, K3 KL, and bias-corrected KL;
• eager and compiled execution;
• input, LM-head weight, and bias gradients;
• exact clipping boundaries and adjacent torch.nextafter values;
• extreme active log-ratio clips and finite gradients;
• rejection of sequence-level importance sampling;
• the defined vllm_is_ratio and delta behavior; and
• existing GRPO and DAPO large-sequence regression paths.

Testing Done

The tested source content is commit 27eb5ce3ab3c56ad90b47e1a4201ae121e09b764.

Targeted correctness tests

The final source-only snapshot was tested on an NVIDIA H100 with:

CUDA_VISIBLE_DEVICES=7 PYTHONPATH="$ROOT/source/src" \
python -m pytest "$ROOT/source/test/chunked_loss/test_grpo_loss.py" \
  -k "gmpo or test_correctness_large_seq_exercises_chunking" -q

28 passed, 1586 deselected, 15 warnings in 5.39s

The command exited with code 0 and had no OOM, NaN, or hang. The exact-boundary test compares the forward loss, clipping metric, and log-probability gradient with rtol=0 and atol=0. The orthogonal tests cover BF16, compilation, KL/bias correction, and multiple outer chunk sizes.

Forward and backward numerical comparison

An additional H100 sweep compared the Liger implementation with a naive full-logits BF16 implementation using a Qwen-0.5B LM-head shape:

• hidden_size=896, vocab_size=151936;
• sequence lengths 512, 1024, 2048, 4096, 8192, 16384, 32768;
• batch sizes 1, 2, 4, 8, 16, 32, 64, 128;
• token-level GMPO, epsilon_low=epsilon_high=0.4, beta=0;
• no bias, deterministic seed 0; and
• compiled=False, with Liger chunk_size=batch_size.

All 56 shapes completed with Liger. The 35 shapes for which the naive implementation fit in memory were compared directly in both forward and backward; the other 21 shapes were Liger-only finite scans because the naive forward pass had already OOMed under the same setup.

Quantity	Worst observed difference
Forward loss, maximum absolute error	1.7881393e-07
Forward loss, maximum relative error	3.6046363e-07
Clipping metric, maximum absolute error	0
Hidden gradient, maximum absolute error	9.5367432e-07
Hidden gradient, maximum relative L2	0.0036441503
Hidden gradient, minimum cosine similarity	0.99999458
Weight gradient, maximum absolute error	6.1035156e-05
Weight gradient, maximum relative L2	0.0036656931
Weight gradient, minimum cosine similarity	0.99999453

No non-finite loss, metric, hidden gradient, or weight gradient was observed in either implementation for any completed run. All 21 larger Liger-only shapes were also finite, including batch_size=128, sequence_length=32768.

H100 memory comparison

The same 56-shape matrix measured peak allocated memory for the LM head plus GMPO loss. This measurement does not include the transformer trunk or optimizer state.

• Liger completed forward and backward for all 56/56 shapes.
• The naive baseline completed 35/56 shapes and OOMed during forward for 21/56.
• Under this fixed setup, all shapes with B*T <= 65,536 fit in the naive implementation, while all shapes with B*T >= 131,072 OOMed during naive forward.
• Across the 35 shapes completed by both implementations, the median full-training peak-memory reduction was 90.6%; the largest was 97.6%.
• At B=2, T=32768, the full peak decreased from 74.58 GiB to 1.77 GiB.
• At the largest tested shape, B=128, T=32768, Liger completed at a 29.44 GiB full peak while the naive implementation OOMed.
• The smallest case, B=1, T=512, is the exception: Liger used 1.34 GiB versus the naive implementation's 1.16 GiB (15.6% higher).

Batch	Sequence	Naive full peak	Liger full peak	Reduction
1	512	1.16 GiB	1.34 GiB	-15.6%
8	2048	15.39 GiB	1.44 GiB	90.6%
16	4096	58.35 GiB	1.77 GiB	97.0%
2	32768	74.58 GiB	1.77 GiB	97.6%
128	32768	OOM	29.44 GiB	N/A

Liger's fused chunked loss computes and stores gradients during its custom autograd forward. Consequently, its principal training-memory peak appears in the API's forward phase, while the naive implementation follows the conventional saved-graph/backward path. The full-peak values above use max(forward_peak, backward_peak) for each provider.

The repository's current benchmark_grpo_loss.py does not expose loss_type="gmpo", so these GMPO-specific measurements were run with an isolated benchmark rather than reporting the existing GRPO/GSPO benchmark as GMPO data.

• Hardware Type: NVIDIA H100 80GB HBM3
• Software: PyTorch 2.11.0+cu130
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[chinoll]

#1315

[chinoll]

@Tcc0403 Thanks for the review! I've made the requested changes.