feat: optional Liger fused RMSNorm/SwiGLU kernels for DFlash training

[jessiewei7]

Summary

Opt-in --use-liger-kernel for scripts/train_dflash.py: patches transformers' Qwen3 RMSNorm/MLP (SwiGLU) with Liger Kernel's fused Triton kernels (the DFlash draft model is Qwen3-based). Measured 3.38 → 3.0 s/step (~11%) on a production DFlash run (GLM target, 5-layer draft, 2×8 GPUs, FSDP, batch size 16) with the flag as the only variable. Off by default; liger-kernel is an optional extra.

Change

• scripts/train_dflash.py: the flag; when set, calls apply_liger_kernel_to_qwen3(rope=False, rms_norm=True, swiglu=True, cross_entropy=False, fused_linear_cross_entropy=False) before model construction. rope is excluded because DFlash applies its own asymmetric rotary embedding via a local apply_rotary_pos_emb the patch cannot reach; the CE kernels because the loss is computed outside the HF forward.
• specforge/modeling/draft/dflash.py: resolve Qwen3RMSNorm/Qwen3MLP through the modeling_qwen3 module at construction time — Liger rebinds the module attributes, so from-imports captured at import time miss the patch depending on import order.
• pyproject.toml: liger optional extra.

Numerics

Liger's kernels are exact, not approximations: same weights and inputs give bit-identical bf16 (training dtype) forward outputs with and without the flag; fp32 gradients match at float-reassociation level (≤2e-6). Verified by building the same draft decoder layer before/after the patch with identical state dicts. Loss curves in the production A/B overlap.

[gemini-code-assist]

Code Review

This pull request adds support for Liger fused Triton kernels to optimize Qwen3 RMSNorm and SwiGLU modules, introducing a new --use-liger-kernel option and updating module resolution to ensure patches apply correctly. The review feedback recommends wrapping the optional liger_kernel import in a try-except block to gracefully handle cases where the dependency is missing or outdated.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Since liger-kernel is an optional dependency (defined as an extra in pyproject.toml), attempting to import it directly when --use-liger-kernel is specified will raise a raw ModuleNotFoundError if the package is not installed. Additionally, if an older version of liger-kernel is installed that does not yet support Qwen3, it will raise an ImportError.

It is highly recommended to wrap this import in a try-except block to provide a clean, user-friendly error message guiding the user on how to install or upgrade the package.

    if args.use_liger_kernel:
        try:
            from liger_kernel.transformers import apply_liger_kernel_to_qwen3
        except ImportError as e:
            raise ImportError(
                "Failed to import 'apply_liger_kernel_to_qwen3' from 'liger_kernel'. "
                "Please ensure you have a compatible version of 'liger-kernel' installed "
                "(e.g., by running 'pip install -U liger-kernel' or 'pip install .[liger]')."
            ) from e

[jessiewei7]

Done in 4c3e929.