Spaces:

Laksh718
/

daedalus-training-space

Paused

fix: add unsloth_zoo to bootstrap step

ad656e7 verified 7 months ago

2.8 kB

"""
HF Training Job entry point for DAEDALUS.
This script has ZERO external imports at the top level — it runs in a
bare Python 3.11 environment and bootstraps the full dependency stack
in the correct order before launching train_hf.py.
Dependency order is critical:
  1. torch 2.7.0 + cu124  — must be first so unsloth detects GPU correctly
  2. unsloth              — must come alone so it pins transformers/trl/peft
  3. everything else      — no torch/unsloth/transformers pins here
The daedalus/ package is provided by the mounted volume at /workspace.
"""

import os
import subprocess
import sys


def run(cmd: list, **kw) -> None:
    print(f">>> {' '.join(str(c) for c in cmd)}", flush=True)
    subprocess.run(cmd, check=True, **kw)


print("=" * 60)
print("DAEDALUS Job Runner — bootstrapping dependencies")
print("=" * 60, flush=True)

# ── Step 1: PyTorch 2.7.0 + CUDA 12.4 ────────────────────────────────────
# torch.int1 needs 2.6+; torch.utils._pytree.register_constant needs 2.7+
# torchvision intentionally omitted — not needed for LLM text training.
print("\n[step 1/3] Installing torch 2.7.0 (cu124)...")
run([
    sys.executable, "-m", "pip", "install",
    "torch==2.7.0",
    "--extra-index-url", "https://download.pytorch.org/whl/cu124",
    "-q",
])

# ── Step 2: Unsloth + unsloth_zoo ────────────────────────────────────────
# unsloth_zoo is a required companion package that the git install does not
# pull in automatically — must be installed explicitly before importing unsloth.
print("\n[step 2/3] Installing unsloth + unsloth_zoo...")
run([
    sys.executable, "-m", "pip", "install",
    "unsloth_zoo",
    "unsloth @ git+https://github.com/unslothai/unsloth.git",
    "-q",
])

# ── Step 3: Remaining dependencies ───────────────────────────────────────
print("\n[step 3/3] Installing remaining dependencies...")
run([
    sys.executable, "-m", "pip", "install",
    "huggingface_hub>=0.26.0",
    "hf_transfer>=0.1.8",
    "datasets>=3.0.0",
    "accelerate>=1.0.0",
    "peft>=0.13.0",
    "trl>=0.12.0",
    "bitsandbytes>=0.45.0",
    "sentencepiece>=0.2.0",
    "protobuf>=4.25.0",
    "python-dotenv>=1.0.0",
    "matplotlib>=3.7.0",
    "numpy>=1.24.0",
    "-q",
])

print("\n[bootstrap] all dependencies installed — launching train_hf.py", flush=True)
print("=" * 60, flush=True)

# The training Space is mounted at /workspace by the Job volume config.
# train_hf.py already has logic to add /workspace to sys.path for daedalus/.
os.chdir("/workspace")
run([sys.executable, "-u", "train_hf.py"])

	"""
	HF Training Job entry point for DAEDALUS.

	This script has ZERO external imports at the top level — it runs in a
	bare Python 3.11 environment and bootstraps the full dependency stack
	in the correct order before launching train_hf.py.

	Dependency order is critical:
	1. torch 2.7.0 + cu124 — must be first so unsloth detects GPU correctly
	2. unsloth — must come alone so it pins transformers/trl/peft
	3. everything else — no torch/unsloth/transformers pins here

	The daedalus/ package is provided by the mounted volume at /workspace.
	"""

	import os
	import subprocess
	import sys


	def run(cmd: list, **kw) -> None:
	print(f">>> {' '.join(str(c) for c in cmd)}", flush=True)
	subprocess.run(cmd, check=True, **kw)


	print("=" * 60)
	print("DAEDALUS Job Runner — bootstrapping dependencies")
	print("=" * 60, flush=True)

	# ── Step 1: PyTorch 2.7.0 + CUDA 12.4 ────────────────────────────────────
	# torch.int1 needs 2.6+; torch.utils._pytree.register_constant needs 2.7+
	# torchvision intentionally omitted — not needed for LLM text training.
	print("\n[step 1/3] Installing torch 2.7.0 (cu124)...")
	run([
	sys.executable, "-m", "pip", "install",
	"torch==2.7.0",
	"--extra-index-url", "https://download.pytorch.org/whl/cu124",
	"-q",
	])

	# ── Step 2: Unsloth + unsloth_zoo ────────────────────────────────────────
	# unsloth_zoo is a required companion package that the git install does not
	# pull in automatically — must be installed explicitly before importing unsloth.
	print("\n[step 2/3] Installing unsloth + unsloth_zoo...")
	run([
	sys.executable, "-m", "pip", "install",
	"unsloth_zoo",
	"unsloth @ git+https://github.com/unslothai/unsloth.git",
	"-q",
	])

	# ── Step 3: Remaining dependencies ───────────────────────────────────────
	print("\n[step 3/3] Installing remaining dependencies...")
	run([
	sys.executable, "-m", "pip", "install",
	"huggingface_hub>=0.26.0",
	"hf_transfer>=0.1.8",
	"datasets>=3.0.0",
	"accelerate>=1.0.0",
	"peft>=0.13.0",
	"trl>=0.12.0",
	"bitsandbytes>=0.45.0",
	"sentencepiece>=0.2.0",
	"protobuf>=4.25.0",
	"python-dotenv>=1.0.0",
	"matplotlib>=3.7.0",
	"numpy>=1.24.0",
	"-q",
	])

	print("\n[bootstrap] all dependencies installed — launching train_hf.py", flush=True)
	print("=" * 60, flush=True)

	# The training Space is mounted at /workspace by the Job volume config.
	# train_hf.py already has logic to add /workspace to sys.path for daedalus/.
	os.chdir("/workspace")
	run([sys.executable, "-u", "train_hf.py"])