Claude Opus 4.8 note: surfaced during review of #555.
Both the GSPO loss kernel (fused_gspo_loss_forward_backward) and the GSPO diagnostic metrics (compute_gspo_metrics) in fast_llm/layers/language_model/loss/policy_gradient.py all-reduce their per-segment buffers over the sequence-data-parallel (sdp_group) and sequence-parallel (sp_group) groups, to reconstruct per-document/segment aggregates when the sequence is split across those ranks.
Gap: the distributed test harness (_run_lm_loss_distributed / _test_gspo_metrics in tests/layers/test_lm_losses.py) only ever drives the vocab-parallel group (it splits the vocab dimension). sdp_group/sp_group stay at their None defaults, so those all_reduce branches never execute under test. There is also no gspo loss subtest in _run_lm_loss_distributed at all — only grpo, grpo_metrics, and gspo_metrics.
Why it needs its own effort: exercising the SDP/SP reduction requires a sequence-split test path (partition tokens/documents across ranks along the sequence dimension, then verify the reduced segment aggregates match a single-rank reference), which is beyond the current 2-rank vocab-split harness.
Single-rank and vocab-parallel paths are tested and verified; this issue tracks the remaining SDP/SP coverage for both the GSPO loss and its metrics.
Claude Opus 4.8 note: surfaced during review of #555.
Both the GSPO loss kernel (
fused_gspo_loss_forward_backward) and the GSPO diagnostic metrics (compute_gspo_metrics) infast_llm/layers/language_model/loss/policy_gradient.pyall-reduce their per-segment buffers over the sequence-data-parallel (sdp_group) and sequence-parallel (sp_group) groups, to reconstruct per-document/segment aggregates when the sequence is split across those ranks.Gap: the distributed test harness (
_run_lm_loss_distributed/_test_gspo_metricsintests/layers/test_lm_losses.py) only ever drives the vocab-parallelgroup(it splits the vocab dimension).sdp_group/sp_groupstay at theirNonedefaults, so thoseall_reducebranches never execute under test. There is also nogspoloss subtest in_run_lm_loss_distributedat all — onlygrpo,grpo_metrics, andgspo_metrics.Why it needs its own effort: exercising the SDP/SP reduction requires a sequence-split test path (partition tokens/documents across ranks along the sequence dimension, then verify the reduced segment aggregates match a single-rank reference), which is beyond the current 2-rank vocab-split harness.
Single-rank and vocab-parallel paths are tested and verified; this issue tracks the remaining SDP/SP coverage for both the GSPO loss and its metrics.