[New Feature] integrate Liger Kernel with fused-CE for NPU + fix conv1d/seed bugs

ROADMAP.md

@@ -55,7 +55,7 @@
 
 - [x] 支持多模态模型
 - [ ] 支持megatron VPP
-- [ ] 支持liger kernel
+- [x] 支持liger kernel
 - [ ] 支持transformers模型的ulysses/ring-attention
 - [ ] 兼容transformers v5的tp、pp
 - [ ] 支持多轮RL
@@ -75,7 +75,7 @@
 
 - [x] Support for multimodal models
 - [ ] Support for Megatron VPP
-- [ ] Support for Liger kernel
+- [x] Support for Liger kernel
 - [ ] Support for Ulysses/Ring-Attention for Transformers models
 - [ ] Compatibility with Transformers v5 TP and PP
 - [ ] Support for multi-turn RL

cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py

@@ -5,6 +5,7 @@
     torchrun --nproc-per-node=8 cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py
 """
 from pathlib import Path
+import os
 
 from peft import LoraConfig
 from transformers import AutoConfig
@@ -17,7 +18,7 @@
 from twinkle.model import TransformersModel
 from twinkle.preprocessor import SelfCognitionProcessor
 from twinkle.utils.framework import Torch
-from twinkle.kernel import kernelize, npu_builtin
+from twinkle.kernel import kernelize, liger_builtin, npu_builtin
 
 logger = get_logger()
 args = CLI.from_args()
@@ -58,13 +59,19 @@ def train():
     if hasattr(text_config, 'use_cache'):
         text_config.use_cache = False
 
-    dataset = Dataset(dataset_meta=DatasetMeta(args.dataset.dataset_id))
+    # Slice the dataset to the training budget (default 320, matching the other
+    # cookbooks) so encoding the full LongAlpaca set (~12k examples) doesn't
+    # dominate runtime; this is a data-amount knob, not a sharding/batch change.
+    _train_samples = args.training.train_samples or 320
+    dataset = Dataset(dataset_meta=DatasetMeta(args.dataset.dataset_id, data_slice=range(_train_samples)))
     try:
         dataset.set_template(args.template.template_cls, model_id=args.model.model_id,
-                             max_length=args.template.max_length)
+                             max_length=args.template.max_length,
+                             truncation_strategy=args.template.truncation_strategy)
     except ValueError:
         dataset.set_template('Qwen3_5Template', model_id=args.model.model_id,
-                             max_length=args.template.max_length)
+                            max_length=args.template.max_length,
+                            truncation_strategy=args.template.truncation_strategy)
     dataset.map(SelfCognitionProcessor(
         args.extra.get('model_name', 'twinkle'),
         args.extra.get('model_author', 'ModelScope'),
@@ -85,9 +92,19 @@ def train():
             }
         },
     )
-    # npu patch
-    if Torch.is_npu_available():
-        model = kernelize(model, npu_builtin(model))
+    # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default,
+    # CANN + FLA) | npu+liger (--enable-liger, Liger per-layer on top of CANN).
+    # Sharding/batch are unchanged across modes.
+    _torch_baseline = os.environ.get('TWINKLE_TORCH_BASELINE', '').lower() in ('1', 'true', 'yes')
+    kernel_mapping = {}
+    if Torch.is_npu_available() and not _torch_baseline:
+        kernel_mapping.update(npu_builtin(model))
+    if args.model.enable_liger and not _torch_baseline:
+        kernel_mapping.update(liger_builtin(model))
+    if kernel_mapping:
+        model = kernelize(model, kernel_mapping)
+    _use_fused_ce = args.model.enable_liger and not _torch_baseline and args.model.enable_fused_ce
+    _task = 'fused_lm_ce' if _use_fused_ce else 'causal_lm'
     lora_cfg = _build_lora_config()
     model.add_adapter_to_model(args.lora.adapter_name, lora_cfg,
                                gradient_accumulation_steps=args.training.gradient_accumulation_steps)
@@ -97,6 +114,9 @@ def train():
         num_warmup_steps=args.scheduler.num_warmup_steps,
         num_training_steps=len(dataloader),
     )
+    if _use_fused_ce:
+        model.set_loss('LigerFusedLinearCrossEntropyLoss', adapter_name=args.lora.adapter_name,
+                       reduction='sum')
 
     if args.training.resume_from_checkpoint:
         checkpoint_path = Path(args.training.resume_from_checkpoint).expanduser().resolve()
@@ -116,10 +136,17 @@ def train():
         f'enable_ep={ENABLE_EP}, output_dir={args.training.output_dir}')
 
     optimizer_group = model.optimizer_group[args.lora.adapter_name]
+    if _use_fused_ce:
+        import torch.distributed as _dist
+        if _dist.is_available() and _dist.is_initialized():
+            _dist.barrier()
+        if Torch.is_npu_available():
+            import torch_npu
+            torch.npu.synchronize()
     for batch in dataloader:
         if callable(batch):
             batch = batch()
-        model.forward_backward(inputs=batch)
+        model.forward_backward(inputs=batch, task=_task)
         model.clip_grad_and_step(max_grad_norm=args.optimizer.max_grad_norm,
                                 gradient_accumulation_steps=args.training.gradient_accumulation_steps)
         cur_step = optimizer_group.cur_step

cookbook/transformers/fsdp2.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 from peft import LoraConfig
@@ -12,7 +13,7 @@
 from twinkle.model import TransformersModel
 from twinkle.preprocessor import SelfCognitionProcessor
 from twinkle.utils.framework import Torch
-from twinkle.kernel import kernelize, npu_builtin
+from twinkle.kernel import kernelize, liger_builtin, npu_builtin
 
 logger = get_logger()
 args = CLI.from_args()
@@ -23,7 +24,9 @@
 
 def build_dataset(num_samples: int) -> Dataset:
     dataset = Dataset(dataset_meta=DatasetMeta(args.dataset.dataset_id, data_slice=range(num_samples)))
-    dataset.set_template(args.template.template_cls, model_id=args.model.model_id)
+    dataset.set_template(args.template.template_cls, model_id=args.model.model_id,
+                        max_length=args.template.max_length,
+                        truncation_strategy=args.template.truncation_strategy)
     dataset.map(SelfCognitionProcessor(
         args.extra.get('model_name', 'twinkle大模型'),
         args.extra.get('model_author', 'ModelScope社区'),
@@ -56,10 +59,36 @@ def train():
     dataset = build_dataset(train_samples)
     dataloader = DataLoader(dataset=dataset, batch_size=args.training.batch_size)
     model = TransformersModel(model_id=args.model.model_id)
-    model.model._no_split_modules = {'Qwen3_5DecoderLayer'}
-    # npu patch
-    if Torch.is_npu_available():
-        model = kernelize(model, npu_builtin(model))
+    # Discover the actual decoder-layer class name(s) from the live model — the
+    # ``model_type.title() + 'DecoderLayer'`` heuristic misnames MoE models
+    # (e.g. ``qwen3_5_moe`` -> ``Qwen3_5_MoeDecoderLayer`` vs the real
+    # ``Qwen3_5MoeDecoderLayer``).
+    discovered = {type(m).__name__ for m in model.model.modules()
+                  if type(m).__name__.endswith('DecoderLayer')}
+    model.model._no_split_modules = list(discovered) or [model.model.config.model_type.title() + 'DecoderLayer']
+    # Compose the kernel mapping: NPU built-ins first, then Liger on top so
+    # `--enable-liger` opts into Liger's cross-device Triton/Ascend kernels
+    # (later keys win on overlap — see twinkle.kernel Kernel.md).
+    kernel_mapping = {}
+    # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default,
+    # CANN + FLA) | npu+liger (--enable-liger, Liger per-layer + fused-CE).
+    _torch_baseline = os.environ.get('TWINKLE_TORCH_BASELINE', '').lower() in ('1', 'true', 'yes')
+    if Torch.is_npu_available() and not _torch_baseline:
+        kernel_mapping.update(npu_builtin(model))
+    if args.model.enable_liger and not _torch_baseline:
+        kernel_mapping.update(liger_builtin(model))
+    if kernel_mapping:
+        model = kernelize(model, kernel_mapping)
+
+    # `--enable-liger` turns on BOTH the per-layer Liger/CANN kernels (above)
+    # AND, by default (`enable_fused_ce=True`), the LigerFusedLinearCrossEntropyLoss
+    # which skips the lm_head GEMM so the (B,T,V) logits tensor is never materialised.
+    # The forward then runs under task='fused_lm_ce' (TransformersFusedCEPatch).
+    # Pass `--no-fused-ce` to keep only the per-layer kernels (standard CE loss).
+    # The loss is device-agnostic: on NPU/CUDA it auto-falls-back to materialised
+    # CE if the fused kernel raises for a given shape (defensive).
+    _use_fused_ce = args.model.enable_liger and args.model.enable_fused_ce
+    _task = 'fused_lm_ce' if _use_fused_ce else 'causal_lm'
 
     lora_config = LoraConfig(**args.get_lora_args())
     model.add_adapter_to_model(
@@ -81,6 +110,9 @@ def train():
         scheduler_cls=args.scheduler.scheduler_cls,
         num_warmup_steps=args.scheduler.num_warmup_steps,
         num_training_steps=len(dataloader))
+    if _use_fused_ce:
+        model.set_loss('LigerFusedLinearCrossEntropyLoss', adapter_name=args.lora.adapter_name,
+                       reduction='sum')
 
     if args.training.resume_from_checkpoint:
         checkpoint_path = Path(args.training.resume_from_checkpoint).expanduser().resolve()
@@ -97,8 +129,31 @@ def train():
     optimizer_group = model.optimizer_group[args.lora.adapter_name]
     best_loss = float('inf')
     eval_interval = args.training.eval_interval or 40
+    if _use_fused_ce:
+        # One-time cross-rank + device drain before the first fused-CE step.
+        # Under accelerate-FSDP2 the first forward lazily runs fully_shard init
+        # (collectives) and the fused-CE loss issues an ad-hoc lm_head.weight
+        # full_tensor() gather after the (identity) lm_head forward; if ranks
+        # enter the loop slightly desynchronised (asymmetric setup), the
+        # collective streams order differently across ranks and the first
+        # backward deadlocks (rank A in autograd backward, rank B already in
+        # the Muon optimizer's DTensor redistribution). A single barrier +
+        # device sync re-aligns the ranks; later steps stay aligned because
+        # the fused-CE + FSDP2 collectives are rank-symmetric. liger_bench.py
+        # avoids this with a per-step _synchronize(); one pre-loop drain is
+        # sufficient (verified on 35B MoE, 4x Ascend, FSDP2).
+        import torch.distributed as _dist
+        if _dist.is_available() and _dist.is_initialized():
+            _dist.barrier()
+        if Torch.is_npu_available():
+            import torch_npu
+            torch.npu.synchronize()
+        if Torch.is_gpu_available():
+            import torch
+            torch.cuda.synchronize()
+        logger.info('[fsdp2] fused-CE: pre-loop barrier+device-sync drain applied')
     for batch in dataloader:
-        model.forward_backward(inputs=batch)
+        model.forward_backward(inputs=batch, task=_task)
         model.clip_grad_and_step()
         cur_step = optimizer_group.cur_step
         if cur_step % args.training.log_interval == 0:

cookbook/transformers/fsdp2.sh

@@ -1,6 +1,12 @@
 #!/bin/sh
 # All training config passed as CLI flags. Override at invocation, e.g.:
 #   sh fsdp2.sh --batch-size 16 --lr 5e-5
+#
+# Liger Kernel: pass --enable-liger to swap in Liger's Triton (CUDA) / Ascend
+# (NPU) kernels via `kernelize(model, liger_builtin(model))`. Without the flag
+# the NPU path uses `npu_builtin()` and the CUDA path is untouched.
+#   sh fsdp2.sh --enable-liger            # enable Liger
+#   sh fsdp2.sh --no-enable-liger         # explicitly disable (default)
 
 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
   torchrun --nproc_per_node=8 fsdp2.py \

cookbook/transformers/sp_fsdp_dense.py

@@ -1,4 +1,5 @@
 from functools import partial
+import os
 from peft import LoraConfig
 
 import twinkle
@@ -9,7 +10,7 @@
 from twinkle.model import TransformersModel
 from twinkle.preprocessor import SelfCognitionProcessor
 from twinkle.utils.framework import Torch
-from twinkle.kernel import kernelize, npu_builtin
+from twinkle.kernel import kernelize, liger_builtin, npu_builtin
 
 logger = get_logger()
 args = CLI.from_args()
@@ -45,7 +46,9 @@ def eval(model):
 def create_dataset(data_slice=None):
     train_samples = args.training.train_samples or 500
     dataset = Dataset(dataset_meta=DatasetMeta(args.dataset.dataset_id, data_slice=data_slice or range(train_samples)))
-    dataset.set_template(args.template.template_cls, model_id=args.model.model_id)
+    dataset.set_template(args.template.template_cls, model_id=args.model.model_id,
+                         max_length=args.template.max_length,
+                         truncation_strategy=args.template.truncation_strategy)
     dataset.map(SelfCognitionProcessor(
         args.extra.get('model_name', 'twinkle模型'),
         args.extra.get('model_author', 'twinkle团队'),
@@ -66,9 +69,19 @@ def train():
         device_mesh=device_mesh,
         strategy=args.model.strategy,
     )
-    # npu patch
-    if Torch.is_npu_available():
-        model = kernelize(model, npu_builtin(model))
+    # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default,
+    # CANN + FLA) | npu+liger (--enable-liger, Liger per-layer on top of CANN).
+    # Sharding/batch are unchanged across modes.
+    _torch_baseline = os.environ.get('TWINKLE_TORCH_BASELINE', '').lower() in ('1', 'true', 'yes')
+    kernel_mapping = {}
+    if Torch.is_npu_available() and not _torch_baseline:
+        kernel_mapping.update(npu_builtin(model))
+    if args.model.enable_liger and not _torch_baseline:
+        kernel_mapping.update(liger_builtin(model))
+    if kernel_mapping:
+        model = kernelize(model, kernel_mapping)
+    _use_fused_ce = args.model.enable_liger and not _torch_baseline and args.model.enable_fused_ce
+    _task = 'fused_lm_ce' if _use_fused_ce else 'causal_lm'
     lora_config = LoraConfig(**args.get_lora_args())
     model.add_adapter_to_model(args.lora.adapter_name, lora_config,
                                gradient_accumulation_steps=args.training.gradient_accumulation_steps)
@@ -81,11 +94,23 @@ def train():
         adapter_name=args.lora.adapter_name,
     )
 
+    if _use_fused_ce:
+        model.set_loss('LigerFusedLinearCrossEntropyLoss', adapter_name=args.lora.adapter_name,
+                       reduction='sum')
+
     logger.info(model.get_train_configs(adapter_name=args.lora.adapter_name))
     logger.info(f'Total steps: {len(dataloader)}')
 
+    if _use_fused_ce:
+        import torch.distributed as _dist
+        if _dist.is_available() and _dist.is_initialized():
+            _dist.barrier()
+        if Torch.is_npu_available():
+            import torch_npu
+            torch.npu.synchronize()
+
     for step, batch in enumerate(dataloader):
-        model.forward_backward(inputs=batch, adapter_name=args.lora.adapter_name)
+        model.forward_backward(inputs=batch, task=_task, adapter_name=args.lora.adapter_name)
         model.clip_grad_and_step(adapter_name=args.lora.adapter_name)
         if step % args.training.log_interval == 0:
             metric = model.calculate_metric(is_training=True, adapter_name=args.lora.adapter_name)

docs/source_en/Components/Kernel/Kernel.md

@@ -4,12 +4,13 @@
 implementation with another collapses to a single `kernelize(model, mapping)`
 call.
 
-The public surface is exactly three symbols:
+The public surface is exactly four symbols:
 
 | Symbol | Purpose |
 | --- | --- |
 | `kernelize(model, mapping=None)` | Apply ``mapping`` to ``model`` (in place) and return it. If ``mapping`` is omitted, it is auto-detected from the current platform (see below) |
 | `npu_builtin(model=None)` | Return the Ascend NPU built-in mapping (composes with user mappings) |
+| `liger_builtin(model=None)` | Return the Liger Kernel built-in mapping — cross-device (CUDA Triton + Ascend NPU). Bare impls, no device gating |
 | `hub(ref, *, revision=None, version=None, backend=None, trust_remote_code=False)` | Build a ``HubRef`` for use as a mapping value; the actual Hub download is deferred to ``kernelize`` |
 
 ## Mapping semantics
@@ -134,6 +135,39 @@ mapping = {
 model = kernelize(model, mapping)
 ```
 
+## Liger Kernel built-in
+
+`liger_builtin(model)` returns a dict of **bare** Liger impls (no device gating) covering RoPE, RMSNorm, SwiGLU/GeGLU, and MoE experts across the Qwen / Llama / Mistral / Mixtral / Phi3 / Gemma / Olmo2 / GLM4 / Granite / InternVL families. Values are bare because Liger self-dispatches across CUDA (Triton) and Ascend NPU (the auto-applied `backends/_ascend` backend) via `liger_kernel.utils.infer_device` — wrapping in `{'cuda': impl}` would wrongly skip Liger on NPU, where it is fully supported.
+
+The fused-linear-CE `forward` replacement and the global `nn.functional.cross_entropy` swap are **excluded** — they belong to the loss layer (`twinkle.loss`), not the kernel layer.
+
+```python
+from twinkle.kernel import kernelize, liger_builtin
+
+model = kernelize(model, liger_builtin(model))
+```
+
+### Composing Liger with the NPU bundle
+
+On NPU both `npu_builtin` and `liger_builtin` are NPU implementations of the same operators. Plain dict merge lets you pick precedence (later keys win):
+
+```python
+from twinkle.kernel import kernelize, liger_builtin, npu_builtin
+
+# Liger wins on overlap
+model = kernelize(model, {**npu_builtin(model), **liger_builtin(model)})
+# Twinkle-NPU wins on overlap
+model = kernelize(model, {**liger_builtin(model), **npu_builtin(model)})
+```
+
+### NPU precedence: CANN wins over Liger-Triton for per-layer ops
+
+On Ascend NPU, `liger_builtin()` post-processes itself via `_prefer_cann_on_npu`: Liger's Triton-on-Ascend kernels for the bandwidth-bound per-layer ops (RMSNorm, SwiGLU, RoPE) are swapped for the faster CANN vendor ops from `npu_impls`, and the `LigerExperts` / `LigerQwen3MoeSwiGLUMLP` class replacements are dropped so `npu_builtin`'s forward-level CANN grouped-matmul MoE expert path takes effect. Non-Qwen families without a CANN equivalent keep their Liger impls. So on NPU `liger_builtin` composes *additively* with `npu_builtin` (it only contributes ops CANN lacks), and `npu + liger` is not slower than `npu` alone for these ops. The fused-linear-CE loss (`twinkle.loss`) is unaffected — Liger's fused kernel is still used there because it has no CANN equivalent. On CUDA the bundle is unchanged.
+
+### RMSNorm attribute migration
+
+Liger's `LigerRMSNorm.forward` reads instance attributes (`offset` / `casting_mode` / `in_place` / `row_mode`) that HuggingFace RMSNorm variants do not define. Liger's monkey-patch sets these eagerly via `_patch_rms_norm_module`; the `liger_impls.rms_norm` adapters do the same **lazily** inside `forward` (with per-family defaults: llama-style `offset=0.0, casting_mode="llama"`, gemma-style `offset=1.0, casting_mode="gemma"`, gemma4 `offset=0.0`). No global state is mutated — contrast with `npu_builtin`'s SDPA install.
+
 ## Environment variables
 
 Only two remain: