|
4 | 4 | import torch |
5 | 5 | from typing import Optional |
6 | 6 | from transformer_engine.pytorch.router import ( |
| 7 | + RoutingMapFormat, |
7 | 8 | fused_topk_with_score_function, |
8 | 9 | fused_compute_score_for_moe_aux_loss, |
9 | 10 | fused_moe_aux_loss, |
@@ -458,6 +459,147 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk): |
458 | 459 | torch.testing.assert_close(probs.grad, probs_clone.grad, atol=atol, rtol=rtol) |
459 | 460 |
|
460 | 461 |
|
| 462 | +# ============================================================================= |
| 463 | +# Test: routing_map BITMAP_U8 vs BYTEMAP parity (fwd + bwd) |
| 464 | +# Mirrors tests/jax/test_fused_router.py::test_topk_bitmap_vs_bytemap. |
| 465 | +# ============================================================================= |
| 466 | + |
| 467 | + |
| 468 | +def _bytemap_to_bitmap_u8(bytemap: torch.Tensor) -> torch.Tensor: |
| 469 | + """Reference packer: bool[T, E] -> uint8[T, ceil(E/8)] LSB-first. |
| 470 | +
|
| 471 | + Matches numpy.packbits(..., bitorder='little'), which is what the JAX-side |
| 472 | + parity test uses. |
| 473 | + """ |
| 474 | + flat = bytemap.to(torch.uint8).cpu().numpy() |
| 475 | + import numpy as np |
| 476 | + |
| 477 | + return torch.from_numpy(np.packbits(flat, axis=-1, bitorder="little")).to(bytemap.device) |
| 478 | + |
| 479 | + |
| 480 | +@pytest.mark.parametrize("dtype", [torch.float32]) |
| 481 | +@pytest.mark.parametrize( |
| 482 | + "num_tokens,num_experts,topk", |
| 483 | + [(128, 32, 4), (256, 128, 8), (256, 130, 8), (128, 1024, 16)], |
| 484 | +) |
| 485 | +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) |
| 486 | +def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function): |
| 487 | + """fused_topk_with_score_function should produce identical probs and an |
| 488 | + LSB-packed bitmap routing_map when routing_map_format=BITMAP_U8, and |
| 489 | + backward gradients should match the bytemap path exactly.""" |
| 490 | + if topk >= num_experts: |
| 491 | + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") |
| 492 | + if score_function in ("sigmoid", "sqrtsoftplus"): |
| 493 | + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 |
| 494 | + logits = ( |
| 495 | + torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 |
| 496 | + ) |
| 497 | + logits = logits.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) |
| 498 | + else: |
| 499 | + logits = ( |
| 500 | + torch.arange( |
| 501 | + -num_tokens * num_experts // 2, |
| 502 | + num_tokens * num_experts // 2, |
| 503 | + device="cuda", |
| 504 | + dtype=dtype, |
| 505 | + ) |
| 506 | + * 1e-4 |
| 507 | + ) |
| 508 | + logits = logits.view(num_tokens, num_experts) |
| 509 | + |
| 510 | + logits_byte = logits.detach().clone().requires_grad_(True) |
| 511 | + logits_bit = logits.detach().clone().requires_grad_(True) |
| 512 | + |
| 513 | + probs_byte, routing_map_byte = fused_topk_with_score_function( |
| 514 | + logits=logits_byte, |
| 515 | + topk=topk, |
| 516 | + use_pre_softmax=False, |
| 517 | + num_groups=None, |
| 518 | + group_topk=None, |
| 519 | + scaling_factor=None, |
| 520 | + score_function=score_function, |
| 521 | + expert_bias=None, |
| 522 | + routing_map_format=RoutingMapFormat.BYTEMAP, |
| 523 | + ) |
| 524 | + probs_bit, routing_map_bit = fused_topk_with_score_function( |
| 525 | + logits=logits_bit, |
| 526 | + topk=topk, |
| 527 | + use_pre_softmax=False, |
| 528 | + num_groups=None, |
| 529 | + group_topk=None, |
| 530 | + scaling_factor=None, |
| 531 | + score_function=score_function, |
| 532 | + expert_bias=None, |
| 533 | + routing_map_format=RoutingMapFormat.BITMAP_U8, |
| 534 | + ) |
| 535 | + |
| 536 | + assert probs_byte.dtype == probs_bit.dtype |
| 537 | + torch.testing.assert_close(probs_byte, probs_bit, atol=0.0, rtol=0.0) |
| 538 | + |
| 539 | + expected_shape = (num_tokens, (num_experts + 7) // 8) |
| 540 | + assert routing_map_bit.shape == expected_shape, ( |
| 541 | + f"Bitmap shape {tuple(routing_map_bit.shape)} != {expected_shape}" |
| 542 | + ) |
| 543 | + assert routing_map_bit.dtype == torch.uint8 |
| 544 | + assert routing_map_byte.dtype == torch.bool |
| 545 | + |
| 546 | + packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) |
| 547 | + torch.testing.assert_close(routing_map_bit, packed_expected, atol=0, rtol=0) |
| 548 | + |
| 549 | + # Backward parity: grad of probs.sum() must be bit-identical across formats. |
| 550 | + probs_byte.sum().backward() |
| 551 | + probs_bit.sum().backward() |
| 552 | + torch.testing.assert_close(logits_byte.grad, logits_bit.grad, atol=0.0, rtol=0.0) |
| 553 | + |
| 554 | + |
| 555 | +@pytest.mark.parametrize("dtype", [torch.float32]) |
| 556 | +@pytest.mark.parametrize( |
| 557 | + "num_tokens,num_experts,topk", |
| 558 | + [(128, 32, 4), (256, 128, 8), (256, 130, 8)], |
| 559 | +) |
| 560 | +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) |
| 561 | +def test_score_for_aux_loss_bitmap_vs_bytemap( |
| 562 | + dtype, num_tokens, num_experts, topk, score_function |
| 563 | +): |
| 564 | + """fused_compute_score_for_moe_aux_loss: bitmap routing_map must equal |
| 565 | + LSB-packed bytemap; scores must be bit-identical across formats.""" |
| 566 | + if topk >= num_experts: |
| 567 | + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") |
| 568 | + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 |
| 569 | + logits = torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 |
| 570 | + logits = logits.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) |
| 571 | + |
| 572 | + logits_byte = logits.detach().clone().requires_grad_(True) |
| 573 | + logits_bit = logits.detach().clone().requires_grad_(True) |
| 574 | + |
| 575 | + routing_map_byte, scores_byte = fused_compute_score_for_moe_aux_loss( |
| 576 | + logits=logits_byte, |
| 577 | + topk=topk, |
| 578 | + score_function=score_function, |
| 579 | + routing_map_format="bytemap", |
| 580 | + ) |
| 581 | + routing_map_bit, scores_bit = fused_compute_score_for_moe_aux_loss( |
| 582 | + logits=logits_bit, |
| 583 | + topk=topk, |
| 584 | + score_function=score_function, |
| 585 | + routing_map_format="bitmap_u8", |
| 586 | + ) |
| 587 | + |
| 588 | + torch.testing.assert_close(scores_byte, scores_bit, atol=0.0, rtol=0.0) |
| 589 | + |
| 590 | + expected_shape = (num_tokens, (num_experts + 7) // 8) |
| 591 | + assert routing_map_bit.shape == expected_shape |
| 592 | + assert routing_map_bit.dtype == torch.uint8 |
| 593 | + assert routing_map_byte.dtype == torch.bool |
| 594 | + packed_expected = _bytemap_to_bitmap_u8(routing_map_byte) |
| 595 | + torch.testing.assert_close(routing_map_bit, packed_expected, atol=0, rtol=0) |
| 596 | + |
| 597 | + # Backward parity through scores. |
| 598 | + scores_byte.sum().backward() |
| 599 | + scores_bit.sum().backward() |
| 600 | + torch.testing.assert_close(logits_byte.grad, logits_bit.grad, atol=0.0, rtol=0.0) |
| 601 | + |
| 602 | + |
461 | 603 | def profile_topk_softmax( |
462 | 604 | dtype, |
463 | 605 | num_tokens, |
|
0 commit comments