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>
This commit is contained in:
Ayushman
GitHub
parent
01f7e14988
commit
c356427f30
3 changed files with 152 additions and 7 deletions
|
|
@ -12,6 +12,7 @@ from __future__ import annotations
|
|||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -71,12 +72,58 @@ def test_default_spec_matches_table(monkeypatch):
|
|||
assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC
|
||||
|
||||
|
||||
def test_skips_torchao_on_windows_rocm():
|
||||
@pytest.mark.parametrize(
|
||||
("rocm_windows_torch_installed", "installed_torch_is_windows_rocm"),
|
||||
[
|
||||
(True, False),
|
||||
(False, True),
|
||||
],
|
||||
)
|
||||
def test_skips_torchao_on_windows_rocm(
|
||||
monkeypatch, tmp_path, rocm_windows_torch_installed, installed_torch_is_windows_rocm
|
||||
):
|
||||
"""The overrides step must skip torchao on Windows ROCm: no working build exists
|
||||
there (it imports an absent c10d backend and crashes transformers.quantizers),
|
||||
so the installer skips it and relies on the runtime stub instead."""
|
||||
source = _INSTALL_SCRIPT.read_text(encoding = "utf-8")
|
||||
# Branches on the Windows-ROCm marker set by _ensure_rocm_torch ...
|
||||
assert "elif _rocm_windows_torch_installed:" in source
|
||||
# ... and reports the skip in the progress label.
|
||||
assert "dependency overrides (skipped, Windows ROCm)" in source
|
||||
mod = _load_module(monkeypatch)
|
||||
installed_specs: list[str] = []
|
||||
progress_labels: list[str] = []
|
||||
|
||||
def _record_pip_install(*args, **kwargs):
|
||||
installed_specs.extend(str(arg) for arg in args)
|
||||
return 0
|
||||
|
||||
unstructured_plugin = tmp_path / "unstructured"
|
||||
github_plugin = tmp_path / "github"
|
||||
unstructured_plugin.mkdir()
|
||||
github_plugin.mkdir()
|
||||
|
||||
subprocess_result = MagicMock()
|
||||
subprocess_result.returncode = 0
|
||||
subprocess_result.stdout = ""
|
||||
|
||||
monkeypatch.setenv("SKIP_STUDIO_BASE", "1")
|
||||
monkeypatch.setattr(mod, "IS_WINDOWS", True)
|
||||
monkeypatch.setattr(mod, "IS_MACOS", False)
|
||||
monkeypatch.setattr(mod, "IS_MAC_ARM", False)
|
||||
monkeypatch.setattr(mod, "NO_TORCH", False)
|
||||
monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed)
|
||||
monkeypatch.setattr(
|
||||
mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm
|
||||
)
|
||||
monkeypatch.setattr(mod, "_bootstrap_uv", lambda: False)
|
||||
monkeypatch.setattr(mod, "_repair_bad_anyio", lambda: None)
|
||||
monkeypatch.setattr(mod, "_ensure_rocm_torch", lambda: None)
|
||||
monkeypatch.setattr(mod, "_ensure_cuda_torch", lambda: None)
|
||||
monkeypatch.setattr(mod, "_has_usable_nvidia_gpu", lambda: True)
|
||||
monkeypatch.setattr(mod, "run", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(mod, "pip_install", _record_pip_install)
|
||||
monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label))
|
||||
monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin)
|
||||
monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin)
|
||||
monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result)
|
||||
|
||||
assert mod.install_python_stack() == 0
|
||||
|
||||
assert not any(spec.startswith("torchao") for spec in installed_specs)
|
||||
assert "dependency overrides (skipped, Windows ROCm)" in progress_labels
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -2296,6 +2296,70 @@ class TestRocmTorchInstalledEnvVar:
|
|||
mock_bnb.assert_not_called()
|
||||
|
||||
|
||||
class TestWindowsRocmTorchaoGuard:
|
||||
"""Verify the torchao skip can detect an installed Windows ROCm torch build."""
|
||||
|
||||
def test_installed_torch_is_windows_rocm_accepts_rocm_probe(self):
|
||||
rv = MagicMock()
|
||||
rv.returncode = 0
|
||||
rv.stdout = "yes"
|
||||
with (
|
||||
patch.object(stack_mod, "IS_WINDOWS", True),
|
||||
patch.object(stack_mod.subprocess, "run", return_value = rv),
|
||||
):
|
||||
assert stack_mod._installed_torch_is_windows_rocm() is True
|
||||
|
||||
def test_installed_torch_is_windows_rocm_rejects_non_rocm_probe(self):
|
||||
rv = MagicMock()
|
||||
rv.returncode = 0
|
||||
rv.stdout = ""
|
||||
with (
|
||||
patch.object(stack_mod, "IS_WINDOWS", True),
|
||||
patch.object(stack_mod.subprocess, "run", return_value = rv),
|
||||
):
|
||||
assert stack_mod._installed_torch_is_windows_rocm() is False
|
||||
|
||||
def test_installed_torch_is_windows_rocm_is_non_windows_noop(self):
|
||||
with patch.object(stack_mod, "IS_WINDOWS", False):
|
||||
assert stack_mod._installed_torch_is_windows_rocm() is False
|
||||
|
||||
@patch.object(stack_mod, "_repair_bad_anyio")
|
||||
@patch.object(stack_mod, "_ensure_rocm_torch")
|
||||
@patch.object(stack_mod, "_ensure_cuda_torch")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "run")
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_install_python_stack_skips_torchao_when_windows_rocm_torch_is_installed(
|
||||
self, mock_pip, mock_run, mock_has_nvidia, mock_cuda, mock_rocm, mock_anyio, tmp_path
|
||||
):
|
||||
unstructured_plugin = tmp_path / "unstructured"
|
||||
github_plugin = tmp_path / "github"
|
||||
unstructured_plugin.mkdir()
|
||||
github_plugin.mkdir()
|
||||
|
||||
subprocess_result = MagicMock()
|
||||
subprocess_result.returncode = 0
|
||||
subprocess_result.stdout = ""
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}),
|
||||
patch.object(stack_mod, "IS_WINDOWS", True),
|
||||
patch.object(stack_mod, "IS_MACOS", False),
|
||||
patch.object(stack_mod, "IS_MAC_ARM", False),
|
||||
patch.object(stack_mod, "NO_TORCH", False),
|
||||
patch.object(stack_mod, "_rocm_windows_torch_installed", False),
|
||||
patch.object(stack_mod, "_bootstrap_uv", return_value = False),
|
||||
patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = True),
|
||||
patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin),
|
||||
patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin),
|
||||
patch.object(stack_mod.subprocess, "run", return_value = subprocess_result),
|
||||
):
|
||||
assert stack_mod.install_python_stack() == 0
|
||||
|
||||
installed_specs = [str(arg) for call in mock_pip.call_args_list for arg in call.args]
|
||||
assert not any("torchao" in arg for arg in installed_specs)
|
||||
|
||||
|
||||
# TEST: worker.py -- Windows ROCm patches (source-level checks)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue