From 76d79325034c7489c180e1b75ff8e36bb192bc13 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 02:07:57 +0000 Subject: [PATCH] Walk DDP/FSDP .module when scanning for the pre-train marker 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. --- unsloth/models/rl.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index cf64ee4e2c..f399a7b3e8 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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: