|
|
| import pytest |
| import torch |
|
|
| from fla.modules.grpo import fused_grpo_loss, grpo_loss_torch |
| from fla.utils import assert_close, device, device_torch_lib, is_nvidia_hopper |
|
|
|
|
| @pytest.mark.parametrize("B", [2]) |
| @pytest.mark.parametrize("T", [16, 1024, 4096]) |
| @pytest.mark.parametrize("V", [32000, 65536, 131072]) |
| @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| @pytest.mark.parametrize("inplace", [True, False]) |
| @pytest.mark.parametrize("repeat", [100]) |
| def test_fused_grpos(B: int, T: int, V: int, dtype: torch.dtype, inplace: bool, repeat: int): |
| device_torch_lib.manual_seed(42) |
| for i in range(repeat): |
| if not is_nvidia_hopper and T == 4096: |
| pytest.skip("Skip test for T=4096 on Intel Alchemist") |
|
|
| def get_random_ref_log_probs(logits, input_ids): |
| with torch.inference_mode(): |
| logits = logits[:, :-1] |
| per_token_logps = [] |
| for logits_row, input_ids_row in zip(logits, input_ids[:, -logits.size(1):], strict=False): |
| log_probs = torch.randn_like(logits_row).log_softmax(dim=-1) |
| token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1) |
| per_token_logps.append(token_log_prob) |
| device_torch_lib.empty_cache() |
| return torch.stack(per_token_logps) |
|
|
| logits = torch.randn(B, T + 1, V, device=device, dtype=dtype) |
| logits.requires_grad_(True) |
| advantages = torch.randn(B, device=device, dtype=torch.float32) |
| input_ids = torch.randint(0, V-1, (B, T + 64), device=device) |
| ref_logp = get_random_ref_log_probs(logits, input_ids) |
| beta = 0.04 |
| completion_mask = torch.ones(B, T, dtype=torch.int32, device=device) |
| completion_mask[::2, T//3: T//2] = 0 |
| save_kl = True |
|
|
| gold_logits = logits.detach().clone().float() |
| gold_logits.requires_grad_(True) |
| gold_ref_logp = ref_logp.clone().float() |
| device_torch_lib.empty_cache() |
| y1 = fused_grpo_loss(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl=save_kl, inplace=inplace) |
| y2 = grpo_loss_torch(gold_logits, gold_ref_logp, input_ids, advantages, beta, completion_mask, save_kl) |
| if save_kl: |
| y1, kl2 = y1 |
| y2, kl3 = y2 |
| assert (kl2-kl3).abs().max() < 1e-3 |
| dy = torch.randn_like(y1) * 10 |
| y1.backward(dy) |
| y2.backward(dy.float()) |
| assert (y1-y2).abs().max() < 1e-3 |
| assert_close(" dlogits", gold_logits.grad, logits.grad, 3e-3) |
|
|
