Skip to content

fix(utils): restore torch cross-entropy state in revert helpers (#1587) - #1611

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

YuvrajSinghBhadoria2 wants to merge 2 commits into
linkedin:mainfrom
YuvrajSinghBhadoria2:fix/cross-entropy-revert

Conversation

@YuvrajSinghBhadoria2

Copy link
Copy Markdown

Fixes #1587

In transformers >= 4.46.1, CrossEntropyLoss is patched into module-level
references and torch.nn.functional.cross_entropy. The revert helper
functions previously only reloaded modeling modules, leaving
F.cross_entropy and nn.CrossEntropyLoss pointing to LigerCrossEntropyLoss.
Consequently, multi-model convergence tests compared two Liger-patched
models against each other after the first test case.

  • Add _ORIGINAL_TORCH_CROSS_ENTROPY capturing the original F.cross_entropy
    before any Liger patching
  • Add revert_liger_cross_entropy_patches() to restore F.cross_entropy,
    reload torch.nn, and remove any lingering CrossEntropyLoss attribute
    on modeling modules bound to LigerCrossEntropyLoss
  • Call revert_liger_cross_entropy_patches() from all revert_liger_kernel_to_*
    helpers in test/utils.py
  • Add test_revert_liger_cross_entropy_patches unit test

All tests verified passing in a CPU-only Docker container (no GPU required).

)

Skip the torch.cuda.device()/torch.xpu.device() context-manager guard
when torch.is_compiling() is true. Dynamo emits incorrect device-restore
nodes during tracing, which leads to wrong output or illegal memory
access (linkedin#1590).

Use getattr(torch, 'is_compiling', lambda: False)() so the check is
robust on torch versions that do not define torch.is_compiling.

Add CPU-runnable unit tests covering the compile-skip path, the
non-compile path, CPU device, and backends without a .device attribute.
…edin#1587)

In transformers >= 4.46.1, CrossEntropyLoss is patched into module-level
references and torch.nn.functional.cross_entropy. The revert helper
functions previously only reloaded modeling modules, leaving
F.cross_entropy and nn.CrossEntropyLoss pointing to LigerCrossEntropyLoss.
Consequently, multi-model convergence tests compared two Liger-patched
models against each other after the first test case.

- Add _ORIGINAL_TORCH_CROSS_ENTROPY capturing the original F.cross_entropy
  before any Liger patching
- Add revert_liger_cross_entropy_patches() to restore F.cross_entropy,
  reload torch.nn, and remove any lingering CrossEntropyLoss attribute
  on modeling modules bound to LigerCrossEntropyLoss
- Call revert_liger_cross_entropy_patches() from all revert_liger_kernel_to_*
  helpers in test/utils.py
- Add test_revert_liger_cross_entropy_patches unit test

Fixes linkedin#1587

This branch has not been deployed

No deployments
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.

1 participant