| """Tests for completion logit slicing shared by GRPO and OPSD.""" | |
| import os | |
| import sys | |
| import torch | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from opsd_utils.opsd_loss import slice_student_completion_logits | |
| def test_slice_matches_legacy_grpo_path(): | |
| logits_to_keep = 4 | |
| full = torch.randn(2, 20, 8) | |
| legacy = full[:, -logits_to_keep - 1 :, :] | |
| legacy = legacy[:, :-1, :] | |
| legacy = legacy[:, -logits_to_keep:, :] | |
| assert torch.allclose(legacy, slice_student_completion_logits(full, logits_to_keep)) | |