Fix ZeroDivisionError in fused linear cross entropy when weight vocab dim is 0

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

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)