tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

fix: accidentally reassigning tensor to weight

Browse Source
This commit is contained in:
NanoCode012
2025-04-09 13:45:29 +07:00
parent 1a85fab2ca
commit 4581d6a8de
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama4.py
View File

@@ -310,13 +310,14 @@ def cce_forward_multimodal(
if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
assert labels is not None
# reset lm head if not already done. linear model has some lm_head weight issue
# reset lm head gradient on first pass.
# linear model has some lm_head weight issue
# see https://github.com/axolotl-ai-cloud/axolotl/pull/2505
global RESET_LM_HEAD # pylint: disable=global-statement
if RESET_LM_HEAD:
RESET_LM_HEAD = False
self.language_model.lm_head.weight = (
self.language_model.lm_head.weight.detach().requires_grad_(True)
)
self.language_model.lm_head.weight.requires_grad_(False) # Detach
self.language_model.lm_head.weight.requires_grad_(True) # Reattach
loss = apply_lce(
hidden_states,