diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 0d3146533..b00ea0956 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -63,6 +63,12 @@ def fused_linear_cross_entropy_forward( # for ex: BT = 4096*4, V = 32000, H = 4096 ==> inc_factor = 8, chunk_size = 2048 BT, H = _input.shape V = weight.shape[0] + assert V > 0, ( + f"weight must have a non-empty vocab dimension, got weight.shape={tuple(weight.shape)}. " + "This usually means the weight tensor has not been materialized yet, e.g. when using " + "DeepSpeed ZeRO-3 and the parameter is accessed directly instead of through the module " + "forward (which triggers the all-gather). Gather the parameter before calling this function." + ) BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) inc_factor = triton.cdiv(V, H) # (V + H - 1) // H diff --git a/test/transformers/test_fused_linear_cross_entropy.py b/test/transformers/test_fused_linear_cross_entropy.py index 380870e0b..5cc390206 100644 --- a/test/transformers/test_fused_linear_cross_entropy.py +++ b/test/transformers/test_fused_linear_cross_entropy.py @@ -8,6 +8,7 @@ from test.utils import set_seed from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction +from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_forward from liger_kernel.transformers.functional import CrossEntropyOutput from liger_kernel.transformers.functional import liger_fused_linear_cross_entropy from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss @@ -1092,3 +1093,22 @@ def liger(): assert_verbose_allclose(lig_gw, ref_gw, atol=atol, rtol=rtol) if bias: assert_verbose_allclose(lig_gb, ref_gb, atol=atol, rtol=rtol) + + +def test_empty_weight_raises_clear_error(): + """ + Regression test for https://github.com/linkedin/Liger-Kernel/issues/767 + + When `weight` has a vocab dimension of 0 (e.g. accessing a DeepSpeed ZeRO-3 + partitioned parameter directly, before it has been gathered), the chunking + math used to divide by `inc_factor = cdiv(V, H)`, which is 0 when V is 0. + This raised a cryptic `ZeroDivisionError` deep inside triton. We now raise a + clear, actionable error instead. + """ + H = 16 + weight = torch.randn(0, H) + _input = torch.randn(4, H, requires_grad=True) + target = torch.randint(0, 10, (4,)) + + with pytest.raises(AssertionError, match="non-empty vocab dimension"): + fused_linear_cross_entropy_forward(_input, weight, target)