ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Walk DDP/FSDP .module when scanning for the pre-train marker

Browse source 
The chain walk followed only .model and .base_model, so a probe that fired on the
model below a DDP/FSDP wrapper (which exposes it via .module) left the marker
undetected and the poisoned compile cache un-reset. Add .module to the walk.
This commit is contained in:
danielhanchen 2026-06-22 02:07:57 +00:00
commit 76d7932503
1 changed files with 5 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

5
unsloth/models/rl.py
View file

@ -417,9 +417,14 @@ def _unsloth_reset_stray_compile_cache(self):
markers.append(_m)
if _m.get("seen"):
seen = True
# Follow the wrapper chain: Unsloth/HF (.model), PEFT (.base_model) and
# DDP / FSDP (.module). A pre-train probe can fire on the model below a
# DDP wrapper, so .module must be walked too or the marker is missed.
_nxt = getattr(_curr, "model", None)
if _nxt is None:
_nxt = getattr(_curr, "base_model", None)
if _nxt is None:
_nxt = getattr(_curr, "module", None)
_curr = _nxt
if seen and os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1":
try: