diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 4f86f99ded..1010e56dac 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -34,13 +34,50 @@ from utils.hardware import apply_gpu_ids def _activate_transformers_version(model_name: str) -> None: - """Activate the correct transformers version BEFORE any ML imports.""" + """Activate the correct transformers version BEFORE any ML imports. + + Uses get_transformers_tier() to decide between .venv_t5_550/ (5.5.0), + .venv_t5_530/ (5.3.0), or the default 4.57.x. + """ + # Ensure backend is on path for utils imports + backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) - from utils.transformers_version import activate_transformers_for_subprocess + from utils.transformers_version import ( + get_transformers_tier, + _resolve_base_model, + _ensure_venv_t5_530_exists, + _ensure_venv_t5_550_exists, + _VENV_T5_530_DIR, + _VENV_T5_550_DIR, + ) - activate_transformers_for_subprocess(model_name) + resolved = _resolve_base_model(model_name) + tier = get_transformers_tier(resolved) + + if tier == "550": + if not _ensure_venv_t5_550_exists(): + raise RuntimeError( + f"Cannot activate transformers 5.5.0: .venv_t5_550 missing at {_VENV_T5_550_DIR}" + ) + if _VENV_T5_550_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_550_DIR) + logger.info("Activated transformers 5.5.0 from %s", _VENV_T5_550_DIR) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_550_DIR + (os.pathsep + _pp if _pp else "") + elif tier == "530": + if not _ensure_venv_t5_530_exists(): + raise RuntimeError( + f"Cannot activate transformers 5.3.0: .venv_t5_530 missing at {_VENV_T5_530_DIR}" + ) + if _VENV_T5_530_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_530_DIR) + logger.info("Activated transformers 5.3.0 from %s", _VENV_T5_530_DIR) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_530_DIR + (os.pathsep + _pp if _pp else "") + else: + logger.info("Using default transformers (4.57.x) for %s", model_name) def _decode_image(image_base64: str): @@ -285,18 +322,13 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: except Exception as e: logger.warning("Could not read adapter_config.json: %s", e) - # Auto-enable trust_remote_code for NemotronH/Nano models only. + # Auto-enable trust_remote_code for Nemotron models only. # NemotronH has config parsing bugs requiring trust_remote_code=True. # Other transformers 5.x models are native and do NOT need it. - # NOTE: Must NOT match Llama-Nemotron (standard Llama architecture). - _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") trust_remote_code = config.get("trust_remote_code", False) if not trust_remote_code: model_name = config["model_name"] - _mn_lower = model_name.lower() - if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/") - ): + if "nemotron" in model_name.lower(): trust_remote_code = True logger.info( "Auto-enabled trust_remote_code for Nemotron model: %s", diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index c3d3daf07e..b412a05cc9 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -86,7 +86,6 @@ def _probe_causal_conv1d_env() -> dict[str, str] | None: "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', " "'torch_mm': torch_mm, " "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', " - "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', " "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()" "}))" ), @@ -238,111 +237,28 @@ def _install_package_wheel_first( else: logger.info("No published %s wheel found: %s", display_name, wheel_url) - is_hip = env and env.get("hip_version") - if is_hip and not shutil.which("hipcc"): - logger.error( - "%s requires hipcc for source compilation on ROCm. " - "Install the ROCm HIP SDK: https://rocm.docs.amd.com", - display_name, - ) - _send_status( - event_queue, - f"{display_name}: hipcc not found (ROCm HIP SDK required)", - ) - return - - if is_hip: - _send_status( - event_queue, - f"Compiling {display_name} from source for ROCm " - "(this may take several minutes)...", - ) - else: - _send_status(event_queue, f"Installing {display_name} from PyPI...") - - # Prefer uv for faster dependency resolution when available - if shutil.which("uv"): - pypi_cmd = [ - "uv", - "pip", - "install", - "--python", - sys.executable, - "--no-build-isolation", - "--no-deps", - ] - # Avoid stale cache artifacts from partial HIP source builds - if is_hip: - pypi_cmd.append("--no-cache") - pypi_cmd.append(f"{pypi_name}=={pypi_version}") - else: - pypi_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--no-build-isolation", - "--no-deps", - "--no-cache-dir", - f"{pypi_name}=={pypi_version}", - ] - - # Source compilation on ROCm can take 10-30 minutes; use a generous - # timeout. Non-HIP installs preserve the pre-existing "no timeout" - # behaviour so unrelated slow installs (e.g. causal-conv1d source - # build on Linux aarch64 or unsupported torch/CUDA combinations) - # are not aborted at 5 minutes by this PR. - _run_kwargs: dict[str, Any] = { - "stdout": _sp.PIPE, - "stderr": _sp.STDOUT, - "text": True, - } - if is_hip: - _run_kwargs["timeout"] = 1800 - - try: - result = _sp.run(pypi_cmd, **_run_kwargs) - except _sp.TimeoutExpired: - logger.error( - "%s installation timed out after %ds", - display_name, - _run_kwargs.get("timeout"), - ) - _send_status( - event_queue, - f"{display_name} installation timed out after " - f"{_run_kwargs.get('timeout')}s", - ) - return - + _send_status(event_queue, f"Installing {display_name} from PyPI...") + pypi_cmd = [ + sys.executable, + "-m", + "pip", + "install", + "--no-build-isolation", + "--no-deps", + "--no-cache-dir", + f"{pypi_name}=={pypi_version}", + ] + result = _sp.run( + pypi_cmd, + stdout = _sp.PIPE, + stderr = _sp.STDOUT, + text = True, + ) if result.returncode != 0: - if is_hip: - # Surface a clear error for ROCm source build failures - error_lines = (result.stdout or "").strip().splitlines() - snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)" - logger.error( - "Failed to compile %s for ROCm:\n%s", - display_name, - result.stdout, - ) - _send_status( - event_queue, - f"Failed to compile {display_name} for ROCm. " - "Check that hipcc and ROCm development headers are installed.\n" - f"{snippet}", - ) - else: - logger.error( - "Failed to install %s from PyPI:\n%s", - display_name, - result.stdout, - ) + logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout) return - if is_hip: - logger.info("Compiled and installed %s from source for ROCm", display_name) - else: - logger.info("Installed %s from PyPI", display_name) + logger.info("Installed %s from PyPI", display_name) def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: @@ -390,15 +306,50 @@ def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None: def _activate_transformers_version(model_name: str) -> None: - """Activate the correct transformers version BEFORE any ML imports.""" + """Activate the correct transformers version BEFORE any ML imports. + + Uses get_transformers_tier() to decide between .venv_t5_550/ (5.5.0), + .venv_t5_530/ (5.3.0), or the default 4.57.x. + """ # Ensure backend is on path for utils imports backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) - from utils.transformers_version import activate_transformers_for_subprocess + from utils.transformers_version import ( + get_transformers_tier, + _resolve_base_model, + _ensure_venv_t5_530_exists, + _ensure_venv_t5_550_exists, + _VENV_T5_530_DIR, + _VENV_T5_550_DIR, + ) - activate_transformers_for_subprocess(model_name) + resolved = _resolve_base_model(model_name) + tier = get_transformers_tier(resolved) + + if tier == "550": + if not _ensure_venv_t5_550_exists(): + raise RuntimeError( + f"Cannot activate transformers 5.5.0: .venv_t5_550 missing at {_VENV_T5_550_DIR}" + ) + if _VENV_T5_550_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_550_DIR) + logger.info("Activated transformers 5.5.0 from %s", _VENV_T5_550_DIR) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_550_DIR + (os.pathsep + _pp if _pp else "") + elif tier == "530": + if not _ensure_venv_t5_530_exists(): + raise RuntimeError( + f"Cannot activate transformers 5.3.0: .venv_t5_530 missing at {_VENV_T5_530_DIR}" + ) + if _VENV_T5_530_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_530_DIR) + logger.info("Activated transformers 5.3.0 from %s", _VENV_T5_530_DIR) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_530_DIR + (os.pathsep + _pp if _pp else "") + else: + logger.info("Using default transformers (4.57.x) for %s", model_name) def run_training_process( @@ -448,17 +399,14 @@ def run_training_process( ) return - # ── 1a. Auto-enable trust_remote_code for NemotronH/Nano models ── + # ── 1a. Auto-enable trust_remote_code for Nemotron models ── # NemotronH has config parsing bugs in transformers that require # trust_remote_code=True as a workaround. Other transformers 5.x models # (Qwen3.5, Gemma 4, etc.) are native and do NOT need it — enabling it # bypasses the compiler (disabling fused CE). - # NOTE: Must NOT match Llama-Nemotron (standard Llama architecture). - _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _lowered = model_name.lower() if ( - any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS) - and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/")) + "nemotron" in _lowered and not config.get("trust_remote_code", False) ): config["trust_remote_code"] = True