Skip to content

fix: restore cross-entropy state in revert_liger_kernel helpers - #1389

Open
atirna wants to merge 2 commits into
linkedin:mainfrom
atirna:fix/revert-liger-cross-entropy-542
Open

fix: restore cross-entropy state in revert_liger_kernel helpers#1389
atirna wants to merge 2 commits into
linkedin:mainfrom
atirna:fix/revert-liger-cross-entropy-542

Conversation

@atirna

@atirna atirna commented Aug 16, 2026

Copy link
Copy Markdown

Summary

revert_liger_kernel_to_* now clears Liger cross-entropy monkey patches so convergence tests don't compare two Liger-patched baselines after the first case.

Fixes #542

Details

  • add revert_liger_cross_entropy_patches() to restore torch.nn.functional.cross_entropy, reload torch.nn (covers nn.CrossEntropyLoss swaps), and drop module-level CrossEntropyLoss bindings when they point at LigerCrossEntropyLoss
  • call it from every revert_liger_kernel_to_* helper in test/utils.py
  • add test_revert_liger_cross_entropy_patches for apply → revert → compare on both patch paths

Testing Done

  • Hardware Type: CPU (macOS arm64)
  • ruff check / ruff format on monkey_patch.py, test/utils.py, test/transformers/test_monkey_patch.py
  • independent review PASS on patch 91b6aba (glm-5.3)
  • make test (needs CUDA/triton; leaving to CI)
  • make test-convergence (needs CUDA; leaving to CI)

Made with Cursor

atirna and others added 2 commits August 17, 2026 00:13
Add revert_liger_cross_entropy_patches() to undo global nn.functional.cross_entropy
and module-level CrossEntropyLoss swaps after convergence-test reverts.

Fixes linkedin#542

Co-authored-by: Cursor <[email protected]>
Assert nn.CrossEntropyLoss restoration and use the saved torch baseline
so the test stays valid even if earlier cases left F.cross_entropy patched.

Co-authored-by: Cursor <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

revert_liger_kernel_to_xxx can't revert LigerCrossEntropyLoss for transformers>=4.46.1

1 participant