[flyingdoubleG]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

LigerFusedLinearCrossEntropyLoss is documented as a drop-in for CrossEntropyLoss, but it differs from it in three places.

1. reduction='mean' normalizes twice. In the fused path the loss always returned num_tokens = #non-ignored tokens. calculate_loss treats a non-zero num_tokens as "sum loss, divide the gradient by the token count", which is the CrossEntropyLoss(reduction='sum') contract. CrossEntropyLoss(reduction='mean') returns num_tokens=0. So with reduction='mean', an option the docstring offers, the already averaged loss was divided by the token count again, and the effective learning rate shrank with batch size. The fused path now returns num_tokens=0 for 'mean'. 'sum' (what the cookbooks use) is unchanged.
2. Micro-batching raised NotImplementedError. The class had no micro_batch_scale, so forward_backward(..., micro_batch_size=...) and dynamic batching hit the base-class NotImplementedError. It now uses the same token-weighted scale as CrossEntropyLoss (1.0 for 'sum').
3. The materialised fallback ignored softcap. After the fused kernel fails, _materialise_ce recomputes logits with F.linear but did not apply the softcap that was passed to the fused kernel. For softcapped models the fallback therefore computed a different loss, despite the "same numerics" note. It now applies softcap * tanh(logits / softcap).

Experiment results

liger_kernel needs CUDA/NPU, so the new CPU tests swap the kernel for a pure-torch module with the same semantics (F.cross_entropy(F.linear(...)), with the optional softcap). They cover:

• test_fused_path_num_tokens_matches_cross_entropy[mean|sum]: loss and num_tokens match CrossEntropyLoss;
• test_micro_batch_scale_matches_cross_entropy[mean|sum];
• test_materialised_fallback_applies_softcap.

Four of the five fail on main. The sum / num_tokens case already passed. All five pass with this change.

The device tests (requires_liger / requires_device) were not run here, because I have no GPU. The fused-kernel call itself is unchanged.

🤖 Generated with Claude Code