[New Feature] integrate Liger Kernel with fused-CE for NPU + fix conv1d/seed bugs#249
[New Feature] integrate Liger Kernel with fused-CE for NPU + fix conv1d/seed bugs#249ys2025-AI wants to merge 2 commits into
Conversation
…bugs Liger Kernel integration: - Add liger_builtin() bundle with _prefer_cann_on_npu() that swaps Liger Triton-on-Ascend kernels for CANN equivalents (faster on NPU) - Add LigerFusedLinearCrossEntropyLoss + TransformersFusedCEPatch that skips lm_head GEMM and [B,T,V] logits materialisation - Wire fused-CE in fsdp2.py/sp_fsdp_dense.py/ep_fsdp2_lora_qwen3_5_moe.py (task='fused_lm_ce' + set_loss + pre-loop barrier) - Add --enable-liger / --no-enable-fused-ce CLI flags Bug fixes: - causal_conv1d.py: fix npu_causal_conv1d_fn axis confusion when conv_dim == seq_len (e.g. 35B-A3B with max_length=8192); add x.shape[1] != D check to disambiguate [B,D,T] vs [B,T,D] layouts - framework.py: fix torch.manual_seed gated behind is_gpu_available() on NPU-only hosts; move to unconditional + independent NPU branch Tests: - test_qwen35_fla_bwd_precision.py: FLA on/off backward precision - test_liger_builtin.py, test_liger_fused_linear_ce.py, test_fused_ce_patch.py Benchmark (35B-A3B, 4x910B3, max_len=8192): - fused-CE saves 12-20 GB/NPU (20-30%) across all 3 scripts - Enables EP batch=8 (was OOM without fused-CE) - SP path 1.6x speedup from skipping logits SP all-gather
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Liger Kernel integration, adding cross-device kernel replacements and a fused linear cross-entropy loss (LigerFusedLinearCrossEntropyLoss) to optimize memory usage. Key feedback includes fixing a critical bug in npu_causal_conv1d_fn where the standard [B, D, T] layout implicitly returns None, and replacing unsafe try...except ImportError blocks with Torch.is_npu_available() across multiple cookbooks to prevent potential AttributeError crashes. Additionally, the bias parameter in the fused cross-entropy fallback should be gathered via full_tensor() if sharded under FSDP, and _no_split_modules in fsdp2.py should be converted from a set to a list to ensure compatibility with HuggingFace Transformers.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
- Replace unsafe try/except ImportError with Torch.is_npu_available() / Torch.is_gpu_available() in fsdp2.py, sp_fsdp_dense.py, ep_fsdp2_lora_qwen3_5_moe.py pre-loop barrier (AttributeError not caught by except ImportError) - Gather bias via full_tensor() in LigerFusedLinearCrossEntropyLoss fallback (bias could be FSDP-sharded DTensor) - Convert _no_split_modules from set to list for HF Transformers compatibility - Gemini modelscope#1 (conv1d missing else) is a false positive: else branch exists at line 1114
1 participant
PR type
PR information
Liger Kernel integration:
Triton-on-Ascend kernels for CANN equivalents (faster on NPU)
skips lm_head GEMM and [B,T,V] logits materialisation
(task='fused_lm_ce' + set_loss + pre-loop barrier)
Bug fixes:
conv_dim == seq_len (e.g. 35B-A3B with max_length=8192); add
x.shape[1] != D check to disambiguate [B,D,T] vs [B,T,D] layouts
on NPU-only hosts; move to unconditional + independent NPU branch
Tests:
Experiment results
Benchmark (Qwen3.5-35B-A3B, 4x910B3, max_len=8192):