ZeroDivisionError when finetuning with DeepSpeed + LigerKernel

Nobody has claimed this yet.

### 🐛 Describe the bug

When trying to finetune `Qwen2.5-1.5B-Instruct` with `DeepSpeed` + LigerKernel, I have encountered the error below:
```bash
Traceback (most recent call last):
File "/home/user/folder/main.py", line 170, in
result = trainer.train()
File "/home/user/folder/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 2240, in train
return inner_training_loop(
File "/home/user/folder/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
File "/home/user/folder/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 3745, in training_step
loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
File "/home/user/folder/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 3810, in compute_loss
outputs = model(**inputs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2063, in forward
loss = self.module(*inputs, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1857, in _call_impl
return inner()
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1805, in inner
result = forward_call(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
return func(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/liger_kernel/transformers/model/qwen2.py", line 203, in lce_forward
loss = LigerForCausalLMLoss(
File "/home/user/folder/.venv/lib/python3.10/site-packages/liger_kernel/transformers/model/loss_utils.py", line 55, in LigerForCausalLMLoss
loss = fixed_fused_linear_cross_entropy(
File "/home/user/folder/.venv/lib/python3.10/site-packages/liger_kernel/transformers/model/loss_utils.py", line 19, in fixed_fused_linear_cross_entropy
loss = F.liger_fused_linear_cross_entropy(
File "/home/user/folder/.venv/lib/python3.10/site-packages/liger_kernel/transformers/functional.py", line 64, in liger_fused_linear_cross_entropy
loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/autograd/function.py", line 575, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
File "/home/user/folder/.venv/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 510, in decorate_fwd
return fwd(*args, **kwargs)
File "/home/user/folder/.venv/lib/python3.10/site-packages/liger_kernel/ops/fused_linear_cross_entropy.py", line 240, in forward
loss, z_loss, grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_forward(
File "/home/user/folder/.venv/lib/python3.10/site-packages/liger_kernel/ops/fused_linear_cross_entropy.py", line 44, in fused_linear_cross_entropy_forward
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor
File "/home/user/folder/.venv/lib/python3.10/site-packages/triton/__init__.py", line 60, in cdiv
return (x + y - 1) // y
ZeroDivisionError: integer division or modulo by zero
```

I run the finetuning script using `uv run deepspeed --num_gpus=2 main.py`.

### Reproduce

With the code below, the `model.lm_head.weight` is `torch.Tensor([])`, which is reasonable, because I assume at this point the model is offloaded to NvME.

```python
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
DataCollatorForLanguageModeling,
TrainerCallback,
TrainerState,
TrainerControl
)
from deepspeed.runtime.utils import see_memory_usage
from datasets import load_dataset
from datasets.formatting.formatting import LazyBatch
from accelerate import init_empty_weights
from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig, HfDeepSpeedConfig
from transformers import Qwen2ForCausalLM, Qwen2Config
from liger_kernel.transformers import AutoLigerKernelForCausalLM
from liger_kernel.transformers import _apply_liger_kernel_to_instance, apply_liger_kernel_to_qwen2

from logging_memory_callback import MemoryLoggingCallback, logger

import deepspeed
import time
import torch
import torch.distributed as dist
import inspect
import os
import json
import traceback

MODEL_PATH = '/usr/local/models/Qwen2.5-1.5B-Instruct'
NUM_GPUS = len(os.getenv('CUDA_VISIBLE_DEVICES', '0').split(','))
logger.info(f"Running on {NUM_GPUS} GPUs")
MAX_SEQ_LEN = 1024

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

ds_config_path = './ds_config.json'

### Must initialize before calling `from_pretrained()`
training_args = TrainingArguments(
output_dir="/home/user/folder/output/deepspeed",
max_steps=5,
per_device_train_batch_size=8,
gradient_accumulation_steps=1,
gradient_checkpointing=True,
logging_dir=f"/home/user/folder/logs",
logging_strategy="steps",
logging_steps=1,
seed=42,
deepspeed=ds_config_path,
learning_rate=0.000007,
adam_beta1=0.9,
adam_beta2=0.95,
weight_decay=0.01,
warmup_steps=20,
fp16=True,
# use_liger_kernel=True, # liger kernel only supports llama
)

apply_liger_kernel_to_qwen2(
rope=True,
swiglu=True,
# cross_entropy=True,
fused_linear_cross_entropy=True
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
attn_implementation="flash_attention_2",
# use_flash_attention_2=True,
torch_dtype=torch.float16, # when initializing `TrainingArguments` before `from_pretrained()` need to set `float16` due to FlashAttention2 error
)
model.gradient_checkpointing_enable() # enable gradient checkpointing
print(f"Signature: {inspect.signature(model.forward)}")
print(f"LM head weights: {model.lm_head.weight}")
print(f"LM head weights shape: {model.lm_head.weight.shape}")

dataset = load_dataset('Dahoas/rm-static', split='train')
samples = dataset.select(range(300))

def format_chat_template(question: str, answer: str) -> str:
"""
Apply chat template to the QA given using the
chat template defined in the model config with `Tokenizer` class

Args:
question: Question of the dataset
answer: Answer of the dataset

Returns:
str: Formatted chat template based on the QA pair given.
"""
messages = [
{
'role': 'user',
'content': question
},
{
'role': 'assistant',
'content': answer
},
]

chat_template = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False,
)

return chat_template

def tokenize_function(examples: LazyBatch) -> torch.Tensor:
"""
Function to tokenize the formatted chat template into tokens

Args:
examples: Batch of dataset from the `map` function
"""
chat_templates = [
format_chat_template(question, answer)
for question, answer in zip(examples['prompt'], examples['response'])
]

return tokenizer(chat_templates, truncation=True, padding='max_length', max_length=MAX_SEQ_LEN, return_tensors='pt')

tokenized_dataset = samples.map(tokenize_function, batched=True)
print("Dataset processed")

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
memory_callback = MemoryLoggingCallback(
world_size=NUM_GPUS,
max_seq_len=MAX_SEQ_LEN,
) # callback func to log metrics required

trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
processing_class=tokenizer,
data_collator=data_collator,
callbacks=[memory_callback],
)
print("Trainer initialized")

try:
print("Start training")
start_time = time.perf_counter()
result = trainer.train()
end_time = time.perf_counter()

logger.info(result)

logger.info(f"Total training time: {end_time - start_time:.4f} seconds")

except Exception as e:
stacktrace = traceback.format_exc()
logger.error(stacktrace)

```

My `ds_config.json` looks something like
```json
{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},

"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},

"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},

"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "nvme",
"pin_memory": true,
"nvme_path": "/mnt/nvme1",
"buffer_count": 5
},
"offload_param": {
"device": "nvme",
"pin_memory": true,
"nvme_path": "/mnt/nvme1",
"buffer_count": 5,
"buffer_size": 5e9,
"max_in_cpu": 1e8
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e6,
"reduce_bucket_size": 2e8,
"stage3_prefetch_bucket_size": 2e7,
"stage3_param_persistence_threshold": 1e6,
"stage3_max_live_parameters": 1e8,
"stage3_max_reuse_distance": 0,
"stage3_gather_16bit_weights_on_model_save": true
},
"activation_checkpointing": {
"partition_activations": true,
"cpu_checkpointing": true,
"contiguous_memory_optimization": false,
"number_checkpoints": null,
"synchronize_checkpoint_boundary": false
},
"compile": {
"deepcompile": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}

```

### Versions

Below is the environment report:

```bash
Environment Report:
-------------------
Operating System: Linux-5.15.0-141-generic-x86_64-with-glibc2.35
Python version: 3.10.12
Liger Kernel version: 0.5.10
PyTorch version: 2.7.0+cu126
CUDA version: 12.6
HIP(ROCm) version: Not available
Triton version: 3.3.0
Transformers version: 4.52.4
XPU version: XPU Not Available
```

Start by reproducing the reported command, `uv run deepspeed --num_gpus=2 main.py`, with the supplied `ds_config.json` and environment versions. Read `liger_kernel/ops/fused_linear_cross_entropy.py` around the `triton.cdiv` call and the Qwen integration shown in `liger_kernel/transformers/model/qwen2.py`. Done means the DeepSpeed plus LigerKernel finetuning path no longer raises the reported ZeroDivisionError.

Written by the indexing model from the issue text.