Guard Windows ROCm torchao override skip (#6837)

* Fix: skip fp16/bf16 validation for full finetuning in RL trainers

When doing full finetuning (FFT) of a bfloat16 model, the fp16/bf16
mismatch validation fires before the corrective logic runs, causing a
misleading error even though the code would properly handle it downstream.
Skip the validation when full_finetuning is active.

Fixes #6731

* Fix: auto-correct fp16/bf16 mismatches for full finetuning before validation

Instead of entirely skipping validation (which could let mismatches
through when mixed_precision_dtype is float32), auto-correct explicit
fp16/bf16 settings that conflict with the model's dtype for FFT. This
way the existing validation still catches real mismatches for non-FFT
cases, and the corrective logic below handles the normalized settings.

Fixes the issue raised in Codex review of PR #6813.

* Guard Windows ROCm torchao override skip

Detect installed ROCm torch directly before applying the torchao override so Windows ROCm environments never install the crashing torchao package even if the earlier ROCm-installed flag is missing.

* Update unsloth/models/rl.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update studio/install_python_stack.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Harden ROCm probe and sync RL precision flags

Tolerate stray stdout noise when probing Windows ROCm torch installs by checking the last non-empty output line, matching the existing torch version probe behavior. Also keep args.fp16 and args.bf16 synchronized with the full-finetuning precision auto-corrections in the RL trainer patch so downstream eval settings see a consistent TrainingArguments state.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add MLX trainer compatibility shims

Patch imported MLXTrainer and MLXTrainingConfig objects to preserve the expected dataclass field ordering and to provide a _train_dataset_for_batches fallback when older trainers or test doubles only expose train_dataset. Also add focused worker tests covering both compatibility paths.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Scope PR to Windows ROCm torchao guard

* Restore PR scope to Windows ROCm guard

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: cover Windows ROCm torchao skip behavior

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Ayushman Paul <ayushman@HP>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>

studio/install_python_stack.py

@@ -181,6 +181,40 @@ def _probe_installed_torch_version() -> str | None:
     return lines[-1] if lines else None
 
 
+def _installed_torch_is_windows_rocm() -> bool:
+    """Return True when the target venv currently has a Windows ROCm torch build.
+
+    This is a belt-and-suspenders guard for the torchao override step: if the
+    earlier ROCm install path failed to set _rocm_windows_torch_installed but the
+    venv already contains a ROCm torch wheel, still skip torchao because it
+    crashes on import on Windows ROCm.
+    """
+    if not IS_WINDOWS:
+        return False
+    try:
+        probe = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "import sys, torch; "
+                    "hip = getattr(getattr(torch, 'version', None), 'hip', None) or ''; "
+                    "ver = getattr(torch, '__version__', '').lower(); "
+                    "sys.stdout.write('yes' if (hip or 'rocm' in ver or 'rocmsdk' in ver) else '')"
+                ),
+            ],
+            stdout = subprocess.PIPE,
+            stderr = subprocess.DEVNULL,
+            text = True,
+            timeout = 90,
+            **_windows_hidden_subprocess_kwargs(),
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return False
+    lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()]
+    return probe.returncode == 0 and bool(lines and lines[-1] == "yes")
+
+
 # constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install
 # from before the cap existed can already be stuck at 4.14+, which later
 # constrained installs won't touch since it already satisfies mcp/fastmcp.
@@ -2256,7 +2290,7 @@ def install_python_stack() -> int:
     #    (no working build; see below).
     if NO_TORCH:
         _progress("dependency overrides (skipped, no torch)")
-    elif _rocm_windows_torch_installed:
+    elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm():
         # No working Windows ROCm torchao build: it imports an absent c10d backend
         # and crashes transformers.quantizers. Studio stubs it at runtime, so
         # installing it only ships a package that crashes on import -- skip it.