gspo: GSPO loss + DeepSpeed parity fixes (loss/grad divisors, SDP, fp32_lm_head, docs_per_step, temperature)

[bigximik]

Summary

This PR adds GSPO loss to fast-LLM along with a suite of supporting fixes that together achieve full metric and training-trajectory parity with DeepSpeed's GRPO/GSPO implementation. Targets the grpo-metrics branch. Six logical units:

1. GSPO loss (sequence-level IS-ratio clipping)

Implements GSPO as an alternative policy-gradient loss alongside the existing per-token GRPO clipping. Controlled via LanguageModelGRPOLossConfig.policy_loss = "gspo".

• New fused_gspo_loss_forward_backward kernel: computes per-segment geometric-mean log-ratio R_s, clips at [1−ε_low, 1+ε_high], and applies R_s × A_s as a uniform per-token gradient within each segment. An all_reduce(SUM) over sequence-data-parallel ranks aggregates (lrn_sum, adv_sum, tok_count) before clipping so the ratio is correct under sequence parallelism.
• New document_index data field and LanguageModelKwargs.document_index kwarg constant to route per-token segment membership through the data pipeline.
• 8 unit tests in tests/layers/test_gspo_loss.py (single-segment, packed sequences, ratio=1 equivalence, clipping, masking, SDP mock, gradient finite-diff, independence from per-token metrics).

2. Dynamic docs_per_step accumulation

Replaces static depth_first_micro_batches with a runtime document-count target — matching DeepSpeed's gradient_accumulation_passes semantics for RL (where each microbatch holds one rollout).

• ScheduleConfig.docs_per_step: when >0, Trainer._prefetch_to_doc_target fetches microbatches one at a time, all-reduces the per-microbatch document count, and stops once the global total ≥ docs_per_step. The final step total is broadcast to all inputs so the normalisation denominator is consistent.
• Trainer._get_or_build_schedule builds and caches a per-N Schedule with _depth_first_override = N // breadth_first_micro_batches, so the existing schedule machinery is reused without changes to the runner.
• Schedule._eff_{depth_first,sequential,num_inputs} properties expose the effective values for a given override.
• 13 unit tests in tests/layers/test_docs_per_step.py.

3. normalize_by_documents

Adds a normalize_by_documents flag to LanguageModelGRPOLossConfig. When True, both the GRPO and GSPO paths divide the loss by num_documents_in_batch (the step-level rollout count) rather than the token count. Matches DeepSpeed's normalization where tokens_weights = 1 / batch_size.

4. Temperature scaling for IS ratio parity

Adds a temperature field to LanguageModelGRPOLossConfig. When set to match the actor's sampling temperature (e.g. 0.7), new log-probabilities are computed at the same temperature as the stored old log-probabilities from vLLM, so the IS ratio starts near 1.0 at step 0 instead of ~1.08. Implementation: _effective_logits_scale = logits_scale_factor / temperature, substituted at all three call-sites in _forward_backward. Default temperature=1.0 preserves existing behaviour exactly.

5. fp32_lm_head precision fix (matches vLLM's bf16_last_layer_fp32)

Adds a fp32_lm_head flag (default False) on LanguageModelHeadConfig. When True, the LM head's logits computation upcasts both input and weight to FP32 before the linear projection, matching vLLM's bf16_last_layer_fp32 quantization. This ensures the trainer computes log-probabilities at the same numerical precision as the actor's sampling, so new_logprobs ≈ old_logprobs at step 0 (IS ratio at training start ≈ 1.0, not artificially inflated by precision mismatch).

• Commit d8cb9ef5: introduces the flag, upcasts input/weight, casts back to BF16 before downstream consumption.
• Commit 0f90f20b: fixes the gradient flow when fp32_lm_head=True. The detached FP32 weight copy has requires_grad=False, which makes output_parallel_linear_backward skip writing to the original weight's grad_buffer. We restore the FSDP gradient contract by computing grad_weight = grad.t() @ saved_input explicitly and accumulating into the BF16 param's grad_buffer via accumulate_gradient.

6. Decoupled loss/gradient divisors and SDP loss double-counting fix

Even with normalize_by_documents=true, fast-LLM's reported grad_norm was ~1024× larger than DeepSpeed's, causing the default gradient_norm_clipping=0.3 to over-clip by ~500× and making training ~10 reward points slower than DS GSPO at the same step count. Two issues, fixed in commit 557a3c4c:

Asymmetric loss/gradient scaling in DS:

• DS loss reported uses /batch_size once (via tokens_weights = 1/batch_size, pipelinerl/finetune/rl/__init__.py:246).
• DS gradient buffer has an ADDITIONAL /(gas × world_size) factor from scale_wrt_gas=True in engine.backward() (deepspeed/runtime/engine.py:1995-1996) and tensor.div_(world_sz) in reduce_scatter_coalesced (deepspeed/runtime/comm/coalesced_collectives.py:124).
• For samples_per_microbatch=1 (PipelineRL standard), gas × world_size = batch_size, so the gradient buffer effectively has 1/batch_size² while the loss metric has 1/batch_size.

Fast-LLM cancels DS's /(gas × world_size) factor structurally via grad_output = data_parallel × grad_scale (runner.py:318) interacting with FSDP's RS-AVG over data_parallel ranks (fsdp.py:396). So we need to apply the second 1/batch_size factor explicitly only to the gradient — keeping the loss metric matched to DS.

Fix: add a grad_divisor parameter to fused_gspo_loss_forward_backward, fused_grpo_loss_forward_backward, and triton_grpo_loss_forward_backward. When normalize_by_documents=true:

• loss divisor = num_documents_in_batch (matches DS rl/loss)
• gradient divisor = num_documents_in_batch² (matches DS grad_norm)

Independent of TP/PP/SDP/DP parallelism and microbatching schedule, because batch_size is invariant under all of them.

SDP loss double-counting:
After the SDP allreduce of lrn_sum/adv_sum/tok_sum in fused_gspo_loss_forward_backward, both SDP ranks compute IDENTICAL per-segment loss values. When LossDef.reduce SUMs across data_group (which includes SDP ranks), the loss metric is double-counted by sdp_size. The gradient is NOT double-counted — each SDP rank contributes gradient from its own LOCAL tokens, with different contributions for different tokens of the same segment.

Fix: divide loss by sdp_size when sdp_group is active. Gradient unaffected.

Verification

End-to-end 7B math run on 4 nodes, GSPO, gradient_norm_clipping=0.3 (default), normalize_by_documents=true, fp32_lm_head=true, temperature=0.7:

Metric	Before unit-6 fix	After unit-6 fix	DS GSPO reference
step 1 grad_norm	141 (1000× DS)	0.135	0.145
step 1 lm_head_loss	-13.7	~-1.7 magnitude	-1.7
step 1 clip_coeff	0.002 (severe over-clip)	1.000 (no clip)	no clip
step 50 newlp	trapped at -0.17	-0.103	-0.105

newlp trajectory tracks DS step-by-step. Both systems show same gradient-spike pattern during warmup ramp-up at steps 14-20 (DS step 16 grad_norm=6.365, fast-LLM step 15=9.005). Match within data variance.

Test plan

• pytest tests/layers/test_gspo_loss.py — GSPO unit tests pass
• pytest tests/layers/test_docs_per_step.py — docs_per_step unit tests pass
• pytest tests/layers/test_lm_losses.py — existing GRPO loss + per-token metrics tests unaffected (the metrics tests previously in test_grpo_metrics.py moved into this file on the base branch)
• End-to-end: 4-node Qwen2.5-7B math run with full config (docs_per_step=1024, temperature=0.7, normalize_by_documents=true, fp32_lm_head=true, default gradient_norm_clipping=0.3) — grad_norm matches DS at step 1, training trajectory matches DS step-by-step through step 50+ (ongoing run validates through step ~410).