mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
LazuriteandVaibhav Jindal
5b0f210cc4
## Summary Fixes #1485. `apply_liger_kernel_to_megatron(cross_entropy=True)` replaces the CE functions only on their defining modules, but Megatron calls copies bound at import time. When Megatron model code is imported before the patch (the usual order, also used by `examples/megatron/run_mode1_monkey_patch.py`), the loss silently stays on Megatron's kernel: - Released megatron-core (0.18 / 0.19): `language_module` binds `fused_vocab_parallel_cross_entropy` with `from ... import` ([L17](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/models/common/language_module/language_module.py#L17)). The unfused branch calls the `megatron.core.tensor_parallel` re-export ([L183](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/models/common/language_module/language_module.py#L183)), which is bound before the `cross_entropy` submodule can be patched. Neither path reaches Liger. - Megatron-LM main (`NVIDIA/Megatron-LM#6766`): `backends` binds the fused function as `_fused_ce` ([L19-L21](https://github.com/NVIDIA/Megatron-LM/blob/ec806b20a0a976924c1668486c3554ffb90c000e/megatron/core/models/backends.py#L19-L21)). The fused path misses Liger; the unfused path already works. This PR rebinds those copies too, but only where they still hold the original function. No version check is needed, and a call site that another framework has already overridden is left alone. If no fused call site ends up on Liger, the patch logs a warning. The docs now say to call the patch before building the model, matching the docstring: Megatron-LM main picks the CE function when the model is constructed. Once the minimum supported megatron-core includes the `BackendSpecProvider.vocab_parallel_cross_entropy` slot, this can switch to overriding that slot on the providers and drop the rebinding for released versions. ## Testing Done - Hardware Type: 1x NVIDIA GPU - [ ] run `make test` to ensure correctness: ran `test/megatron` (173 passed, 48 skipped multi-GPU tests) - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence: N/A, no kernel change The new tests in `test_monkey_patch.py` compute the loss through stub call sites laid out like each Megatron version; 4 of 5 fail on current `main`. A real 1-layer `GPTModel` (TP=1) was run for one forward and backward pass, with model code imported before the patch. Calls into Liger CE (fused / unfused): | | before | after | |---|---|---| | megatron-core 0.18.2 | 0 / 0 | 1 / 1 | | Megatron-LM main | 0 / 1 | 1 / 1 | When the patch runs first, both cases also give 1 / 1 after this change (before: 1 / 0 on 0.18.2). Co-authored-by: Vaibhav Jindal <[email protected]>