| import torch |
|
|
|
|
| @torch.no_grad() |
| def test_rope(lit_llama, orig_llama) -> None: |
| torch.manual_seed(1) |
|
|
| bs, seq_len, n_head, n_embed = 1, 6, 2, 8 |
| x = torch.randint(0, 10000, size=(bs, seq_len, n_head, n_embed // n_head)).float() |
|
|
| freqs_cis = orig_llama.precompute_freqs_cis(n_embed // n_head, seq_len) |
| llama_rope_cache = lit_llama.build_rope_cache(seq_len, n_embed // n_head, dtype=x.dtype, device=x.device) |
| torch.testing.assert_close(freqs_cis, torch.view_as_complex(llama_rope_cache)) |
|
|
| llama_x_rope = lit_llama.apply_rope(x, llama_rope_cache) |
| orig_llama_x_rope, _ = orig_llama.apply_rotary_emb(x, x, freqs_cis) |
| torch.testing.assert_close(llama_x_rope, orig_llama_x_rope) |
|
|
