diff --git a/ROADMAP.md b/ROADMAP.md index e11205e7..f7cb712f 100644 --- a/ROADMAP.md +++ b/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 diff --git a/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py b/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py index e1cbd2c8..71302b3f 100644 --- a/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py +++ b/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 diff --git a/cookbook/transformers/fsdp2.py b/cookbook/transformers/fsdp2.py index 43ff78e0..397df79a 100644 --- a/cookbook/transformers/fsdp2.py +++ b/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: diff --git a/cookbook/transformers/fsdp2.sh b/cookbook/transformers/fsdp2.sh index 77ae00a5..5a9a7772 100644 --- a/cookbook/transformers/fsdp2.sh +++ b/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 \ diff --git a/cookbook/transformers/sp_fsdp_dense.py b/cookbook/transformers/sp_fsdp_dense.py index b41a0400..dc7f6bc6 100644 --- a/cookbook/transformers/sp_fsdp_dense.py +++ b/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) diff --git a/docs/source_en/Components/Kernel/Kernel.md b/docs/source_en/Components/Kernel/Kernel.md index 972fbf33..ce7a2768 100644 --- a/docs/source_en/Components/Kernel/Kernel.md +++ b/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: diff --git "a/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" "b/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" index a97c778e..f484f677 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" @@ -2,12 +2,13 @@ `twinkle.kernel` 提供一个 mapping 驱动的内核替换接口，把“用一种实现替换模型里的另一种实现”压缩为一次 `kernelize(model, mapping)` 调用。 -公开符号只有三个： +公开符号只有四个： | 符号 | 作用 | | --- | --- | | `kernelize(model, mapping=None)` | 在 `model` 上应用 `mapping`，原地修改后返回。省略 `mapping` 时按当前平台自动检测（见下文） | | `npu_builtin(model=None)` | 返回 Ascend NPU 内置替换的 mapping dict（可与用户 mapping 自由组合） | +| `liger_builtin(model=None)` | 返回 Liger Kernel 内置替换的 mapping dict —— 跨设备（CUDA Triton + Ascend NPU）。值为裸 impl，不做设备门控 | | `hub(ref, *, revision=None, version=None, backend=None, trust_remote_code=False)` | 构造一个 `HubRef`，用作 mapping value；真实下载推迟到 `kernelize` 执行 | ## Mapping 语义 @@ -132,6 +133,39 @@ mapping = { model = kernelize(model, mapping) ``` +## Liger Kernel 内置 + +`liger_builtin(model)` 返回**裸** Liger impl（不做设备门控）的 mapping，覆盖 Qwen / Llama / Mistral / Mixtral / Phi3 / Gemma / Olmo2 / GLM4 / Granite / InternVL 各族的 RoPE、RMSNorm、SwiGLU/GeGLU 及 MoE experts。值为裸 impl，是因为 Liger 自身就跨设备分派：CUDA 走 Triton、Ascend NPU 走自动应用的 `backends/_ascend` 后端（经 `liger_kernel.utils.infer_device`）。若包成 `{'cuda': impl}` 会在 NPU 上错误跳过——而 Liger 在 NPU 上是完全支持的。 + +融合线性 CE 的 `forward` 替换与全局 `nn.functional.cross_entropy` 替换**不在此 bundle 内**——它们属于 loss 层（`twinkle.loss`），不属于 kernel 层。 + +```python +from twinkle.kernel import kernelize, liger_builtin + +model = kernelize(model, liger_builtin(model)) +``` + +### Liger 与 NPU bundle 组合 + +在 NPU 上 `npu_builtin` 与 `liger_builtin` 都是同一批算子的 NPU 实现。普通 dict 合并即可选择优先级（后写的 key 胜出）： + +```python +from twinkle.kernel import kernelize, liger_builtin, npu_builtin + +# 重叠算子由 Liger 胜出 +model = kernelize(model, {**npu_builtin(model), **liger_builtin(model)}) +# 重叠算子由 Twinkle-NPU 胜出 +model = kernelize(model, {**liger_builtin(model), **npu_builtin(model)}) +``` + +### NPU 优先级：逐层算子上 CANN 胜过 Liger-Triton + +在 Ascend NPU 上，`liger_builtin()` 会通过 `_prefer_cann_on_npu` 做后处理：对带宽敏感的逐层算子（RMSNorm、SwiGLU、RoPE），Liger 的 Triton-on-Ascend 内核会被替换为 `npu_impls` 里更快的 CANN 厂商算子；`LigerExperts` / `LigerQwen3MoeSwiGLUMLP` 的类替换也会被丢弃，从而让 `npu_builtin` 的 forward 级 CANN 分组矩阵乘 MoE expert 路径生效。没有 CANN 对应实现的非 Qwen 族仍保留 Liger impl。因此在 NPU 上 `liger_builtin` 与 `npu_builtin` 是**叠加**关系（只贡献 CANN 缺少的算子），`npu + liger` 在这些算子上不会比单独 `npu` 更慢。融合线性 CE loss（`twinkle.loss`）不受影响——该处没有 CANN 对应实现，仍用 Liger 的融合内核。CUDA 上 bundle 不变。 + +### RMSNorm 属性迁移 + +Liger 的 `LigerRMSNorm.forward` 读取的实例属性（`offset` / `casting_mode` / `in_place` / `row_mode`）在 HuggingFace 的 RMSNorm 变体上并不存在。Liger 的 monkey-patch 通过 `_patch_rms_norm_module` 急切地设置这些属性；`liger_impls.rms_norm` 适配器改为在 `forward` 内**懒设置**（按族默认值：llama 风格 `offset=0.0, casting_mode="llama"`，gemma 风格 `offset=1.0, casting_mode="gemma"`，gemma4 `offset=0.0`）。不污染任何全局状态——与 `npu_builtin` 的 SDPA 全局 install 形成对比。 + ## 环境变量 只有两个保留： diff --git a/src/twinkle/cli/cli.py b/src/twinkle/cli/cli.py index ce720a9f..a12594d0 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -34,6 +34,17 @@ class ModelArgs: ddp_config: dict[str, Any] | None = None fsdp_config: dict[str, Any] | None = None grad_scaler_config: dict[str, Any] | None = None + # Liger Kernel toggle: when True, the cookbook applies `liger_builtin()` via + # `kernelize`. Off by default — opt in with --enable-liger / TWINKLE_ENABLE_LIGER. + enable_liger: bool = False + # Fused-linear-CE loss toggle. Only meaningful when `enable_liger` is True. + # Defaults True so `--enable-liger` turns on BOTH the per-layer Liger/CANN + # kernels AND the LigerFusedLinearCrossEntropyLoss (skip-lm_head-GEMM, no + # (B,T,V) logits). Pass `--no-fused-ce` to opt out of the fused-CE loss and + # keep only the per-layer kernels (the loss falls back to standard CE). The + # loss itself is device-agnostic: on NPU/CUDA it auto-falls-back to + # materialised CE if the fused kernel raises for a given shape. + enable_fused_ce: bool = True @dataclass diff --git a/src/twinkle/kernel/__init__.py b/src/twinkle/kernel/__init__.py index f1499680..eb738edc 100644 --- a/src/twinkle/kernel/__init__.py +++ b/src/twinkle/kernel/__init__.py @@ -1,13 +1,15 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Mapping-driven kernel replacement. -Three public symbols: +Public symbols: - :func:`kernelize` apply ``mapping`` to a model - :func:`hub` build a Hub kernel reference - :func:`npu_builtin` the Ascend NPU built-in bundle +- :func:`liger_builtin` the Liger Kernel built-in bundle (cross-device) """ from .builtin import npu_builtin from .core import hub, kernelize +from .liger import liger_builtin -__all__ = ['kernelize', 'hub', 'npu_builtin'] +__all__ = ['kernelize', 'hub', 'npu_builtin', 'liger_builtin'] diff --git a/src/twinkle/kernel/causal_conv1d.py b/src/twinkle/kernel/causal_conv1d.py index b18c014f..17c4db57 100644 --- a/src/twinkle/kernel/causal_conv1d.py +++ b/src/twinkle/kernel/causal_conv1d.py @@ -1093,10 +1093,21 @@ def npu_causal_conv1d_fn( backend: Optional[str] = None, cu_seqlens: Optional[torch.Tensor] = None, ): - """Adapter matching twinkle's ``causal_conv1d_fn`` call signature.""" - del seq_idx, backend + """Adapter for twinkle's ``causal_conv1d_fn`` call signature. - if x.dim() == 3 and weight.dim() == 2 and x.shape[-1] == weight.shape[0] and x.shape[-1] != weight.shape[-1]: + Handles two input layouts: standard Qwen3.5 path passes ``x=[B, D, T]`` + (needs transpose), SP path passes ``x=[B, T, D]`` (no transpose needed). + Detected via ``x.shape[-1] == D and x.shape[1] != D``; when ``T == D`` + (ambiguous) defaults to transposing (standard path). + """ + del seq_idx, backend + D, W = weight.shape[0], weight.shape[1] + is_bt_d_layout = ( + x.dim() == 3 and weight.dim() == 2 and x.shape[-1] == D # last dim is the channel dim + and x.shape[1] != D # second dim is NOT D -> genuinely [B, T, D] + and D != W # sanity: D != kernel_size + ) + if is_bt_d_layout: weight_t = weight.transpose(0, 1).contiguous() y_t, _ = causal_conv1d(x=x, weight=weight_t, bias=bias, activation=activation, cu_seqlens=cu_seqlens) return y_t diff --git a/src/twinkle/kernel/chunk_gated_delta_rule.py b/src/twinkle/kernel/chunk_gated_delta_rule.py index 2d0beee7..18fe9e32 100644 --- a/src/twinkle/kernel/chunk_gated_delta_rule.py +++ b/src/twinkle/kernel/chunk_gated_delta_rule.py @@ -6,15 +6,29 @@ import torch import warnings -from mindspeed.lite.ops.triton.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -from mindspeed.lite.ops.triton.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o -from mindspeed.lite.ops.triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from mindspeed.lite.ops.triton.cumsum import chunk_local_cumsum -from mindspeed.lite.ops.triton.solve_tril import solve_tril -from mindspeed.lite.ops.triton.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard -from mindspeed.lite.ops.triton.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd from typing import Optional +# MindSpeed relocated its triton ops between package layouts across versions: +# newer exposes them under ``mindspeed.ops.triton.*``; older under +# ``mindspeed.lite.ops.triton.*``. Try the newer path first and fall back to +# the legacy ``.lite.`` layout so this wrapper works across MindSpeed versions. +try: + from mindspeed.ops.triton.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h + from mindspeed.ops.triton.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o + from mindspeed.ops.triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd + from mindspeed.ops.triton.cumsum import chunk_local_cumsum + from mindspeed.ops.triton.solve_tril import solve_tril + from mindspeed.ops.triton.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + from mindspeed.ops.triton.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +except ImportError: + from mindspeed.lite.ops.triton.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h + from mindspeed.lite.ops.triton.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o + from mindspeed.lite.ops.triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd + from mindspeed.lite.ops.triton.cumsum import chunk_local_cumsum + from mindspeed.lite.ops.triton.solve_tril import solve_tril + from mindspeed.lite.ops.triton.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + from mindspeed.lite.ops.triton.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd + def _torch_l2norm_fwd( x: torch.Tensor, diff --git a/src/twinkle/kernel/liger.py b/src/twinkle/kernel/liger.py new file mode 100644 index 00000000..c8366a71 --- /dev/null +++ b/src/twinkle/kernel/liger.py @@ -0,0 +1,325 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""``liger_builtin()`` returns the Liger Kernel replacement bundle. + +This mirrors ``npu_builtin()`` but for Liger Kernel. Two key differences from +the NPU bundle, both consequences of Liger being cross-device: + + 1. **No device gating.** Values are *bare* impls (not ``{'cuda': impl}``). + Liger's modules self-dispatch across CUDA (Triton) and Ascend NPU (the + auto-applied ``backends/_ascend`` backend) via ``infer_device`` / + ``select_impl`` in ``liger_kernel``. Device-conditional wrapping would + wrongly skip Liger on NPU, where it is fully supported. + + 2. **No side effects.** Unlike ``npu_builtin`` (which installs global SDPA + and per-instance FLA), this bundle contains only class/function + replacements consumable by ``kernelize``. The process-level + ``nn.functional.cross_entropy`` swap and fused-linear-CE ``forward`` + replacement are deliberately *excluded* — they belong to the loss layer, + not the kernel layer (see ROADMAP / ``twinkle.loss``). + +Composing with ``npu_builtin`` on NPU: both bundles are NPU implementations of +the same operators. Plain dict merge lets the user pick precedence (later keys +win) — see ``Kernel.md``: + + model = kernelize(model, {**npu_builtin(model), **liger_builtin(model)}) # Liger wins + model = kernelize(model, {**liger_builtin(model), **npu_builtin(model)}) # Twinkle-NPU wins + model = kernelize(model, liger_builtin(model)) # Liger only + +RMSNorm uses Liger-aware adapter classes (``liger_impls.rms_norm``) that set the +Liger-specific instance attributes lazily in ``forward``; SwiGLU/GeGLU use +function-level forward replacements; RoPE and the Qwen3-MoE expert/MLP classes +are wired as bare Liger classes (their forwards read only attributes the HF +instances already provide). +""" +from __future__ import annotations + +import importlib +import torch.nn as nn +from typing import Any + +from twinkle import get_logger +from twinkle.utils.device_mesh import Platform + +logger = get_logger() + + +def _import_optional(name: str): + try: + return importlib.import_module(name) + except ImportError: + return None + + +def _has(maybe_none, attr_path: str) -> bool: + """Return True if ``attr_path`` (may contain one dot, e.g. ``Cls.member``) + exists on ``maybe_none`` (a module or None).""" + if maybe_none is None: + return False + if '.' in attr_path: + head, _, tail = attr_path.partition('.') + owner = getattr(maybe_none, head, None) + return owner is not None and hasattr(owner, tail) + return hasattr(maybe_none, attr_path) + + +def _add_class_if_present(bundle, module_path, class_name, impl_cls): + mod = _import_optional(module_path) + if mod is None: + return + cls = getattr(mod, class_name, None) + if isinstance(cls, type): + bundle[cls] = impl_cls + + +def _add_fwd_if_present(bundle, module_path, class_name, fn): + """Function-level forward replacement: ``bundle['..forward'] = fn``. + + Only added when the HF module imports AND the target class exists, so the + bundle stays safe when a model family is not installed. + """ + mod = _import_optional(module_path) + if mod is None: + return + cls = getattr(mod, class_name, None) + if isinstance(cls, type): + bundle[f'{module_path}.{class_name}.forward'] = fn + + +def _add_attr_if_present(bundle, module_path, attr_name, impl): + mod = _import_optional(module_path) + if mod is None: + return + if _has(mod, attr_name): + bundle[f'{module_path}.{attr_name}'] = impl + + +# ── family registrations ────────────────────────────────────────────────────── +# Each helper mirrors the subset of liger_kernel's apply_liger_kernel_to_ +# that is expressible as a pure class/function replacement (RoPE, RMSNorm, +# SwiGLU/GeGLU). Fused-linear-CE and the global cross_entropy swap are excluded. + + +def _add_llama_style(bundle, module_path, rms_name, mlp_name, with_rope=True): + """llama / qwen / mistral / mixtral / phi3 / olmo2 / glm4 / granite / internvl: + llama-cast RMSNorm + SwiGLU MLP + (optional) RoPE.""" + from liger_kernel.transformers import liger_rotary_pos_emb + + from .liger_impls import LigerRMSNormReplacement, liger_swiglu_forward + + _add_class_if_present(bundle, module_path, rms_name, LigerRMSNormReplacement) + _add_fwd_if_present(bundle, module_path, mlp_name, liger_swiglu_forward) + if with_rope: + _add_attr_if_present(bundle, module_path, 'apply_rotary_pos_emb', liger_rotary_pos_emb) + + +def _add_gemma_style(bundle, module_path, rms_name, mlp_name): + """gemma / gemma2 / gemma3_text: gemma-cast RMSNorm (offset=1.0) + GeGLU MLP. + RoPE differs per gemma variant and is left to the model-specific entries.""" + from .liger_impls import LigerRMSNormGemmaReplacement, liger_geglu_forward + + _add_class_if_present(bundle, module_path, rms_name, LigerRMSNormGemmaReplacement) + _add_fwd_if_present(bundle, module_path, mlp_name, liger_geglu_forward) + + +def _add_qwen2(bundle): + _add_llama_style(bundle, 'transformers.models.qwen2.modeling_qwen2', 'Qwen2RMSNorm', 'Qwen2MLP') + + +def _add_qwen3(bundle): + _add_llama_style(bundle, 'transformers.models.qwen3.modeling_qwen3', 'Qwen3RMSNorm', 'Qwen3MLP') + + +def _add_qwen3_moe(bundle): + base = 'transformers.models.qwen3_moe.modeling_qwen3_moe' + from liger_kernel.transformers import LigerExperts, LigerQwen3MoeSwiGLUMLP, liger_rotary_pos_emb + + from .liger_impls import LigerRMSNormReplacement, liger_swiglu_forward + + _add_class_if_present(bundle, base, 'Qwen3MoeRMSNorm', LigerRMSNormReplacement) + _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', liger_rotary_pos_emb) + # transformers v5+: experts are a separate fused class; v4 uses Qwen3MoeMLP. + _add_class_if_present(bundle, base, 'Qwen3MoeExperts', LigerExperts) + _add_fwd_if_present(bundle, base, 'Qwen3MoeMLP', liger_swiglu_forward) + # Qwen3MoeMLP (v4) is structurally identical to LigerQwen3MoeSwiGLUMLP; offer + # class replacement too so v4 users get the fused SiLU kernel without a + # forward-level swap shadowing it. + mod = _import_optional(base) + if mod is not None and isinstance(getattr(mod, 'Qwen3MoeMLP', None), type): + if not any(k == getattr(mod, 'Qwen3MoeMLP') for k in bundle): + bundle[getattr(mod, 'Qwen3MoeMLP')] = LigerQwen3MoeSwiGLUMLP + + +def _add_qwen3_5(bundle): + from .liger_impls import LigerRMSNormQwen35Replacement, liger_swiglu_forward + + base = 'transformers.models.qwen3_5.modeling_qwen3_5' + if _import_optional(base) is None: + return + _add_class_if_present(bundle, base, 'Qwen3_5RMSNorm', LigerRMSNormQwen35Replacement) + _add_class_if_present(bundle, base, 'Qwen3_5VisionRMSNorm', LigerRMSNormQwen35Replacement) + # RoPE intentionally NOT replaced: Qwen3.5 uses partial_rotary_factor=0.25 + + # mrope_interleaved=True; Liger's liger_rotary_pos_emb assumes full-rotation + # with the rotate_half convention, which is incompatible. Let npu_builtin's + # npu_apply_rotary_pos_emb (which handles Partial-RoPE) take precedence. + _add_fwd_if_present(bundle, base, 'Qwen3_5MLP', liger_swiglu_forward) + _add_fwd_if_present(bundle, base, 'Qwen3_5VisionMLP', liger_swiglu_forward) + + +def _add_qwen3_5_moe(bundle): + from liger_kernel.transformers import LigerExperts + + from .liger_impls import LigerRMSNormQwen35Replacement, liger_swiglu_forward + + base = 'transformers.models.qwen3_5_moe.modeling_qwen3_5_moe' + if _import_optional(base) is None: + return + _add_class_if_present(bundle, base, 'Qwen3_5MoeRMSNorm', LigerRMSNormQwen35Replacement) + # RoPE intentionally NOT replaced (same reason as _add_qwen3_5). + _add_class_if_present(bundle, base, 'Qwen3_5MoeExperts', LigerExperts) + _add_fwd_if_present(bundle, base, 'Qwen3_5MoeMLP', liger_swiglu_forward) + + +def _add_qwen2_5_vl(bundle): + from liger_kernel.transformers import liger_rotary_pos_emb + + from .liger_impls import LigerRMSNormReplacement, liger_swiglu_forward + + base = 'transformers.models.qwen2_5_vl.modeling_qwen2_5_vl' + if _import_optional(base) is None: + return + _add_class_if_present(bundle, base, 'Qwen2_5_VLRMSNorm', LigerRMSNormReplacement) + _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', liger_rotary_pos_emb) + _add_fwd_if_present(bundle, base, 'Qwen2MLP', liger_swiglu_forward) + _add_fwd_if_present(bundle, base, 'Qwen2_5_VLMLP', liger_swiglu_forward) + + +def _add_gemma4(bundle): + base = 'transformers.models.gemma4.modeling_gemma4' + if _import_optional(base) is None: + return + from .liger_impls import LigerRMSNormGemma4Replacement, liger_geglu_forward + + _add_class_if_present(bundle, base, 'Gemma4RMSNorm', LigerRMSNormGemma4Replacement) + _add_fwd_if_present(bundle, base, 'Gemma4TextMLP', liger_geglu_forward) + + +def liger_builtin(model: nn.Module | None = None) -> dict[Any, Any]: + """Return the Liger Kernel built-in mapping; composes with ``kernelize``. + + Args: + model: Accepted for API symmetry with ``npu_builtin(model)``; Liger's + pure class/function replacements have no per-instance side effects, + so the model is not traversed here. + + Returns: + A ``dict`` whose keys are HF ``nn.Module`` subclasses (class replacement) + or dotted paths (function/forward replacement) and whose values are + *bare* Liger impls (no device gating). Missing model families are + silently skipped — the bundle only contains entries for installed + transformers model modules. + + On NPU, Liger's Triton-on-Ascend kernels are slower than the CANN vendor + ops in ``npu_impls`` for the bandwidth-bound per-layer ops (RMSNorm, + SwiGLU, RoPE). ``_prefer_cann_on_npu`` post-processes the bundle to swap + those Liger impls for their CANN equivalents and drops the LigerExperts + class replacement (which shadows ``npu_builtin``'s faster forward-level + MoE expert replacement). Non-Qwen families without a CANN equivalent keep + their Liger impls. On CUDA, the bundle is unchanged. + + Raises: + ImportError: if ``liger_kernel`` is not importable (caught at the first + ``from liger_kernel ...`` statement inside the family helpers). + """ + bundle: dict[Any, Any] = {} + + # ── Qwen family (primary on Twinkle) ────────────────────────────────────── + _add_qwen2(bundle) + _add_qwen3(bundle) + _add_qwen3_moe(bundle) + _add_qwen3_5(bundle) + _add_qwen3_5_moe(bundle) + _add_qwen2_5_vl(bundle) + + # ── Llama-style dense / MoE families ───────────────────────────────────── + _add_llama_style(bundle, 'transformers.models.llama.modeling_llama', 'LlamaRMSNorm', 'LlamaMLP') + _add_llama_style(bundle, 'transformers.models.mistral.modeling_mistral', 'MistralRMSNorm', 'MistralMLP') + _add_llama_style(bundle, 'transformers.models.mixtral.modeling_mixtral', 'MixtralRMSNorm', + 'MixtralBlockSparseTop2MLP') + _add_llama_style(bundle, 'transformers.models.phi3.modeling_phi3', 'Phi3RMSNorm', 'Phi3MLP') + _add_llama_style(bundle, 'transformers.models.glm4.modeling_glm4', 'Glm4RMSNorm', 'Glm4MLP') + _add_llama_style(bundle, 'transformers.models.olmo2.modeling_olmo2', 'Olmo2RMSNorm', 'Olmo2MLP') + _add_llama_style(bundle, 'transformers.models.granite.modeling_granite', 'GraniteRMSNorm', 'GraniteMLP') + _add_llama_style(bundle, 'transformers.models.internvl.modeling_internvl', 'InternVLRMSNorm', 'InternVLMLP') + + # ── Gemma family (gemma-cast RMSNorm + GeGLU) ──────────────────────────── + _add_gemma_style(bundle, 'transformers.models.gemma.modeling_gemma', 'GemmaRMSNorm', 'GemmaMLP') + _add_gemma_style(bundle, 'transformers.models.gemma2.modeling_gemma2', 'Gemma2RMSNorm', 'Gemma2MLP') + _add_gemma_style(bundle, 'transformers.models.gemma3.modeling_gemma3', 'Gemma3RMSNorm', 'Gemma3MLP') + _add_gemma4(bundle) + + if not bundle: + logger.warning('[liger_builtin] No Liger entries were registered — ' + 'is liger_kernel installed and are transformers model modules importable?') + + # On NPU, prefer CANN vendor ops over Liger's Triton-on-Ascend kernels for + # the per-layer ops where CANN is significantly faster (RMSNorm: single-pass + # vs Liger's 2-pass tiled; SwiGLU/RoPE: one CANN op vs Triton launch + extra + # allocations). Also drop LigerExperts class replacements that shadow + # npu_builtin's faster forward-level MoE expert replacement. + if Platform.device_prefix() == 'npu': + _prefer_cann_on_npu(bundle) + return bundle + + +def _prefer_cann_on_npu(bundle: dict[Any, Any]) -> None: + """In-place: swap Liger Triton impls for CANN vendor ops on NPU. + + Replaces: + - LigerRMSNorm* class values -> NpuRMSNorm (single-pass CANN aclnn) + - liger_swiglu_forward values -> npu_swiglu_forward (one CANN fused op) + - liger_rotary_pos_emb values -> npu_apply_rotary_pos_emb (no Q/K copies) + + Removes: + - LigerExperts / LigerQwen3MoeSwiGLUMLP class keys — they replace + ``m.__class__`` which shadows npu_builtin's forward-level + ``npu_packed_moe_experts_forward`` (a different dict key), leaving + the fast CANN grouped-matmul MoE path unused. + """ + from .npu_impls import NpuRMSNorm, npu_apply_rotary_pos_emb, npu_swiglu_forward + + # Stringify once for fast membership checks. + def _val_name(v) -> str: + return getattr(v, '__name__', getattr(v, '__qualname__', '')) + + keys_to_drop: list[Any] = [] + for key, val in list(bundle.items()): + name = _val_name(val) + + # Drop LigerExperts / LigerQwen3MoeSwiGLUMLP class replacements. + if isinstance(val, type) and any( + s in name for s in ('LigerExperts', 'LigerQwen3MoeSwiGLUMLP', 'LigerQwen3_5MoeSwiGLUMLP')): + keys_to_drop.append(key) + continue + + # Swap LigerRMSNorm* class replacements -> NpuRMSNorm. + if isinstance(val, type) and name.startswith('LigerRMSNorm'): + bundle[key] = NpuRMSNorm + continue + + # Swap liger_swiglu_forward -> npu_swiglu_forward. + if name == 'liger_swiglu_forward': + bundle[key] = npu_swiglu_forward + continue + + # Swap liger_rotary_pos_emb -> npu_apply_rotary_pos_emb. + if name == 'liger_rotary_pos_emb': + bundle[key] = npu_apply_rotary_pos_emb + continue + + for key in keys_to_drop: + del bundle[key] + + if keys_to_drop: + logger.info( + '[liger_builtin] NPU: dropped %d LigerExperts/LigerMoeMLP class replacement(s) ' + 'so npu_builtin\'s CANN grouped-matmul MoE forward takes effect', len(keys_to_drop)) diff --git a/src/twinkle/kernel/liger_impls/__init__.py b/src/twinkle/kernel/liger_impls/__init__.py new file mode 100644 index 00000000..0ea8f34d --- /dev/null +++ b/src/twinkle/kernel/liger_impls/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Class-replacement adapters that bridge Liger Kernel modules onto +HuggingFace instances when applied via ``kernelize``. + +``kernelize`` swaps ``m.__class__ = impl_cls`` *without* calling ``__init__`` +(see ``Kernel.md`` caveats). Liger's own ``LigerRMSNorm.forward`` reads +instance attributes (``offset`` / ``casting_mode`` / ``in_place`` / ``row_mode``) +that HuggingFace RMSNorm variants do not define. Liger's monkey-patch path +sets those attributes eagerly via ``_patch_rms_norm_module``; the adapters here +do the same lazily inside ``forward`` so the class-replacement contract is +honoured and no global state is mutated. + +SwiGLU / Experts / RoPE need no adapter — Liger's classes/functions read only +attributes (``gate_proj`` / ``up_proj`` / ``down_proj`` / ``weight``) that the +HuggingFace instances already provide, so they are re-exported verbatim. +""" +from .rms_norm import (LigerRMSNormGemma4Replacement, LigerRMSNormGemmaReplacement, LigerRMSNormQwen35Replacement, + LigerRMSNormReplacement) +from .swiglu import liger_geglu_forward, liger_swiglu_forward + +__all__ = [ + 'LigerRMSNormReplacement', + 'LigerRMSNormGemmaReplacement', + 'LigerRMSNormGemma4Replacement', + 'LigerRMSNormQwen35Replacement', + 'liger_swiglu_forward', + 'liger_geglu_forward', +] diff --git a/src/twinkle/kernel/liger_impls/rms_norm.py b/src/twinkle/kernel/liger_impls/rms_norm.py new file mode 100644 index 00000000..f9d06470 --- /dev/null +++ b/src/twinkle/kernel/liger_impls/rms_norm.py @@ -0,0 +1,114 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Class-replacement RMSNorm adapters for Liger Kernel. + +Designed for ``kernelize`` class replacement: do not define ``__init__``; rely +on the attributes already present on the original HuggingFace instance +(``weight``, ``variance_epsilon`` *or* ``eps``). The Liger-specific attributes +(``offset`` / ``casting_mode`` / ``in_place`` / ``row_mode``) are set lazily on +the first ``forward`` with family-appropriate defaults, mirroring Liger's own +``_patch_rms_norm_module`` (monkey_patch.py) so no global state is mutated. + +Per-family defaults (matching ``liger_kernel.transformers.monkey_patch``): + + - llama / qwen / phi3 / olmo2 / glm4 / granite / internvl: + offset=0.0, casting_mode="llama", in_place=True + - gemma / gemma2 / gemma3_text / paligemma: + offset=1.0, casting_mode="gemma", in_place=False + - gemma4: + offset=0.0, casting_mode="gemma", in_place=False + - qwen3_5 / qwen3_5_moe: + offset=1.0, casting_mode="gemma", in_place=False + (residual parameterization: scale = 1 + weight, weight init=0) +""" +from __future__ import annotations + +import torch.nn as nn +from liger_kernel.ops import LigerRMSNormFunction + + +class _LigerRMSNormBase(nn.Module): + """Base adapter; subclasses set the family-specific defaults below.""" + + _liger_offset: float = 0.0 + _liger_casting_mode: str = 'llama' + _liger_in_place: bool = True + _liger_row_mode = None + + def _ensure_liger_attrs(self) -> None: + if getattr(self, '_liger_attrs_set', False): + return + # ``variance_epsilon`` is the Liger name; HF names it either + # ``variance_epsilon`` (newer) or ``eps`` (older). ``_patch_rms_norm_module`` + # resolves the same way. + if not hasattr(self, 'variance_epsilon'): + self.variance_epsilon = getattr(self, 'eps', 1e-6) + self.offset = self._liger_offset + self.casting_mode = self._liger_casting_mode + self.in_place = self._liger_in_place + self.row_mode = self._liger_row_mode + self._liger_attrs_set = True + + def forward(self, hidden_states): + self._ensure_liger_attrs() + return LigerRMSNormFunction.apply( + hidden_states, + self.weight, + self.variance_epsilon, + self.offset, + self.casting_mode, + self.in_place, + self.row_mode, + ) + + def extra_repr(self): + self._ensure_liger_attrs() + return (f'weight_shape={tuple(self.weight.shape) if self.weight is not None else None}, ' + f'eps={self.variance_epsilon}, offset={self.offset}, ' + f'in_place={self.in_place}, row_mode={self.row_mode}') + + +class LigerRMSNormReplacement(_LigerRMSNormBase): + """llama / qwen / phi3 family: offset=0.0, casting_mode='llama', in_place=True.""" + + _liger_offset = 0.0 + _liger_casting_mode = 'llama' + _liger_in_place = True + + +class LigerRMSNormGemmaReplacement(_LigerRMSNormBase): + """gemma / gemma2 / gemma3_text: offset=1.0, casting_mode='gemma', in_place=False.""" + + _liger_offset = 1.0 + _liger_casting_mode = 'gemma' + _liger_in_place = False + + +class LigerRMSNormGemma4Replacement(_LigerRMSNormBase): + """gemma4: offset=0.0, casting_mode='gemma', in_place=False. + + Note: Gemma4RMSNorm has a ``with_scale=False`` variant (no weight, used for + attention ``v_norm``). That path falls back to a plain torch RMSNorm; this + adapter is only wired onto the scale-bearing variants — see ``liger_builtin``. + """ + + _liger_offset = 0.0 + _liger_casting_mode = 'gemma' + _liger_in_place = False + + +class LigerRMSNormQwen35Replacement(_LigerRMSNormBase): + """qwen3_5 / qwen3_5_moe: offset=1.0, casting_mode='gemma', in_place=False. + + Qwen3.5 uses residual parameterization (``weight`` init=0, scale = 1 + weight), + matching the gemma casting-mode kernel which computes + ``norm(x) * (offset + weight)`` = ``(1 + weight) * norm(x)`` with offset=1.0. + The gemma precision path (fp32 throughout, cast back at end) also matches + HF ``Qwen3_5RMSNorm.forward`` which does ``_norm(x.float()) * (1 + w.float())``. + + Using ``LigerRMSNormReplacement`` (llama cast, offset=0) here produces + ``weight * norm(x)`` ≈ 0 (since weight≈0), causing NaN in attention and loss. + """ + + _liger_offset = 1.0 + _liger_casting_mode = 'gemma' + _liger_in_place = False diff --git a/src/twinkle/kernel/liger_impls/swiglu.py b/src/twinkle/kernel/liger_impls/swiglu.py new file mode 100644 index 00000000..97b134b4 --- /dev/null +++ b/src/twinkle/kernel/liger_impls/swiglu.py @@ -0,0 +1,32 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""SwiGLU forward-replacement for Liger Kernel. + +Used as a *function-level* mapping value (string key ``'..forward'``) +so it composes with the existing SwiGLU forward-replacement pattern in +``builtin.py`` (``_add_swiglu_if_present``). Reads only ``gate_proj`` / +``up_proj`` / ``down_proj``, which every HuggingFace SwiGLU MLP variant +(Qwen2MLP, Qwen3MLP, LlamaMLP, MistralMLP, ...) already defines, so no +``__init__`` and no per-instance attribute setup is required. + +For Qwen3-MoE the expert/MLP classes differ; class replacement with Liger's +own ``LigerQwen3MoeSwiGLUMLP`` / ``LigerExperts`` is wired directly in +``liger_builtin`` since those read matching attributes. +""" +from __future__ import annotations + +from liger_kernel.ops import LigerGELUMulFunction, LigerSiLUMulFunction + + +def liger_swiglu_forward(self, x): + return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x))) + + +def liger_geglu_forward(self, x): + """GeGLU forward replacement for the gemma family (gemma / gemma2 / gemma3 / gemma4). + + Reads only ``gate_proj`` / ``up_proj`` / ``down_proj`` — the same attributes + HuggingFace GeGLU MLP variants define — so it is a safe function-level + mapping value. Uses the tanh GELU approximation, matching Liger's own + ``LigerGEGLUMLP`` and HF's gemma activation choice. + """ + return self.down_proj(LigerGELUMulFunction.apply(self.gate_proj(x), self.up_proj(x))) diff --git a/src/twinkle/loss/__init__.py b/src/twinkle/loss/__init__.py index 8e1d0e2a..d062a3fa 100644 --- a/src/twinkle/loss/__init__.py +++ b/src/twinkle/loss/__init__.py @@ -6,12 +6,14 @@ from .gkd import GKDLoss from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, SAPOLoss from .infonce import InfonceLoss +from .liger_fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss from .mse import MSELoss torch_loss_mapping = { 'mse': MSELoss, 'chunked_cross_entropy': ChunkedCrossEntropyLoss, 'cross_entropy': CrossEntropyLoss, + 'liger_fused_linear_cross_entropy': LigerFusedLinearCrossEntropyLoss, # KD losses 'gkd': GKDLoss, # RL losses diff --git a/src/twinkle/loss/liger_fused_linear_cross_entropy.py b/src/twinkle/loss/liger_fused_linear_cross_entropy.py new file mode 100644 index 00000000..408d7304 --- /dev/null +++ b/src/twinkle/loss/liger_fused_linear_cross_entropy.py @@ -0,0 +1,208 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Liger Fused Linear Cross-Entropy loss. + +Fuses the final ``lm_head`` matmul with cross-entropy so the full +``(B, T, V)`` logits tensor is never materialised — the memory + bandwidth +win for large-vocab models like Qwen3.5-35B-A3B (V ~= 151k). + +Contract with the model forward +-------------------------------- +This loss sets ``require_logits = True`` / ``require_logps = False`` so +``TransformersModel.forward`` keeps ``outputs['logits']`` and skips the +``selective_log_softmax`` pass. Under ``TransformersFusedCEPatch`` +(applied via ``task='fused_lm_ce'``) the lm_head is replaced by identity, so +``outputs['logits']`` is actually the last hidden state ``[B, T, H]``, and the +patch stashes the lm_head module under ``outputs['lm_head']``. This loss then +reads both and calls Liger's fused kernel: + + LigerFusedLinearCrossEntropy(lin_weight=lm_head.weight, + _input=hidden.view(-1, H), + target=labels.view(-1)) + +If ``outputs['lm_head']`` is absent (e.g. an eval forward that did not pass +``task='fused_lm_ce'``), the loss transparently falls back to the standard CE +computed from materialised logits — same value, just without the fusion. This +makes the loss safe to register once and use across mixed train/eval task +contexts. + +No token shift here +------------------- +Twinkle templates pre-shift labels via ``_roll_labels`` (``template/base.py``), +so ``labels[i]`` is already the target for the prediction made at position +``i``. Liger's own ``LigerForCausalLMLoss`` shifts because HF passes +unshifted labels — we must NOT re-shift, so we call Liger's +``LigerFusedLinearCrossEntropyLoss`` module (which does not shift) directly +rather than ``LigerForCausalLMLoss`` (which does). + +Hardware +-------- +Device selection is delegated to Liger's ``infer_device`` / ``select_impl`` +dispatch: this file contains no ``torch.cuda`` / ``Platform.is_npu`` probes. +The same loss class runs on CUDA (Triton kernel) and Ascend NPU (the +``backends/_ascend`` fused-linear-CE backend), mirroring the bare-impl +philosophy of ``liger_builtin``. + +FSDP2 collective ordering +------------------------ +Under accelerate-FSDP2 the first model forward lazily runs ``fully_shard`` init +(collectives), and this loss issues an ad-hoc ``lm_head.weight.full_tensor()`` +gather after the (identity) lm_head forward. If ranks enter the first fused-CE +step slightly desynchronised (asymmetric setup work), the two collective +streams enqueue in different orders across ranks and the first backward +deadlocks — e.g. rank A stuck in autograd backward while rank B has already +reached the optimizer's DTensor redistribution (Muon ``_zeropower_via_newtonschulz``). +This is a cookbook-level timing requirement, not a loss bug: the loss's +``full_tensor()`` gather is rank-symmetric. The fix is a one-time +``dist.barrier()`` + device ``synchronize()`` before the first fused-CE +training step — ``liger_bench.py`` already does a per-step ``_synchronize()`` +and ``cookbook/transformers/fsdp2.py`` does a single pre-loop drain (both +verified sufficient). This loss stays device-agnostic: it does NOT issue the +drain itself (no ``torch.cuda`` / ``Platform.is_npu`` probes); the cookbook owns +it. +""" +from __future__ import annotations + +import torch.nn.functional as F +from typing import Optional + +from twinkle import get_logger +from twinkle.data_format import LossOutput, ModelOutput +from .base import Loss + +logger = get_logger() + +# Lazily import Liger so the module loads even when liger_kernel is absent; +# the loss then raises a clear error only if actually constructed/used. +_LigerFLCEModule = None + + +def _get_liger_module(): + global _LigerFLCEModule + if _LigerFLCEModule is not None: + return _LigerFLCEModule + from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss as _Mod # noqa: F401 + _LigerFLCEModule = _Mod + return _LigerFLCEModule + + +class LigerFusedLinearCrossEntropyLoss(Loss): + """Fused lm_head + cross-entropy loss (Liger kernel). + + Args: + ignore_index: Label id treated as padding (excluded from loss); + forwarded to Liger which masks it inside the fused kernel. + reduction: 'sum' (default, matching Twinkle's CrossEntropyLoss model + default) or 'mean'. With 'sum' the loss is the per-token CE summed + over non-ignored tokens; ``calculate_loss`` accumulates this and the + metric divides by ``num_tokens`` to recover the per-token value, and + the gradient is divided by ``num_tokens`` to give a per-token-mean + gradient — identical to ``CrossEntropyLoss(reduction='sum')``. + label_smoothing: 0.0 (default) for hard labels; forwarded to Liger. + softcap: Optional logit softcapping value (e.g. Gemma); forwarded to + Liger. None disables it. + + Defensive fallback (device-agnostic) + ------------------------------------- + The Liger fused kernel is wrapped in a try/except. If it raises for any + reason — an unsupported shape on the Ascend backend, a Triton compile + failure on CUDA, a version mismatch, an OOM inside the fused kernel — the + loss logs one warning, marks the fused path broken (so subsequent calls skip + straight to the fallback without re-throwing), materialises the logits via + ``F.linear(hidden, lm_head.weight)`` (NOT ``lm_head.forward`` — that is + identity under the patch), and computes standard cross-entropy on them. So + ``--enable-liger`` is always safe on both NPU and CUDA: best case the fused + kernel saves the logits memory, worst case it degrades transparently to the + standard CE path (same numerics, just without the memory win). No + ``torch.cuda`` / ``Platform.is_npu`` probes anywhere — the guard is purely + exception-based. + """ + + # Keep outputs['logits'] (hidden states under the patch); skip logps. + require_logits = True + require_logps = False + + def __init__(self, + ignore_index: int = -100, + reduction: str = 'sum', + label_smoothing: float = 0.0, + softcap: float | None = None, + **kwargs): + super().__init__() + assert reduction in ('mean', 'sum'), f"reduction must be 'mean' or 'sum', got {reduction!r}" + self.ignore_index = ignore_index + self.reduction = reduction + self.label_smoothing = label_smoothing + self.softcap = softcap + self._liger = _get_liger_module()( + ignore_index=ignore_index, + reduction=reduction, + label_smoothing=label_smoothing, + softcap=softcap, + ) + self._fused_broken = False + self._warned = False + + def _materialise_ce(self, inputs, outputs, lm_head, hidden, **kwargs): + """Materialise logits from the stashed lm_head weight (NOT its forward, + which is identity under the patch) and run standard cross-entropy.""" + from .cross_entropy import CrossEntropyLoss + weight = lm_head.weight + if hasattr(weight, 'full_tensor'): + weight = weight.full_tensor() + bias = getattr(lm_head, 'bias', None) + if bias is not None and hasattr(bias, 'full_tensor'): + bias = bias.full_tensor() + logits = F.linear(hidden.reshape(-1, hidden.shape[-1]), weight, bias=bias) + out = dict(outputs) + out['logits'] = logits.view(*hidden.shape[:-1], -1) + return CrossEntropyLoss(ignore_index=self.ignore_index, reduction=self.reduction)(inputs, out, **kwargs) + + def __call__(self, inputs, outputs: ModelOutput, **kwargs): + labels = inputs['labels'] + lm_head = outputs.get('lm_head') + have_head = lm_head is not None and getattr(lm_head, 'weight', None) is not None + + if have_head and not self._fused_broken: + # ── Fused path: hidden states + lm_head.weight ── + hidden = outputs['logits'] # [B, T, H] under TransformersFusedCEPatch + hidden_size = hidden.shape[-1] + hidden_flat = hidden.reshape(-1, hidden_size) + labels_flat = labels.reshape(-1).to(hidden_flat.device) + # Under EP/FSDP2 the lm_head weight may be a sharded DTensor: the + # identity-forward patch (TransformersFusedCEPatch) skips the normal + # lm_head forward that would auto-unshard it, so materialise the + # full weight here. ``full_tensor()`` is a collective gather; under + # LoRA the lm_head is frozen so only the forward gather is needed + # (grad flows to hidden, connecting back to the LoRA backbone). The + # duck-typed ``hasattr`` check keeps this device/strategy-agnostic — + # regular tensors (no EP/FSDP) pass through unchanged. + weight = lm_head.weight + if hasattr(weight, 'full_tensor'): + weight = weight.full_tensor() + try: + loss = self._liger(weight, hidden_flat, labels_flat) + except Exception as e: # noqa: BLE001 — defensive, device-agnostic + if not self._warned: + self._warned = True + logger.warning( + '[LigerFusedLinearCrossEntropyLoss] fused kernel raised %r; ' + 'falling back to materialised cross-entropy (device-agnostic defensive ' + 'fallback). Subsequent calls skip the fused path. This is expected when ' + 'the device backend does not support the shape (e.g. Ascend chunked path ' + 'on some shapes, CUDA Triton compile failure).', e) + self._fused_broken = True + return self._materialise_ce(inputs, outputs, lm_head, hidden, **kwargs) + num_tokens = (labels != self.ignore_index).to(loss.device).sum().clamp(min=1) + return LossOutput(loss=loss, num_tokens=num_tokens) + + if have_head: + # fused path broken on a previous call -> materialise via the stashed weight + return self._materialise_ce(inputs, outputs, lm_head, outputs['logits'], **kwargs) + + # ── Fallback: standard CE from materialised logits (eval w/o patch) ── + # No `lm_head` in outputs means the forward did not run under + # TransformersFusedCEPatch (e.g. an eval pass); outputs['logits'] then + # already holds real logits. label_smoothing is a fused-kernel-only knob + # here; the eval fallback uses hard labels (the common eval case). + from .cross_entropy import CrossEntropyLoss + return CrossEntropyLoss(ignore_index=self.ignore_index, reduction=self.reduction)(inputs, outputs, **kwargs) diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index d45ac618..264503e8 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -997,6 +997,27 @@ def gather_loss_tensors( return inputs, outputs labels = inputs.get('labels') logps = outputs.get('logps') + # Fused-linear-CE path: the model forward skipped the lm_head GEMM and + # stashed the lm_head module under ``outputs['lm_head']`` (see + # ``TransformersFusedCEPatch``), so ``outputs['logps']`` is absent and + # ``outputs['logits']`` holds the last hidden state (a per-rank seq + # shard under SP). Gather hidden states + labels across the sequence + # group so every rank computes the same full loss — mirrors the logps + # gather below. The hidden tensor ([B, T_local, H]) is tiny vs logits, + # and the fused loss still chunks the GEMM internally, so the memory + # win is preserved. + if (logps is None and labels is not None and 'lm_head' in outputs and torch.is_tensor(outputs.get('logits')) + and outputs['logits'].dim() >= 2): + hidden = outputs['logits'] + inputs = copy(inputs) + outputs = copy(outputs) + real_position_ids = sequence_parallel.real_position_ids + gathered_hidden, gathered_labels = GatherLoss.apply(hidden, labels, 1, real_position_ids) + gathered_hidden = self._trim_gathered_sequence_padding(gathered_hidden, real_position_ids) + gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) + outputs['logits'] = gathered_hidden + inputs['labels'] = gathered_labels + return inputs, outputs if labels is None or logps is None: return inputs, outputs if not torch.is_tensor(logps) or logps.dim() < 2: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index d88c9684..017a515b 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -60,7 +60,12 @@ def _resolve_task_context(model, task): if task == 'embedding': from twinkle.patch.transformers_emb import TransformersEmbeddingPatch return apply_context(model, TransformersEmbeddingPatch()) - raise ValueError(f'Unknown task={task!r}; expected one of: causal_lm, embedding.') + if task == 'fused_lm_ce': + # Skip the lm_head GEMM so a fused-linear-CE loss can take over (Liger). + # Device-agnostic: Liger self-dispatches the fused kernel across CUDA/NPU. + from twinkle.patch.transformers_fused_ce import TransformersFusedCEPatch + return apply_context(model, TransformersFusedCEPatch()) + raise ValueError(f'Unknown task={task!r}; expected one of: causal_lm, embedding, fused_lm_ce.') @dataclass diff --git a/src/twinkle/patch/transformers_fused_ce.py b/src/twinkle/patch/transformers_fused_ce.py new file mode 100644 index 00000000..a1eeb47a --- /dev/null +++ b/src/twinkle/patch/transformers_fused_ce.py @@ -0,0 +1,126 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Patch a HF transformers causal LM to skip the lm_head matmul so the loss +layer can fuse it with cross-entropy (Liger Fused Linear-CE). + +Two mutations applied to the model (both reverted by ``unpatch``): + +1. ``lm_head.forward`` is replaced with identity, so the wrapped model returns + the final hidden states under ``output.logits`` instead of the full + ``(B, T, V)`` logits tensor. This skips the vocab projection GEMM and the + materialisation of the logits tensor, which is the memory + bandwidth win + that fused-linear-CE unlocks. +2. A forward hook on the lm-head-bearing submodule preserves the full HF + ``ModelOutput`` (so MoE ``aux_loss`` / ``router_logits`` / ``past_key_values`` + survive) and additionally stashes a reference to the lm_head module under + ``outputs['lm_head']`` so the loss can read ``.weight`` without holding the + model. + +This patch is deliberately hardware-agnostic: it contains no device probes +(``torch.cuda`` / ``Platform.is_npu`` / ``infer_device``). Device selection of +the fused-CE kernel is handled by Liger's own ``infer_device`` / +``select_impl`` dispatch (CUDA-Triton vs Ascend backend), exactly like the +``liger_builtin`` bundle emits bare impls. See ``liger.py`` docstring. + +The patch is applied per-forward via ``apply_context`` (see +``_resolve_task_context`` with ``task='fused_lm_ce'``), so it composes with +Twinkle's SP / routing-replay / FSDP2 / processor-postprocess orchestration +without touching the forward body. Under FSDP2 the lm_head params are unsharded +by FSDP's pre-forward hook before the (identity) forward runs, so +``lm_head.weight`` read by the loss is the full weight. + +FSDP2 collective ordering: under accelerate-FSDP2 the first fused-CE forward +lazily runs ``fully_shard`` init (collectives), and the loss's ad-hoc +``lm_head.weight.full_tensor()`` gather (issued after the identity forward) can +race with them if ranks enter the first fused-CE step slightly desynchronised +(asymmetric setup work) — the two collective streams enqueue in different orders +across ranks and the first backward deadlocks (rank A in autograd backward, +rank B already in the optimizer's DTensor redistribution). Cookbooks using +``task='fused_lm_ce'`` under accelerate-FSDP2 should issue a one-time +``dist.barrier()`` + device ``synchronize()`` before the first fused-CE training +step (see ``liger_fused_linear_cross_entropy.py`` and +``cookbook/transformers/fsdp2.py``; ``liger_bench.py`` already does a per-step +``_synchronize()``). This patch itself stays device-free (no ``torch.cuda`` / +``Platform.is_npu`` probes — enforced by ``test_patch_is_device_free``). + +Composing with ``LigerFusedLinearCrossEntropyLoss`` (see ``twinkle.loss``): +the loss dispatches on the presence of ``outputs['lm_head']`` — present => +fused path (hidden_states + lm_head.weight); absent => standard CE from +materialised logits. So eval forwards that omit ``task='fused_lm_ce'`` still +work with the same loss instance. +""" +from types import MethodType +from typing import TYPE_CHECKING, Optional + +from twinkle.patch import Patch + +if TYPE_CHECKING: + import torch + +# Reuse the lm-head discovery list from the embedding patch verbatim — same +# families (Qwen / Llama / Phi3 with 'lm_head', GPT-NeoX with 'embed_out', ...). +from twinkle.patch.transformers_emb import get_lm_head_model + +_LM_HEADS = ['lm_head', 'embed_out', 'output_layer', 'output'] + + +def _identity_forward(self, hidden_states): + """Return hidden states unchanged; the vocab GEMM is deferred to the loss.""" + return hidden_states + + +class TransformersFusedCEPatch(Patch): + """Skip the lm_head matmul so a fused-linear-CE loss can take over. + + Reversible via ``unpatch``. Device-agnostic: no ``cuda``/``npu`` branching. + """ + + def __call__(self, module, *args, **kwargs): + from torch.nn import Module + lm_head_model = get_lm_head_model(module, lm_heads=_LM_HEADS) + + head: Optional[Module] = None + for name in _LM_HEADS: + if hasattr(lm_head_model, name): + head = getattr(lm_head_model, name) + break + assert head is not None, 'Cannot find the proper lm_head name' + + # Save originals BEFORE mutation so unpatch can restore them verbatim. + self._head = head + self._origin_forward = head.forward + head.forward = MethodType(_identity_forward, head) + + # Closure captures `head` so the hook can stash the module reference + # without polluting the lm_head_model with attributes. + def _stash_lm_head_hook(module, args, kwargs, output): + # Preserve the full HF ModelOutput (aux_loss / router_logits / + # past_key_values / ...); only add the lm_head module reference. + # output.logits is hidden_states here (identity forward). + if hasattr(output, 'items'): + out = dict(output.items()) + elif isinstance(output, dict): + out = dict(output) + elif isinstance(output, (tuple, list)): + # HF non-return_dict path: (logits, ...) — keep logits only. + out = {'logits': output[0] if len(output) else None} + else: + out = {'logits': output} + out['lm_head'] = head + return out + + self._hook_handle = lm_head_model.register_forward_hook(_stash_lm_head_hook, with_kwargs=True) + return module + + def unpatch(self, module, *args, **kwargs): + handle = getattr(self, '_hook_handle', None) + if handle is not None: + handle.remove() + self._hook_handle = None + + head = getattr(self, '_head', None) + origin = getattr(self, '_origin_forward', None) + if head is not None and origin is not None: + head.forward = origin + self._origin_forward = None + self._head = None + return module diff --git a/src/twinkle/utils/framework.py b/src/twinkle/utils/framework.py index 7bdb9b64..5bd1c0a3 100644 --- a/src/twinkle/utils/framework.py +++ b/src/twinkle/utils/framework.py @@ -163,27 +163,27 @@ def set_device(local_rank: Union[int, str] = None) -> None: def seed_everything(seed: Optional[int] = 42, deterministic: bool = False): random.seed(seed) np.random.seed(seed) + import torch + torch.manual_seed(seed) + if Torch.is_gpu_available(): - import torch - torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) - if Torch.is_npu_available(): - import torch_npu - torch.npu.manual_seed_all(seed) + if Torch.is_npu_available(): + import torch_npu + torch.npu.manual_seed_all(seed) - if deterministic: - torch.use_deterministic_algorithms(True) + if deterministic: + torch.use_deterministic_algorithms(True, warn_only=True) + if Torch.is_gpu_available(): os.environ['CUDA_LAUNCH_BLOCKING'] = '1' os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8' os.environ['FLASH_ATTENTION_DETERMINISTIC'] = '1' - torch.use_deterministic_algorithms(True, warn_only=True) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False - - if Torch.is_npu_available(): - os.environ['ASCEND_LAUNCH_BLOCKING'] = '1' - os.environ['HCCL_DETERMINISTIC'] = '1' + if Torch.is_npu_available(): + os.environ['ASCEND_LAUNCH_BLOCKING'] = '1' + os.environ['HCCL_DETERMINISTIC'] = '1' @staticmethod def to_local_tensor(tensor: 'torch.Tensor') -> 'torch.Tensor': diff --git a/tests/kernel/test_liger_builtin.py b/tests/kernel/test_liger_builtin.py new file mode 100644 index 00000000..893df205 --- /dev/null +++ b/tests/kernel/test_liger_builtin.py @@ -0,0 +1,116 @@ +import os + +import torch +import torch.nn as nn +import pytest + +import twinkle.kernel as k + +try: + import liger_kernel # noqa: F401 + _HAS_LIGER = True +except ImportError: + _HAS_LIGER = False + +requires_liger = pytest.mark.skipif(not _HAS_LIGER, reason='liger_kernel not installed') + + +@requires_liger +def test_liger_builtin_returns_dict(): + bundle = k.liger_builtin() + assert isinstance(bundle, dict) + assert len(bundle) > 0 + + +@requires_liger +def test_liger_builtin_values_are_bare_impls(): + """unlike npu_builtin (which wraps in {'npu': impl}), liger_builtin emits + *bare* impls because Liger self-dispatches across CUDA/NPU internally.""" + for key, value in k.liger_builtin().items(): + assert not isinstance(value, dict), f'value for {key!r} must be bare, got a device-dict' + + +@requires_liger +def test_liger_builtin_compose_with_user_override(): + sentinel = object() + merged = {**k.liger_builtin(), 'fake.module.path.fn': sentinel} + assert merged['fake.module.path.fn'] is sentinel + + +@requires_liger +def test_liger_builtin_skips_missing_modeling_modules(monkeypatch): + """If every transformers modeling module is unimportable, the bundle must + still return a dict (empty) rather than raising.""" + import twinkle.kernel.liger as liger_mod + + monkeypatch.setattr(liger_mod, '_import_optional', lambda name: None) + bundle = liger_mod.liger_builtin() + assert isinstance(bundle, dict) + assert bundle == {} + + +@requires_liger +def test_liger_builtin_no_global_side_effects(): + """Liger's bundle must not mutate process-global state — contrast with + npu_builtin which installs a global SDPA override. Concretely, the + ``torch.nn.functional.cross_entropy`` callable (the thing a fused-linear-CE + swap would touch) must be byte-identical before and after.""" + import torch.nn.functional as F + + before = F.cross_entropy + k.liger_builtin() + assert F.cross_entropy is before + + +def test_kernelize_class_replacement_applies(): + """kernelize must swap __class__ on a module whose type matches a bundle + class key (independent of which families are installed).""" + import torch.nn as nn + + class _FakeHFNorm(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(4)) + + def forward(self, x): + return x + + class _FakeLigerImpl(nn.Module): + def forward(self, x): + return x + 1 + + parent = nn.Sequential(_FakeHFNorm()) + out = k.kernelize(parent, {_FakeHFNorm: _FakeLigerImpl}) + assert out is parent + assert type(parent[0]) is _FakeLigerImpl + + +# ── CLI toggle (liger-independent: part of twinkle core) ────────────────────── + + +def test_enable_liger_defaults_false(): + from twinkle.cli import ModelArgs + + assert ModelArgs().enable_liger is False + + +def test_enable_liger_toggle_via_cli(): + from twinkle.cli import CLI + + on = CLI.from_args(['--enable-liger']) + off = CLI.from_args(['--no-enable-liger']) + assert on.model.enable_liger is True + assert off.model.enable_liger is False + + +def test_enable_liger_env_var(): + from twinkle.cli import CLI + + old = os.environ.copy() + try: + os.environ.clear() + os.environ['TWINKLE_ENABLE_LIGER'] = 'true' + assert CLI.from_args(argv=[]).model.enable_liger is True + finally: + os.environ.clear() + os.environ.update(old) diff --git a/tests/kernel/test_public_api.py b/tests/kernel/test_public_api.py index f9a17a2a..2d01f116 100644 --- a/tests/kernel/test_public_api.py +++ b/tests/kernel/test_public_api.py @@ -1,8 +1,9 @@ -def test_public_exports_exactly_three_symbols(): +def test_public_exports_exactly_four_symbols(): import twinkle.kernel as k - assert sorted(k.__all__) == ['hub', 'kernelize', 'npu_builtin'] + assert sorted(k.__all__) == ['hub', 'kernelize', 'liger_builtin', 'npu_builtin'] assert callable(k.kernelize) assert callable(k.npu_builtin) + assert callable(k.liger_builtin) assert callable(k.hub) diff --git a/tests/loss/test_liger_fused_linear_ce.py b/tests/loss/test_liger_fused_linear_ce.py new file mode 100644 index 00000000..9437f617 --- /dev/null +++ b/tests/loss/test_liger_fused_linear_ce.py @@ -0,0 +1,263 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tests for LigerFusedLinearCrossEntropyLoss. + +CPU tests (always run) cover registration + the eval fallback path (standard +CE from materialised logits). Device tests (Liger kernel on CUDA/NPU) cover the +fused-path numerical correctness vs CrossEntropyLoss. +""" +import pytest +import torch +import torch.nn.functional as F + +from twinkle.loss import CrossEntropyLoss, LigerFusedLinearCrossEntropyLoss, torch_loss_mapping + +try: + import liger_kernel # noqa: F401 + _HAS_LIGER = True +except ImportError: + _HAS_LIGER = False + +requires_liger = pytest.mark.skipif(not _HAS_LIGER, reason='liger_kernel not installed') + + +def _has_accelerator(): + if torch.cuda.is_available(): + return 'cuda' + try: + import torch_npu # noqa: F401 + if torch_npu.npu.is_available(): + return 'npu' + except ImportError: + pass + return None + + +requires_device = pytest.mark.skipif(not _has_accelerator(), reason='no CUDA/NPU available for Liger kernel') + + +# ── CPU: registration + flags ──────────────────────────────────────────────── + + +def test_loss_registered_in_mapping(): + assert torch_loss_mapping['liger_fused_linear_cross_entropy'] is LigerFusedLinearCrossEntropyLoss + + +def test_loss_flags_keep_logits_skip_logps(): + # require_logits=True so forward keeps outputs['logits'] (= hidden states under the patch) + # require_logps=False so forward skips selective_log_softmax + assert LigerFusedLinearCrossEntropyLoss.require_logits is True + assert LigerFusedLinearCrossEntropyLoss.require_logps is False + + +# ── CPU: fallback path (no lm_head in outputs -> standard CE) ────────────────── + + +def _make_logits_batch(bs=4, seq=8, vocab=20, seed=42): + torch.manual_seed(seed) + logits = torch.randn(bs, seq, vocab, dtype=torch.float32) + labels = torch.randint(0, vocab, (bs, seq)) + # mark some tokens ignored + labels[:, seq // 2:] = -100 + return labels, logits + + +def test_fallback_matches_cross_entropy_loss_when_no_lm_head(): + labels, logits = _make_logits_batch() + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + ce = CrossEntropyLoss(ignore_index=-100, reduction='mean') + out_fused = fused({'labels': labels}, {'logits': logits.clone()}) + out_ce = ce({'labels': labels}, {'logits': logits.clone()}) + assert torch.allclose(out_fused['loss'], out_ce['loss'], atol=1e-6) + + +def test_fallback_reduction_sum_num_tokens(): + labels, logits = _make_logits_batch() + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='sum') + out = fused({'labels': labels}, {'logits': logits.clone()}) + expected_tokens = (labels != -100).sum().item() + assert out['num_tokens'].item() == expected_tokens + + +def test_fallback_all_ignored_yields_zero_loss(): + labels = torch.tensor([[-100, -100], [-100, -100]]) + logits = torch.randn(2, 2, 5, dtype=torch.float32) + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + out = fused({'labels': labels}, {'logits': logits}) + assert out['loss'].item() == pytest.approx(0.0, abs=1e-6) + + +# ── Device: fused-path numerical correctness vs CrossEntropyLoss ─────────────── + + +@requires_liger +@requires_device +def test_fused_path_matches_cross_entropy_on_device(): + device = _has_accelerator() + torch.manual_seed(123) + bs, seq, hidden, vocab = 2, 16, 8, 32 + # hidden states + a real lm_head weight + hidden_states = torch.randn(bs, seq, hidden, device=device, dtype=torch.float32, requires_grad=True) + lm_head = torch.nn.Linear(hidden, vocab, bias=False, device=device, dtype=torch.float32) + labels = torch.randint(0, vocab, (bs, seq), device=device) + labels[:, seq // 2:] = -100 + + # Reference: materialise logits via lm_head, then CrossEntropyLoss. + with torch.no_grad(): + logits = lm_head(hidden_states.detach()) + ref = CrossEntropyLoss(ignore_index=-100, reduction='mean')({'labels': labels}, {'logits': logits}) + + # Fused: feed hidden states + lm_head.weight, no logits materialised. + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + out = fused({'labels': labels}, {'logits': hidden_states.detach(), 'lm_head': lm_head}) + + # Liger accumulates in fp32 internally; CrossEntropyLoss uses the (fp32 here) + # log_softmax path. They must agree to a tight tolerance. + assert torch.allclose(out['loss'].cpu(), ref['loss'].cpu(), atol=1e-4, rtol=1e-3), ( + f'fused={out["loss"].item():.6f} vs ref={ref["loss"].item():.6f}') + + +@requires_liger +@requires_device +def test_fused_path_backward_grad_flows_to_hidden_states(): + """Gradient must reach hidden_states (connects back to the LoRA backbone); + lm_head is frozen so grad_weight need not be populated.""" + device = _has_accelerator() + torch.manual_seed(7) + bs, seq, hidden, vocab = 2, 8, 8, 16 + hidden_states = torch.randn(bs, seq, hidden, device=device, dtype=torch.float32, requires_grad=True) + lm_head = torch.nn.Linear(hidden, vocab, bias=False, device=device, dtype=torch.float32) + for p in lm_head.parameters(): + p.requires_grad_(False) # frozen lm_head, LoRA-style + labels = torch.randint(0, vocab, (bs, seq), device=device) + + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + out = fused({'labels': labels}, {'logits': hidden_states, 'lm_head': lm_head}) + out['loss'].backward() + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + assert torch.isfinite(hidden_states.grad).all() + + +@requires_liger +@requires_device +def test_fused_path_ignores_padding_tokens(): + """Labels == -100 must be excluded exactly (Liger's ignore_index).""" + device = _has_accelerator() + torch.manual_seed(99) + bs, seq, hidden, vocab = 1, 6, 8, 12 + hidden_states = torch.randn(bs, seq, hidden, device=device, dtype=torch.float32, requires_grad=True) + lm_head = torch.nn.Linear(hidden, vocab, bias=False, device=device, dtype=torch.float32) + labels = torch.randint(0, vocab, (bs, seq), device=device) + # blank half the tokens + mask = torch.zeros(bs, seq, dtype=torch.bool, device=device) + mask[:, seq // 2:] = True + labels_full = labels.clone() + labels_masked = labels.clone() + labels_masked[mask] = -100 + + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + out_masked = fused({'labels': labels_masked}, {'logits': hidden_states.detach(), 'lm_head': lm_head}) + out_full = fused({'labels': labels_full}, {'logits': hidden_states.detach(), 'lm_head': lm_head}) + # masked loss is over half the tokens; both finite and non-negative. + assert out_masked['loss'].item() >= 0 + assert out_full['loss'].item() >= 0 + assert torch.isfinite(out_masked['loss']) and torch.isfinite(out_full['loss']) + + +@requires_liger +def test_construction_without_device_is_lazy(): + """Building the loss must not require a device — the kernel is only + dispatched when __call__ runs under the fused path.""" + loss = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean', label_smoothing=0.1) + assert loss.reduction == 'mean' + assert loss.label_smoothing == 0.1 + + +# ── Device-agnostic defensive fallback (fused kernel raises) ──────────────── + + +@requires_liger +@requires_device +def test_fused_kernel_failure_falls_back_to_materialised_ce(): + """When the Liger fused kernel raises (any device/shape), the loss must + fall back to materialised-CE via the stashed lm_head weight (NOT its + identity-patched forward) and match CrossEntropyLoss; subsequent calls must + skip the fused path (broken flag).""" + device = _has_accelerator() + torch.manual_seed(5) + bs, seq, hidden, vocab = 2, 6, 8, 14 + hidden_states = torch.randn(bs, seq, hidden, device=device, dtype=torch.float32, requires_grad=True) + lm_head = torch.nn.Linear(hidden, vocab, bias=False, device=device, dtype=torch.float32) + for p in lm_head.parameters(): + p.requires_grad_(False) + labels = torch.randint(0, vocab, (bs, seq), device=device) + + # Reference: materialise logits via lm_head, standard CE. + with torch.no_grad(): + ref_logits = lm_head(hidden_states.detach()) + ref = CrossEntropyLoss(ignore_index=-100, reduction='mean')({'labels': labels}, {'logits': ref_logits}) + + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + + # Force the fused kernel to raise on the first call by swapping the Liger + # module for a raising callable (assigning .__call__ on an nn.Module instance + # doesn't work — Module.__call__ is class-level). + class _Boom: + def __call__(self, *a, **k): + raise RuntimeError('simulated fused-kernel failure (e.g. unsupported shape)') + + original_liger = fused._liger + fused._liger = _Boom() + + # The fallback warning is emitted via twinkle's logger (not the `warnings` + # module), so capture log records. + import logging as _logging + from twinkle import get_logger as _get_logger + + class _LogCapture(_logging.Handler): + def __init__(self): + super().__init__(level=_logging.WARNING) + self.records: list[str] = [] + + def emit(self, record): + self.records.append(record.getMessage()) + + _cap = _LogCapture() + _get_logger().addHandler(_cap) + try: + out = fused({'labels': labels}, {'logits': hidden_states.detach(), 'lm_head': lm_head}) + warned = any('falling back to materialised cross-entropy' in m for m in _cap.records) + finally: + _get_logger().removeHandler(_cap) + + # Fallback matches the materialised-CE reference. + assert torch.allclose(out['loss'].cpu(), ref['loss'].cpu(), atol=1e-4, rtol=1e-3), ( + f'fallback={out["loss"].item():.6f} vs ref={ref["loss"].item():.6f}') + assert warned, 'expected the defensive-fallback warning' + + # Restore the real kernel; the broken flag must persist (no retry). + fused._liger = original_liger + out2 = fused({'labels': labels}, {'logits': hidden_states.detach(), 'lm_head': lm_head}) + assert fused._fused_broken is True + assert torch.allclose(out2['loss'].cpu(), ref['loss'].cpu(), atol=1e-4, rtol=1e-3), ( + 'subsequent call must go straight to the fallback without re-invoking the fused kernel') + + +@requires_liger +@requires_device +def test_fused_kernel_success_does_not_mark_broken(): + """Sanity: a successful fused call must NOT set the broken flag (fallback + path stays inactive).""" + device = _has_accelerator() + torch.manual_seed(8) + bs, seq, hidden, vocab = 2, 8, 8, 16 + hidden_states = torch.randn(bs, seq, hidden, device=device, dtype=torch.float32, requires_grad=True) + lm_head = torch.nn.Linear(hidden, vocab, bias=False, device=device, dtype=torch.float32) + for p in lm_head.parameters(): + p.requires_grad_(False) + labels = torch.randint(0, vocab, (bs, seq), device=device) + fused = LigerFusedLinearCrossEntropyLoss(ignore_index=-100, reduction='mean') + fused({'labels': labels}, {'logits': hidden_states.detach(), 'lm_head': lm_head}) + assert fused._fused_broken is False + + diff --git a/tests/patch/__init__.py b/tests/patch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/patch/test_fused_ce_patch.py b/tests/patch/test_fused_ce_patch.py new file mode 100644 index 00000000..74a19e3c --- /dev/null +++ b/tests/patch/test_fused_ce_patch.py @@ -0,0 +1,118 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tests for TransformersFusedCEPatch (device-free: runs on CPU). + +These cover the patch mechanics — identity lm_head forward, lm_head module +stashed into outputs, full HF output preserved (aux_loss etc.), and reversibility. +The fused-loss numerical correctness (which needs the Liger kernel on a GPU/NPU) +lives in tests/loss/test_liger_fused_linear_ce.py. +""" +import torch +import torch.nn as nn +import pytest + +from twinkle.patch import apply_context +from twinkle.patch.transformers_fused_ce import TransformersFusedCEPatch, _identity_forward +from twinkle.patch.transformers_emb import get_lm_head_model + + +class _FakeLMHeadModel(nn.Module): + """Minimal HF-ForCausalLM-like module: a backbone returning hidden states + via .logits, plus an lm_head. Mirrors the structure the patch expects.""" + + def __init__(self, hidden_dim=8, vocab=16): + super().__init__() + self.embed = nn.Linear(hidden_dim, hidden_dim, bias=False) + self.lm_head = nn.Linear(hidden_dim, vocab, bias=False) + + def forward(self, input_ids=None, **kwargs): + bs, seq = input_ids.shape + h = self.embed(torch.randn(bs, seq, self.embed.in_features, dtype=torch.float32)) + # identity under the patch -> .logits is h; otherwise full logits + logits = self.lm_head(h) + return {'logits': logits, 'aux_loss': torch.tensor(0.7), 'past_key_values': None} + + +def _make_model(): + torch.manual_seed(0) + return _FakeLMHeadModel() + + +def test_identity_forward_returns_input_unchanged(): + head = nn.Linear(4, 4) + x = torch.randn(2, 4) + y = _identity_forward(head, x) + assert y is x # exact same tensor object; no copy, no vocab GEMM + + +def test_patch_replaces_lm_head_with_identity_and_stashes_head(): + model = _make_model() + input_ids = torch.randint(0, 10, (1, 3)) + with apply_context(model, TransformersFusedCEPatch()): + # Under the patch, lm_head.forward is identity. + assert model.lm_head.forward.__func__ is _identity_forward + out = model(input_ids=input_ids) + # The hook stashes the lm_head module and preserves aux_loss. + assert out['lm_head'] is model.lm_head + assert torch.equal(out['aux_loss'], torch.tensor(0.7)) + # logits == hidden states (identity, no vocab GEMM): vocab dim dropped. + assert out['logits'].shape[-1] == 8 # hidden, not vocab + assert 'past_key_values' in out # full HF output preserved + + +def test_patch_unpatch_restores_original_forward(): + """Functional check (bound-method `is` is fragile): under the patch the + vocab GEMM is skipped; after unpatch it runs again and the lm_head stashed + reference is gone.""" + model = _make_model() + input_ids = torch.randint(0, 10, (1, 3)) + + # Before patch: full vocab logits. + out_before = model(input_ids=input_ids) + assert out_before['logits'].shape[-1] == 16 # vocab + assert 'lm_head' not in out_before + + with apply_context(model, TransformersFusedCEPatch()): + out_patched = model(input_ids=input_ids) + assert out_patched['logits'].shape[-1] == 8 # hidden (identity) + assert 'lm_head' in out_patched + + # After unpatch: back to full vocab logits, no stashed lm_head. + out_after = model(input_ids=input_ids) + assert out_after['logits'].shape[-1] == 16 # vocab restored + assert 'lm_head' not in out_after + + +def test_patch_unpatch_on_exception(): + model = _make_model() + input_ids = torch.randint(0, 10, (1, 3)) + with pytest.raises(RuntimeError): + with apply_context(model, TransformersFusedCEPatch()): + raise RuntimeError('boom') + # unpatch must run even on exception: vocab GEMM restored. + out = model(input_ids=input_ids) + assert out['logits'].shape[-1] == 16 + assert 'lm_head' not in out + + +def test_patch_is_device_free(): + """The patch source must contain no device probes in actual code + (hardware-unaware contract). Parsed via ast so docstrings/comments don't + trip the check.""" + import ast + import inspect + from twinkle.patch import transformers_fused_ce as mod + tree = ast.parse(inspect.getsource(mod)) + probes = {'torch', 'cuda', 'npu', 'is_npu_available', 'infer_device', 'select_impl'} + for node in ast.walk(tree): + # Attribute access like torch.cuda.is_available / Platform.is_npu + if isinstance(node, ast.Attribute): + chain = [] + cur = node + while isinstance(cur, ast.Attribute): + chain.append(cur.attr) + cur = cur.value + if isinstance(cur, ast.Name): + chain.append(cur.id) + # Names that must not appear as the root of an attribute chain. + assert cur.id not in ('torch', 'Platform'), \ + f'patch references device probe: {".".join(reversed(chain))}' diff --git a/tests/transformers/test_qwen35_fla_bwd_precision.py b/tests/transformers/test_qwen35_fla_bwd_precision.py new file mode 100644 index 00000000..6bdd83c8 --- /dev/null +++ b/tests/transformers/test_qwen35_fla_bwd_precision.py @@ -0,0 +1,300 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""FLA on/off backward precision test for Qwen3.5 linear attention. + +Compares FLA ON (MindSpeed Triton bf16) vs FLA OFF (torch fp32 reference) +on a tiny single-device model. Uses mean-abs-diff as bias detector and +aggregate grad-norm as scale guard. +""" +from __future__ import annotations + +import copy +import os +import sys +import pytest +import torch +import torch.nn.functional as F + +try: + import torch_npu # noqa: F401 + _HAS_NPU = True +except ImportError: + _HAS_NPU = False + +try: + from transformers import Qwen3_5ForCausalLM, Qwen3_5MoeForCausalLM, Qwen3_5TextConfig, Qwen3_5MoeTextConfig + _HAS_QWEN35 = True +except Exception: + Qwen3_5ForCausalLM = None + Qwen3_5MoeForCausalLM = None + Qwen3_5TextConfig = None + Qwen3_5MoeTextConfig = None + _HAS_QWEN35 = False + +# Tolerances: bf16 kernel vs fp32 reference. Hard-fail on mean-abs-diff +# (systematic bias) and aggregate grad-norm (scale). Max-abs is diagnostic only. +LOSS_ATOL = 5e-3 +LOGITS_MEAN_ATOL = 2e-2 +GRAD_MEAN_ATOL = 5e-3 +GRAD_NORM_RTOL = 2.5e-1 + + +def _device() -> torch.device: + return torch.device('npu:0') if _HAS_NPU else torch.device('cpu') + + +def _dtype() -> torch.dtype: + return torch.bfloat16 + + +def _seed(seed: int) -> None: + os.environ.setdefault('PYTHONHASHSEED', str(seed)) + torch.manual_seed(seed) + if _HAS_NPU: + torch_npu.npu.manual_seed(seed) + # Flatten determinism for CANN matmul/conv kernels where supported. + try: + torch_npu.npu.set_allow_internal_format(False) + except Exception: + pass + + +def _common_config_kwargs(layer_types): + return dict( + vocab_size=128, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + linear_conv_kernel_dim=4, + linear_key_head_dim=16, + linear_value_head_dim=16, + linear_num_key_heads=2, + linear_num_value_heads=4, + layer_types=layer_types, + intermediate_size=256, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + attention_dropout=0.0, + use_cache=False, + ) + + +def _build_tiny(model_kind: str, layer_types, device: torch.device, seed: int) -> torch.nn.Module: + _seed(seed) + if model_kind == 'dense': + cfg = Qwen3_5TextConfig(**_common_config_kwargs(layer_types)) + elif model_kind == 'moe': + cfg = Qwen3_5MoeTextConfig( + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + num_experts=4, + num_experts_per_tok=2, + output_router_logits=False, + **_common_config_kwargs(layer_types), + ) + else: + raise ValueError(f'Unknown model kind: {model_kind}') + cfg._attn_implementation = 'sdpa' + cls = Qwen3_5ForCausalLM if model_kind == 'dense' else Qwen3_5MoeForCausalLM + model = cls(cfg) + model.to(device=device, dtype=_dtype()) + model.eval() + _force_fla_off(model) + return model + + +def _build_batch(device: torch.device, seq_len: int): + torch.manual_seed(0) + ids = torch.randint(3, 128, (2, seq_len), device=device, dtype=torch.long) + labels = torch.randint(3, 128, (2, seq_len), device=device, dtype=torch.long) + return ids, labels + + +def _force_fla_off(model: torch.nn.Module) -> None: + from transformers.models.qwen3_5.modeling_qwen3_5 import torch_chunk_gated_delta_rule + for layer in model.model.layers: + la = getattr(layer, 'linear_attn', None) + if la is None: + continue + la.chunk_gated_delta_rule = torch_chunk_gated_delta_rule + la.causal_conv1d_fn = None + if getattr(la, '_twinkle_npu_patched', False): + delattr(la, '_twinkle_npu_patched') + + +def _force_fla_on(model: torch.nn.Module) -> int: + import importlib + import twinkle.kernel.npu_impls.fla as fla_mod + importlib.reload(fla_mod) + return fla_mod.apply_qwen3_5_fla(model) + + +def _run_forward_backward(model: torch.nn.Module, ids, labels): + out = model(input_ids=ids).logits + loss = F.cross_entropy(out.float().reshape(-1, out.size(-1)), labels.reshape(-1)) + # Zero existing grads + model.zero_grad(set_to_none=True) + loss.backward() + grads = {} + for i, layer in enumerate(model.model.layers): + la = getattr(layer, 'linear_attn', None) + if la is None: + continue + for name in ('in_proj_qkv', 'in_proj_z', 'in_proj_b', 'in_proj_a', 'out_proj', 'conv1d'): + sub = getattr(la, name, None) + if sub is None or not hasattr(sub, 'weight'): + continue + g = sub.weight.grad + key = f'l{i}.{name}.weight' + grads[key] = g.detach().float().cpu() if g is not None else None + return out.detach().float().cpu(), loss.detach().float().cpu().item(), grads + + +def _max_diff(a: torch.Tensor, b: torch.Tensor): + diff = (a.float() - b.float()).abs() + denom = b.float().abs() + 1e-9 + rel = (diff / denom) + # Filter out near-zero denominators to avoid the rel-diff blow-up. + mask = b.float().abs() > 1e-3 + rel_masked = rel[mask] if mask.any() else rel + return { + 'max_abs': diff.max().item(), + 'mean_abs': diff.mean().item(), + 'max_rel': float(rel_masked.max()) if rel_masked.numel() else 0.0, + 'mean_rel': float(rel_masked.mean()) if rel_masked.numel() else 0.0, + } + + +def _assert_mean_close_report(name: str, a, b, mean_atol, norm_rtol=None): + if a is None or b is None: + raise AssertionError(f'{name}: one side is None (a={a is None}, b={b is None})') + stats = _max_diff(a, b) + ok = stats['mean_abs'] < mean_atol + if norm_rtol is not None and a.numel() > 0: + nrel = abs(a.float().norm() - b.float().norm()) / (b.float().norm() + 1e-9) + stats['norm_rel'] = nrel.item() + # Per-tensor norm_rel is noisy on small-magnitude tensors in MoE; treat + # as diagnostic. The aggregate grad-norm check below is the real + # scale guard. + flag = 'PASS' if ok else 'FAIL' + extra = f' norm_rel={stats.get("norm_rel", 0):.4e}' if 'norm_rel' in stats else '' + print(f' [{flag}] {name:32s} mean_abs={stats["mean_abs"]:.4e} max_abs={stats["max_abs"]:.4e}' + f' max_rel={stats["max_rel"]:.4e}{extra} (mean_atol={mean_atol:.0e})') + if not ok: + raise AssertionError( + f'{name} mismatch: mean_abs={stats["mean_abs"]:.4e} max_abs={stats["max_abs"]:.4e} ' + f'max_rel={stats["max_rel"]:.4e} (mean_atol={mean_atol})') + return stats + + +def _assert_close_report(name: str, a, b, atol, rtol, required: bool = True): + if a is None or b is None: + if required: + raise AssertionError(f'{name}: one side is None (a={a is None}, b={b is None})') + return None + stats = _max_diff(a, b) + ok = bool(torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol)) + flag = 'PASS' if ok else 'FAIL' + print(f' [{flag}] {name:32s} max_abs={stats["max_abs"]:.4e} mean_abs={stats["mean_abs"]:.4e}' + f' max_rel={stats["max_rel"]:.4e} (atol={atol:.0e}, rtol={rtol:.0e})') + if not ok and required: + raise AssertionError( + f'{name} mismatch: max_abs={stats["max_abs"]:.4e} mean_abs={stats["mean_abs"]:.4e} ' + f'max_rel={stats["max_rel"]:.4e} (atol={atol}, rtol={rtol})') + return stats + + +def _compare_models(model_kind: str, layer_types, seed: int, seq_len: int): + device = _device() + base = _build_tiny(model_kind, layer_types, device, seed) + ids, labels = _build_batch(device, seq_len) + + # FLA OFF (ground truth): torch reference path + os.environ['TWINKLE_NPU_FLA'] = '0' + m_off = copy.deepcopy(base).to(device) + _force_fla_off(m_off) + from transformers.models.qwen3_5.modeling_qwen3_5 import torch_chunk_gated_delta_rule as _torch_ref + with torch.no_grad(): + for layer in m_off.model.layers: + la = getattr(layer, 'linear_attn', None) + if la is not None: + assert la.causal_conv1d_fn is None, 'FLA OFF path must have causal_conv1d_fn=None' + assert la.chunk_gated_delta_rule is _torch_ref, \ + f'FLA OFF path must use torch reference, got {la.chunk_gated_delta_rule}' + logits_off, loss_off, grads_off = _run_forward_backward(m_off, ids, labels) + + # FLA ON: MindSpeed Triton kernel via twinkle patch + os.environ['TWINKLE_NPU_FLA'] = '1' + m_on = copy.deepcopy(base).to(device) + _force_fla_off(m_on) # start clean so apply_qwen3_5_fla re-patches + n_patched = _force_fla_on(m_on) + assert n_patched > 0, 'apply_qwen3_5_fla patched 0 instances; FLA ON path not active' + logits_on, loss_on, grads_on = _run_forward_backward(m_on, ids, labels) + + print(f'\n=== {model_kind} {layer_types} seed={seed} seq_len={seq_len} (patched={n_patched}) ===') + print(f' loss: off={loss_off:.6f} on={loss_on:.6f} diff={abs(loss_on - loss_off):.4e}') + stats = {'loss': {'off': loss_off, 'on': loss_on, 'abs_diff': abs(loss_on - loss_off)}} + _assert_mean_close_report('logits', logits_on, logits_off, LOGITS_MEAN_ATOL) + assert abs(loss_on - loss_off) < LOSS_ATOL, f'loss diff {abs(loss_on - loss_off)} >= {LOSS_ATOL}' + grad_stats = {} + total_norm_on = torch.tensor(0.0) + total_norm_off = torch.tensor(0.0) + for key in grads_off: + s = _assert_mean_close_report(key, grads_on.get(key), grads_off.get(key), + GRAD_MEAN_ATOL, GRAD_NORM_RTOL) + if s is not None: + grad_stats[key] = s + if grads_on.get(key) is not None: + total_norm_on += grads_on[key].float().pow(2).sum() + if grads_off.get(key) is not None: + total_norm_off += grads_off[key].float().pow(2).sum() + # Aggregate grad-norm check (robust to per-tensor MoE routing noise). + total_norm_on = total_norm_on.sqrt() + total_norm_off = total_norm_off.sqrt() + total_rel = abs(total_norm_on - total_norm_off) / (total_norm_off + 1e-9) + ok = total_rel < GRAD_NORM_RTOL + print(f' [{"PASS" if ok else "FAIL"}] total_grad_norm_rel ' + f'on={total_norm_on.item():.4e} off={total_norm_off.item():.4e} ' + f'rel={total_rel.item():.4e} (rtol={GRAD_NORM_RTOL})') + if not ok: + raise AssertionError( + f'aggregate grad-norm rel diff {total_rel.item():.4e} >= {GRAD_NORM_RTOL}') + stats['grads'] = grad_stats + return stats + + +# --------------------------------------------------------------------------- +# Test cases +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not _HAS_NPU, reason='requires Ascend NPU (torch_npu)') +@pytest.mark.skipif(not _HAS_QWEN35, reason='requires transformers Qwen3.5') +class TestQwen35FlaBwdPrecision: + """Compare FLA ON (bf16 Triton kernel) vs FLA OFF (fp32 torch reference). + + These tests do NOT require multiple devices and are the precise analogue of + the FLA toggle used by ``cookbook/transformers/fsdp2.sh`` (which sets the + same ``TWINKLE_NPU_FLA`` env var via ``npu_builtin(model)``). + """ + + @pytest.mark.parametrize('seed', [1234, 2024]) + @pytest.mark.parametrize('seq_len', [64, 256]) + def test_dense_linear_attention_fla_on_vs_off(self, seed, seq_len): + _compare_models('dense', ['linear_attention', 'linear_attention'], seed, seq_len) + + @pytest.mark.parametrize('seed', [1234]) + def test_dense_mixed_attention_fla_on_vs_off(self, seed): + _compare_models('dense', ['full_attention', 'linear_attention'], seed, 128) + + @pytest.mark.parametrize('seed', [1234, 2024]) + def test_moe_linear_attention_fla_on_vs_off(self, seed): + _compare_models('moe', ['linear_attention', 'linear_attention'], seed, 128) + + +if __name__ == '__main__': + _compare_models('dense', ['linear_attention', 'linear_attention'], 1234, 64) + _compare_models('dense', ['linear_attention', 'linear_attention'], 2024, 256) + _compare_models('moe', ['linear_attention', 'linear_attention'], 1234, 128) + print('\nAll comparisons finished.')