mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Yulong Cao
40a9d8a61b
## Summary Adds Liger Kernel support for DeepSeek-V3 (`deepseek_v3` model type), closing the gap left open by #623. - `model/deepseek_v3.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_v3` with class-level and instance-level patching, registered in `MODEL_TYPE_TO_APPLY_LIGER_FN` (`AutoLigerKernelForCausalLM` picks it up automatically). - RMSNorm: `DeepseekV3RMSNorm` → `LigerRMSNorm`, including attention `q_a_layernorm` / `kv_a_layernorm`. - SwiGLU: dense MLPs and shared experts (`DeepseekV3MLP`) → `LigerQwen3MoeSwiGLUMLP` on all supported versions. Routed experts: batched `DeepseekV3Experts` → fused `LigerExperts` on transformers v5+; per-expert patch on transformers v4. - RoPE: unsupported — DeepSeek-V3 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. - Tests: instance monkey-patch test (asserts both v5 and v4 branch outcomes), 4 convergence tests (bf16/fp32 × FLCE/logits) driven by a mini `deepseek_v3` config, `revert_liger_kernel_to_deepseek_v3` util, README entry. ## Details - The transformers v4 branch (per-expert patching) mirrors the existing mixtral / qwen3_moe patterns but could not be exercised locally (only transformers 5.16.1 is installed here); it is covered by the CI v4 matrix. - On transformers v5+, `LigerExperts` accepts `DeepseekV3Config` (`num_local_experts` aliases `n_routed_experts`), and its batched expert weights (`(N, 2*I, H)` / `(N, H, I)`) match the HF layout, so pretrained checkpoints load unchanged. ## Testing Done | Check | Status | |---|---| | Import | PASS | | Instance Patching | PASS | | Convergence bf16 FLCE | PASS | | Convergence bf16 with_logits | PASS | | Convergence fp32 FLCE | PASS | | Convergence fp32 with_logits | PASS | | Convergence multimodal | SKIP (text-only model) | | Checkstyle | PASS | Exact commands run on this change: - `pytest test/transformers/test_monkey_patch.py -k deepseek_v3` → 1 passed - `pytest test/convergence/{bf16,fp32}/test_mini_models.py test/convergence/{bf16,fp32}/test_mini_models_with_logits.py -k deepseek_v3` → 4 passed - `ruff check` + `ruff format --check` clean on all changed files ### Memory Single A100-SXM4-80GB, bf16, forward + backward, no optimizer step. Randomly initialised (no public checkpoint). Real dims: `vocab_size=129280`, `hidden_size=7168`, `num_attention_heads=128`, `intermediate_size=18432`, `moe_intermediate_size=2048` and the MLA projections. Slimmed and disclosed: 2 layers with `first_k_dense_replace=1` (real: 61 / 3), 16 routed experts with `num_experts_per_tok=2`, `n_group=2`, `topk_group=1`, `routed_scaling_factor=1.0` (real: 256 / 8, 8 / 4, 2.5). | batch | seq | peak mem HF | peak mem Liger | saved | |------:|----:|------------:|---------------:|------:| | 1 | 2048 | 13.1 GiB | 12.61 GiB | 4% | | 1 | 4096 | 16.24 GiB | 13.02 GiB | 20% | | 1 | 8192 | 26.18 GiB | 16.4 GiB | 37% | | 1 | 16384 | 46.06 GiB | 23.15 GiB | 50% | | 1 | 32768 | OOM | 36.64 GiB | | | 4 | 8192 | OOM | 36.67 GiB | | Savings grow with the token count: HF's loss pipeline materializes the full `[tokens, 129280]` logits tensor and its fp32 upcast/softmax (~2.4 GiB per additional 1K tokens), while FLCE replaces it with chunked transients plus a fixed 3.45 GiB fp32 weight-gradient accumulator — 4% at 2K, 50% at 16K, and at 32K tokens (batch 1 x 32K or 4 x 8K) HF no longer fits the 80 GB GPU while Liger holds ~36.7 GiB. ### Throughput Same hardware and dtype, batch 1, seq 4096, varying depth: | layers | lm_head share | tok/s HF | tok/s Liger | ratio | |-------:|--------------:|---------:|------------:|------:| | 2 | 27% | 8,457 | 8,172 | 0.97x | | 4 | 18% | 5,117 | 5,039 | 0.98x | | 8 | 10% | 2,871 | 2,862 | 1.00x | The ratio converges to 1.00x as depth dilutes the loss segment. The shipping 61-layer model puts `lm_head` at ~0.14% of parameters, so the end-to-end cost is ~1.00x. (On a 15 GB T4 in fp16 — sm_75 has no bf16 tensor cores — slimmer variants show the same pattern with the fixed accumulator dominating until ~4K tokens; happy to share those numbers too.) - Hardware Type: RTX 4060 Laptop GPU (8GB) — transformers 5.16.1, torch 2.11.0+cu128; benchmarks: NVIDIA A100-SXM4-80GB — transformers 5.8.1, torch 2.11.0+cu130 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Closes #623