| import torch | |
| from verl.utils.torch_functional import logprobs_from_logits_v2 | |
| def test_hint_next_token_labels_match_causal_logits(): | |
| batch_size = 3 | |
| sequence_length = 8 | |
| vocab_size = 16 | |
| input_ids = torch.randint(vocab_size, (batch_size, sequence_length)) | |
| responses = input_ids[:, 1:] | |
| logits = torch.randn(batch_size, sequence_length, vocab_size) | |
| causal_logits = logits[:, -responses.size(1) - 1:-1, :] | |
| assert causal_logits.shape[:-1] == responses.shape | |
| assert logprobs_from_logits_v2(causal_logits, responses).shape == responses.shape | |