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

[nileshpatil6]

Summary

Fixes #767. This reopens #1289, which I accidentally closed for good by deleting my fork during a repo cleanup (GitHub won't let a PR reopen once the head repo is gone). Same fix, rebased onto current main.

fused_linear_cross_entropy_forward computes chunking as:

V = weight.shape[0]
inc_factor = triton.cdiv(V, H)
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))

When weight has a vocab dim of 0, inc_factor is 0 and the next cdiv divides by zero. The user just sees ZeroDivisionError: integer division or modulo by zero raised from inside triton with nothing pointing at the real cause. That is what the reporter in #767 hit, and it comes up with DeepSpeed ZeRO-3 when a partitioned parameter is touched directly before it has been gathered.

This adds an assert before the chunking math with a message that names the actual problem and what to do about it, instead of letting it fall through to triton.

Details

Kept it as an assert to match the existing validation style at the top of the same function (return_z_loss, return_token_accuracy, return_predicted_tokens all assert). Happy to switch it to a ValueError if you would rather it not be strippable under -O.

Testing Done

Added test_empty_weight_raises_clear_error in test/transformers/test_fused_linear_cross_entropy.py. It runs on CPU since the assert fires before any kernel launch, so it does not need a GPU.

I do not have a GPU on this machine, so I want to be upfront about what I could and could not verify:

• New test passes, and I confirmed it actually catches the bug: reverting just the source change makes it fail with ZeroDivisionError: integer division or modulo by zero from triton/__init__.py, and restoring the fix makes it pass.

• No regressions: unmodified main gives 141 failed / 0 passed on this file, and with my change it is 141 failed / 1 passed. Same failures either way, all of them RuntimeError: 0 active drivers from triton because there is no CUDA device here. The only difference is my new test passing.

• ruff check and ruff format --check both clean on the two changed files.

• Could not run make test or make test-convergence for real, both need a GPU. Would appreciate CI covering those.

• Hardware Type: CPU only (no GPU available, see note above)

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

[nileshpatil6]

Giving this one a bump since it's been quiet since the resubmit. It's a small guard, one ZeroDivisionError fix with a test, rebased on current main and mergeable. CI never ran on this reopened version, so if a maintainer can approve the workflow run I'll take care of anything it flags right away.