From 97ea310df2e121c5cca1e64c73f875d9d4e2914e Mon Sep 17 00:00:00 2001 From: Haozhe Zhang Date: Tue, 23 Jun 2026 16:00:53 +0800 Subject: [PATCH] test(distillation): add value-level CPU tests for forward_kl_topk OPD loss The existing CPU coverage only asserts the overlap diagnostics; the loss value, student/teacher mass, the default-on log_prob_min_clamp branch, and the use_chunked_topk path had no value-level assertions. Pin them against a from-scratch forward-KL reference on tiny CPU tensors. --- ...test_forward_kl_topk_loss_values_on_cpu.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 tests/workers/test_forward_kl_topk_loss_values_on_cpu.py diff --git a/tests/workers/test_forward_kl_topk_loss_values_on_cpu.py b/tests/workers/test_forward_kl_topk_loss_values_on_cpu.py new file mode 100644 index 00000000000..54229afb17a --- /dev/null +++ b/tests/workers/test_forward_kl_topk_loss_values_on_cpu.py @@ -0,0 +1,105 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Value-level CPU tests for ``compute_forward_kl_topk`` (FSDP backend). + +The existing CPU coverage (``test_distillation_topk_symmetry_on_cpu``) only +asserts the two diagnostic outputs (``overlap_count`` / +``overlap_token_advantage``). The load-bearing outputs of the GKD-OPD loss -- +``distillation_losses`` (the objective), ``student_mass`` / ``teacher_mass``, +the default-on ``log_prob_min_clamp`` branch, and the ``use_chunked_topk`` +path -- had no value-level assertions, so a regression in ``kl_divergence`` or +the clamp/chunk handling would pass CI. These tests pin the loss math against +a from-scratch reference on tiny CPU tensors. +""" + +from types import SimpleNamespace + +import torch + +from verl.trainer.distillation.fsdp.losses import compute_forward_kl_topk + + +def _cfg(log_prob_min_clamp=None, use_chunked_topk=False): + return SimpleNamespace( + distillation_loss=SimpleNamespace( + log_prob_min_clamp=log_prob_min_clamp, + use_chunked_topk=use_chunked_topk, + chunked_topk_chunk_size=4096, + ) + ) + + +def _nested(rows: torch.Tensor) -> torch.Tensor: + # compute_forward_kl_topk expects jagged nested teacher tensors and reads + # ``.values().unsqueeze(0)``; one batch row of (seqlen, topk) reproduces a + # single-sequence micro-batch. + return torch.nested.nested_tensor([rows], layout=torch.jagged) + + +def _make_inputs(seqlen=4, vocab=10, topk=3, seed=0): + g = torch.Generator().manual_seed(seed) + student_logits = torch.randn(1, seqlen, vocab, generator=g) + teacher_log_probs = torch.log_softmax(torch.randn(seqlen, vocab, generator=g), dim=-1) + t_lp, t_ids = torch.topk(teacher_log_probs, topk, dim=-1) + return student_logits, t_lp, t_ids.long() + + +def _forward_kl_reference(student_logits, t_lp, t_ids, clamp=None): + s_at_t = torch.gather(torch.log_softmax(student_logits, dim=-1), dim=-1, index=t_ids.unsqueeze(0)) + t_lp_b = t_lp.unsqueeze(0) + if clamp is not None: + t_lp_b, s_at_t = t_lp_b.clamp_min(clamp), s_at_t.clamp_min(clamp) + loss = (t_lp_b.exp() * (t_lp_b - s_at_t)).sum(dim=-1) + return loss, s_at_t, t_lp_b + + +def test_forward_kl_topk_loss_matches_reference(): + student_logits, t_lp, t_ids = _make_inputs() + out = compute_forward_kl_topk(student_logits, _nested(t_lp), _nested(t_ids), _cfg(log_prob_min_clamp=None), "thd") + # student_mass / teacher_mass are taken on the *unclamped* log-probs. + s_at_t = torch.gather(torch.log_softmax(student_logits, dim=-1), dim=-1, index=t_ids.unsqueeze(0)) + ref_loss, _, _ = _forward_kl_reference(student_logits, t_lp, t_ids, clamp=None) + torch.testing.assert_close(out["distillation_losses"], ref_loss) + torch.testing.assert_close(out["student_mass"], s_at_t.exp().sum(dim=-1)) + torch.testing.assert_close(out["teacher_mass"], t_lp.unsqueeze(0).exp().sum(dim=-1)) + + +def test_log_prob_min_clamp_branch_changes_loss(): + # Force student log-probs far below the clamp at the teacher's top-k ids: + # a hard peak on id 0 drives log-prob ~ -60 at every other id. + seqlen, vocab, clamp = 2, 12, -10.0 + student_logits = torch.full((1, seqlen, vocab), -30.0) + student_logits[..., 0] = 30.0 + t_ids = torch.tensor([[1, 2, 3], [4, 5, 6]]) + t_lp = torch.log_softmax(torch.randn(seqlen, vocab), dim=-1).gather(dim=-1, index=t_ids) + + out = compute_forward_kl_topk( + student_logits, _nested(t_lp), _nested(t_ids.long()), _cfg(log_prob_min_clamp=clamp), "thd" + ) + ref_clamped, _, _ = _forward_kl_reference(student_logits, t_lp, t_ids, clamp=clamp) + ref_unclamped, _, _ = _forward_kl_reference(student_logits, t_lp, t_ids, clamp=None) + torch.testing.assert_close(out["distillation_losses"], ref_clamped) + # The clamp is on by default and materially changes the loss here. + assert not torch.allclose(ref_clamped, ref_unclamped) + + +def test_chunked_topk_matches_default_path(): + student_logits, t_lp, t_ids = _make_inputs(seqlen=6, vocab=16, topk=4, seed=1) + default = compute_forward_kl_topk( + student_logits, _nested(t_lp), _nested(t_ids), _cfg(use_chunked_topk=False), "thd" + ) + chunked = compute_forward_kl_topk(student_logits, _nested(t_lp), _nested(t_ids), _cfg(use_chunked_topk=True), "thd") + assert set(chunked) == set(default) + for key in default: + torch.testing.assert_close(chunked[key], default[key])