Make GSPO loss length-proportional

[jlamypoirier]

Codex GPT-5 note:

Summary

• make GSPO apply segment loss once per labeled token instead of once per document
• keep per-document label counts for the geometric-mean ratio and advantage computation
• update the GSPO reference test to match the length-proportional loss

Rationale

PipelineRL DeepSpeed GSPO uses length-proportional sequence weighting when group_normalization=false: each segment contributes in proportion to its labeled-token count. Fast-LLM was still using mask / num_labels_in_seq as both the geometric-mean normalizer and the loss weight, making each document contribute uniformly regardless of length. This keeps num_labels_in_seq for the ratio/advantage means but uses the label mask as the loss/gradient weight.

Test

• FAST_LLM_TEST_RESULTS_PATH=/tmp/fast_llm_tests/gspo_length_proportional /Users/joel.lamy-poirier/Projects/Fast-LLM/venv/bin/python -m pytest -v -n 4 tests/layers/test_lm_losses.py -k gspo