File size: 2,588 Bytes
c335050 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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)
|