mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Yulong CaoandVaibhav Jindal
8343b3cf7b
## Summary Adds Liger Kernel support for DeepSeek-V2 (`deepseek_v2` model type). Closes #129 (on the maintainer 2024 Q4 roadmap) and supersedes the stale PR #499, which predates the `deepseek_v2` module in transformers. - `model/deepseek_v2.py`: new `lce_forward`. In training with labels, `skip_logits` defaults to `True`, so logits are never materialized (FLCE path via `LigerForCausalLMLoss`); otherwise logits are materialized and the HF-native `self.loss_function` is used. Returns `LigerCausalLMOutputWithPast` with `token_accuracy` / `predicted_tokens`. - `monkey_patch.py`: `apply_liger_kernel_to_deepseek_v2` with class-level and instance-level patching, registered in `MODEL_TYPE_TO_APPLY_LIGER_FN` (`AutoLigerKernelForCausalLM` picks it up automatically; previously `deepseek_v2` models silently got no Liger kernels). - RMSNorm: `DeepseekV2RMSNorm` → `LigerRMSNorm`, including the MLA attention `q_a_layernorm` / `kv_a_layernorm`. - SwiGLU: dense MLPs and shared experts (`DeepseekV2MLP`) → `LigerQwen3MoeSwiGLUMLP` on all supported versions. Routed experts: batched `DeepseekV2Experts` → fused `LigerExperts` on transformers v5+; per-expert patch on transformers v4 (the `deepseek_v2` module landed in v4.54). - RoPE: unsupported — DeepSeek-V2 uses interleaved partial RoPE, incompatible with `liger_rotary_pos_emb`. `rope=True` emits a `warning_once` and skips (no silent behavior change); RoPE is intentionally not listed in the README table. - `mlp_bias`: `LigerQwen3MoeSwiGLUMLP` is bias-free, matching the `DeepseekV2Config` default (`mlp_bias=False`); `mlp_bias=True` checkpoints are not supported and noted in the docstring. - Tests: instance monkey-patch test (asserts both v5 and v4 branch outcomes), 4 convergence tests (bf16/fp32 × FLCE/logits) driven by a mini `deepseek_v2` config, `revert_liger_kernel_to_deepseek_v2` util, README entry. ## Details - transformers v5 renames the V2 MoE block `DeepseekV2MoE` → `DeepseekV2Moe` (V3 kept `DeepseekV3MoE`), so the apply function resolves the class with a try/except import. - Shared-code change: `LigerExperts`' expert-count lookup previously handled `num_experts` (qwen3_moe/qwen3_next) and `num_local_experts` (mixtral/deepseek_v3). `DeepseekV2Config` on early transformers v5 only exposes `n_routed_experts` (the `num_experts` alias was added later), so a third `elif` branch was added. Existing branches are untouched; regression coverage for both neighboring branches ran in the same validation (see below). - The transformers v4 branch (per-expert patching) mirrors the existing mixtral / qwen3_moe / deepseek_v3 patterns; it is exercised by CI's v4 matrix where the `deepseek_v2` module is available (≥ 4.54) and skipped via `is_deepseek_v2_available()` otherwise. - On transformers v5+, `LigerExperts`' batched expert weights (`(N, 2*I, H)` / `(N, H, I)`) match the HF `DeepseekV2Experts` layout, so pretrained checkpoints load unchanged. ## Testing Done | Check | Status | |---|---| | Instance patching (CPU, transformers 5.16.1) | PASS | | Convergence bf16 FLCE (Kaggle T4) | PASS | | Convergence bf16 with_logits (Kaggle T4) | PASS | | Convergence fp32 FLCE (Kaggle T4) | PASS | | Convergence fp32 with_logits (Kaggle T4) | PASS | | Regression: mini_deepseek_v3 + mini_qwen3_moe (Kaggle T4) | PASS (pre-rebase base; see Post-rebase re-verification below) | | Instance patching deepseek_v3 (Kaggle T4) | PASS | | deepseek_v4 sibling tests | SKIP (module not in transformers 5.0.0) | | Checkstyle (ruff) | PASS | Exact commands run on this change: - `pytest test/transformers/test_monkey_patch.py -k deepseek` → 3 passed (transformers 5.16.1, CPU) - `pytest test/transformers/test_monkey_patch.py test/convergence/{bf16,fp32}/test_mini_models.py test/convergence/{bf16,fp32}/test_mini_models_with_logits.py -k 'deepseek or qwen3_moe'` on Kaggle T4 (torch 2.10.0+cu128, transformers 5.0.0) → **13 passed, 7 skipped, 0 failed** ([kernel log](https://www.kaggle.com/code/cauliyulong/liger-dsv2-129-test)) ### Post-rebase re-verification The branch was rebased onto `main` at `6ad077c` (it was 17 commits behind). Everything below was re-run on the rebased commit: - `pytest test/transformers/test_monkey_patch.py -k deepseek` → **3 passed** (CPU, transformers 5.16.1) - `pytest test/convergence/{bf16,fp32}/test_mini_models{,_with_logits}.py -k deepseek_v2` → **4 passed** (RTX 4060, torch 2.11.0+cu128) - `... -k deepseek_v3` → **4 passed** - `... -k qwen3_moe` → 3 passed, 1 failed: `bf16/test_mini_models_with_logits.py::test_mini_model[mini_qwen3_moe-...]`, `[Loss]` mismatch at index `(0, 27)` (`5.145391464` vs `4.135959148`). **Not caused by this PR**: the identical mismatch (same index, same values) reproduces on a clean `origin/main` checkout at `6ad077c` in the same local environment (transformers 5.16.1 / torch 2.11.0), and this PR does not touch the qwen3_moe path — `LigerExperts` takes the `hasattr(config, "num_experts")` branch for qwen3_moe both before and after the change. It does not reproduce on the Kaggle T4 run above (transformers 5.0.0), so it looks environment-specific rather than a regression from this change. - `ruff check` on all changed files → all checks passed - Hardware Type: NVIDIA T4 (Kaggle, sm_75) + RTX 4060 (sm_89); CPU for the instance-patching tests - [x] run `make checkstyle` to ensure code style (`ruff check` clean on all changed files) - [ ] run `make test` to ensure correctness - [ ] run `make test-convergence` to ensure convergence Note on the two unchecked boxes: `make test` / `make test-convergence` were not run end-to-end on this change. The targeted subsets listed above are what was exercised; I can run the full targets on request. Happy to add a memory/throughput table on A100 for `DeepSeek-V2-Lite` dims if reviewers find it useful. Co-authored-by: Vaibhav Jindal <[email protected]>