"""CPU training backend — real SFT/LoRA via TRL, no GPU required. The closed-loop case: mindX produces a small JSONL dataset (dream cycle), mindXtrain needs to fine-tune a tiny base model on it locally without provisioning a MI300X droplet. This backend exists so that mindX agents can trigger self-training on commodity hardware; it is also the smoke lane for any new recipe before burning AMD credits. Slow but produces a *real* checkpoint and is compatible with the rest of the pipeline (`mindxtrain quantize`, `mindxtrain receipt`, `publish` — all expect a HF-format checkpoint directory). This module follows the project lazy-import contract: importing the module must succeed on a base install, but calling `run_trl_cpu` requires `uv sync --extra ml`. """ from __future__ import annotations import os from collections.abc import Callable, Iterator from dataclasses import dataclass from pathlib import Path from typing import Any from mindxtrain.autotune.plan import AutotunePlan from mindxtrain.config.schema import ( LoraMethod, QLoraMethod, XTrainConfig, resolve_thread_count, ) def _apply_cpu_throttle(cfg: XTrainConfig, sink: Callable[[str], None]) -> int: """Apply the recipe's `cfg.train.cpu_throttle` to the current process. Resolves percent → thread count against the host's actual core count, sets every thread-pool env var the downstream stack respects (torch, OpenMP, MKL, OpenBLAS), optionally pins OpenMP threads to cores via OMP_PROC_BIND=close + OMP_PLACES=cores (Ryzen-friendly: keeps threads on the same CCX chiplet, reduces cross-CCX cache traffic for the small matmuls CPU training does), and shifts the process's POSIX nice level so the rest of the laptop stays usable. Must be called BEFORE torch is imported / first used, otherwise the thread-pool size is locked at whatever torch saw on first init. Returns the resolved thread count for logging. """ throttle = cfg.train.cpu_throttle total_cores = os.cpu_count() or 1 threads = resolve_thread_count(throttle.percent, total_cores) # Set every thread-pool env var the downstream stack reads. PyTorch's # ATen reads OMP_NUM_THREADS at first init; MKL and OpenBLAS each # have their own knob. All must agree to actually cap the workload. os.environ["OMP_NUM_THREADS"] = str(threads) os.environ["MKL_NUM_THREADS"] = str(threads) os.environ["OPENBLAS_NUM_THREADS"] = str(threads) os.environ["NUMEXPR_NUM_THREADS"] = str(threads) # tokenizers (HF Rust lib) has its own parallelism that BLAS doesn't # cap. Disable it during throttled runs — for a 135M smoke on 1-2 # threads, the tokenizer-side parallelism only thrashes the cache. os.environ.setdefault("TOKENIZERS_PARALLELISM", "false" if threads <= 2 else "true") if throttle.omp_proc_bind: # CCX-aware pinning. `close` = threads adjacent to the master; # `cores` = one thread per physical core. Both safe on non-AMD # CPUs; ignored by OpenMP runtimes that don't honor them. os.environ["OMP_PROC_BIND"] = "close" os.environ["OMP_PLACES"] = "cores" if throttle.nice_level != 0: try: os.nice(throttle.nice_level) sink(f"[trl_cpu] nice level set to {throttle.nice_level}") except (PermissionError, OSError) as exc: # Negative nice needs CAP_SYS_NICE. Surface, don't fail. sink( f"[trl_cpu] nice({throttle.nice_level}) refused " f"({type(exc).__name__}: {exc}) — continuing without it", ) sink( f"[trl_cpu] throttle: {throttle.percent}% of {total_cores} cores " f"→ {threads} thread(s); OMP_PROC_BIND={'close' if throttle.omp_proc_bind else 'off'}", ) return threads def _require_ml_deps() -> dict[str, Any]: """Import TRL + transformers + peft + datasets eagerly; surface a single message.""" missing: list[str] = [] try: from datasets import Dataset # type: ignore except ImportError: missing.append("datasets") Dataset = None # type: ignore try: from transformers import AutoModelForCausalLM, AutoTokenizer # type: ignore except ImportError: missing.append("transformers") AutoModelForCausalLM = AutoTokenizer = None # type: ignore try: from trl import SFTConfig, SFTTrainer # type: ignore except ImportError: missing.append("trl") SFTConfig = SFTTrainer = None # type: ignore try: from peft import LoraConfig # type: ignore except ImportError: # peft only required for LoRA/QLoRA — kept optional here LoraConfig = None # type: ignore if missing: msg = ( f"CPU training backend requires {', '.join(missing)} — " "run `uv sync --extra ml`." ) raise RuntimeError(msg) return { "Dataset": Dataset, "AutoModelForCausalLM": AutoModelForCausalLM, "AutoTokenizer": AutoTokenizer, "SFTConfig": SFTConfig, "SFTTrainer": SFTTrainer, "LoraConfig": LoraConfig, } def _stream_dataset_rows(cfg: XTrainConfig) -> Iterator[dict[str, Any]]: """Yield raw dataset rows for the configured DataCfg.source.""" from mindxtrain.data.curate import load_streaming_dataset yield from load_streaming_dataset(cfg.data) def _materialize_dataset(cfg: XTrainConfig, Dataset: Any) -> Any: """Pull the stream into an in-memory `datasets.Dataset` for TRL. CPU corpora are small by construction; we don't try to stream into TRL here since `SFTTrainer` wants a `Dataset` with `__len__`. Cap at `cfg.data.max_samples or 50000` to keep memory bounded. """ cap = cfg.data.max_samples if cfg.data.max_samples is not None else 50_000 rows: list[dict[str, Any]] = [] for row in _stream_dataset_rows(cfg): rows.append(row) if len(rows) >= cap: break if not rows: msg = ( f"data.source={cfg.data.source!r} yielded zero examples — " "check the `path` / `hf_id` and that the source has data." ) raise RuntimeError(msg) return Dataset.from_list(rows) def _build_lora_config(cfg: XTrainConfig, LoraConfig: Any) -> Any | None: method = cfg.train.method if isinstance(method, LoraMethod) or isinstance(method, QLoraMethod): if LoraConfig is None: msg = "peft not installed; `uv sync --extra ml` or drop method.kind to 'full'." raise RuntimeError(msg) return LoraConfig( r=method.r, lora_alpha=method.alpha, lora_dropout=method.dropout, target_modules=list(method.target_modules), bias="none", task_type="CAUSAL_LM", ) return None def _build_event_callback( on_event: Callable[[dict[str, Any]], None], sink: Callable[[str], None], ) -> Any: """Return a TrainerCallback that fires `on_event(dict)` per HF Trainer log. Bridges in-process Trainer logs into the operator's SSE event stream without going through stdout regex parsing. Each `on_log` call carries a `{loss, learning_rate, grad_norm, …}` dict — we lift that into the same shape `StepEvent` expects (step / loss / lr / grad_norm), and emit `eval` + `status` events on the other Trainer hooks. Lazy-imported `transformers.TrainerCallback` so this helper is only realised when the backend actually runs (consistent with the rest of this module). """ from transformers import TrainerCallback # type: ignore class _CB(TrainerCallback): # type: ignore[misc, valid-type] def on_log( self, args: Any, state: Any, control: Any, logs: dict[str, float] | None = None, **_kw: Any, ) -> None: if not logs: return # HF Trainer emits multiple kinds of log: train step (has 'loss'), # final summary (has 'train_loss'), and eval (has 'eval_loss'). # Map step-level logs to StepEvent. if "loss" in logs: on_event({ "kind": "step", "step": int(state.global_step), "loss": float(logs["loss"]), "lr": float(logs["learning_rate"]) if "learning_rate" in logs else None, "grad_norm": float(logs["grad_norm"]) if "grad_norm" in logs else None, "tokens_per_s": None, # Realtime-feedback fields for the Coach progress bar + # "is it learning" chart. state.max_steps is the # authoritative total HF resolved after packing. "total_steps": ( int(state.max_steps) if getattr(state, "max_steps", 0) else None ), "mean_token_accuracy": ( float(logs["mean_token_accuracy"]) if "mean_token_accuracy" in logs else None ), "entropy": ( float(logs["entropy"]) if "entropy" in logs else None ), }) sink( f"[trl_cpu] step={state.global_step} loss={logs['loss']:.4f}" + (f" lr={logs['learning_rate']:.2e}" if "learning_rate" in logs else "") + (f" grad_norm={logs['grad_norm']:.3f}" if "grad_norm" in logs else ""), ) def on_evaluate( self, args: Any, state: Any, control: Any, metrics: dict[str, float] | None = None, **_kw: Any, ) -> None: if not metrics: return clean = {k: float(v) for k, v in metrics.items() if isinstance(v, (int, float))} if clean: on_event({ "kind": "eval", "step": int(state.global_step), "suite": "mid_train", "metrics": clean, }) return _CB() def _env_force_cpu() -> bool: """`MINDXTRAIN_FORCE_CPU` forces the CPU fallback even when a GPU exists. The escape hatch for parity testing and for pinning the deterministic CPU path on a box that happens to have an accelerator. """ return os.environ.get("MINDXTRAIN_FORCE_CPU", "").strip().lower() in { "1", "true", "yes", "on", } @dataclass(frozen=True) class _DevicePlan: """Resolved device/dtype decisions for one in-process training run.""" label: str # human tag e.g. "cuda (bfloat16)" / "cpu (float32)" device_map: dict[str, str] # passed to from_pretrained torch_dtype: Any # a torch.dtype attn_impl: str # "sdpa" on GPU, "eager" on CPU bf16: bool fp16: bool apply_cpu_throttle: bool # only the CPU path throttles threads batch_cap: int | None # cap per-device batch (CPU=2); None = uncapped grad_checkpointing: bool def _resolve_device(cfg: XTrainConfig, torch_mod: Any, *, force_cpu: bool) -> _DevicePlan: """Pick the device/dtype for a run. GPU when one is visible (ROCm exposes itself through `torch.cuda`) and not forced off; otherwise the exact CPU defaults the legacy `trl_cpu` lane used. """ if not force_cpu and torch_mod.cuda.is_available(): bf16 = bool(torch_mod.cuda.is_bf16_supported()) return _DevicePlan( label=f"cuda ({'bfloat16' if bf16 else 'float16'})", device_map={"": "cuda"}, torch_dtype=torch_mod.bfloat16 if bf16 else torch_mod.float16, attn_impl="sdpa", bf16=bf16, fp16=not bf16, apply_cpu_throttle=False, batch_cap=None, grad_checkpointing=cfg.train.gradient_checkpointing, ) return _DevicePlan( label="cpu (float32)", device_map={"": "cpu"}, torch_dtype=torch_mod.float32, attn_impl="eager", bf16=False, fp16=False, apply_cpu_throttle=True, batch_cap=2, grad_checkpointing=False, # CPU + checkpointing is pathologically slow ) def _capped_batch(per_device: int, cap: int | None) -> int: """Clamp the per-device batch to `cap` (CPU=2); `None` leaves it as-is.""" return max(1, per_device if cap is None else min(per_device, cap)) def run_trl_cpu( cfg: XTrainConfig, plan: AutotunePlan, out_dir: Path, *, on_line: Callable[[str], None] | None = None, on_event: Callable[[dict[str, Any]], None] | None = None, ) -> Path: """Force-CPU wrapper around `run_trl_local` (the mindX self-training lane). Preserves the deterministic CPU behaviour the dream-cycle loop depends on: float32, eager attention, CPU throttle, batch cap 2, no gradient checkpointing. """ return run_trl_local( cfg, plan, out_dir, force_cpu=True, on_line=on_line, on_event=on_event, ) def run_trl_local( cfg: XTrainConfig, plan: AutotunePlan, out_dir: Path, *, force_cpu: bool = False, on_line: Callable[[str], None] | None = None, on_event: Callable[[dict[str, Any]], None] | None = None, ) -> Path: """Run a TRL SFT job on the local device; return the checkpoint directory. Auto-detects an accelerator: uses the GPU (CUDA or ROCm, bf16/fp16) when one is visible, otherwise falls back to CPU (float32) with a logged warning. Set `force_cpu=True` (or `MINDXTRAIN_FORCE_CPU=1`) to pin the CPU path. Same `on_line` / `on_event` streaming contract as the other lanes so the Coach UI's live log + loss curve work identically on GPU and CPU. `on_line` mirrors the axolotl backend signature so the Coach UI can stream log lines uniformly across lanes. TRL doesn't emit one-line-per- step by default, but we forward `transformers` log records via a tiny handler so the streaming surface stays consistent. `on_event` is the *structured* counterpart: each HF Trainer log fires a dict with `{kind: "step", step, loss, lr, grad_norm, ...}` so the Coach UI can populate its Chart.js loss curve directly, without stdout-parsing the way the axolotl subprocess path does. When provided, the loss chart fills in in real time during training. On the CPU path, applies `cfg.train.cpu_throttle` before torch initializes so the thread-pool size actually caps the workload (BLAS layers snapshot their thread count on first use). The GPU path skips throttling entirely. """ sink = on_line if on_line is not None else (lambda _line: None) forced_cpu = force_cpu or _env_force_cpu() tag = "trl_cpu" if forced_cpu else "trl_local" # When the CPU path is known up front (force_cpu), throttle BEFORE importing # the ML stack — exactly the legacy ordering the dream-loop relies on. threads: int | None = None if forced_cpu: threads = _apply_cpu_throttle(cfg, sink) deps = _require_ml_deps() Dataset = deps["Dataset"] AutoModelForCausalLM = deps["AutoModelForCausalLM"] AutoTokenizer = deps["AutoTokenizer"] SFTConfig = deps["SFTConfig"] SFTTrainer = deps["SFTTrainer"] LoraConfig = deps["LoraConfig"] import torch # type: ignore device = _resolve_device(cfg, torch, force_cpu=forced_cpu) if device.apply_cpu_throttle: if threads is None: # Auto-detected CPU fallback (no accelerator visible). Throttle now; # `set_num_threads` still caps at runtime even post-import. if not forced_cpu: sink(f"[{tag}] no accelerator detected → CPU fallback (float32)") threads = _apply_cpu_throttle(cfg, sink) # ATen reads OMP_NUM_THREADS at init, but `set_num_threads` is the # canonical knob — call both belt-and-suspenders. torch.set_num_threads(threads) try: torch.set_num_interop_threads(threads) except RuntimeError: # Must be called before any aten op runs; OMP env var is the fallback. pass else: sink(f"[{tag}] accelerator detected → {device.label}") out_dir = Path(out_dir) checkpoint_dir = out_dir / "checkpoint" checkpoint_dir.mkdir(parents=True, exist_ok=True) sink(f"[{tag}] base={cfg.model.name} method={cfg.train.method.kind}") tokenizer = AutoTokenizer.from_pretrained(cfg.model.name, use_fast=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # Some base tokenizers (e.g., SmolLM2-135M) don't ship a chat # template. The dream corpus is ChatML-shaped already, so we set # ChatML explicitly when missing — matches what mindX's machine_ # dreaming phase 5b emits. Qwen / Llama base models keep their own # template untouched. if getattr(tokenizer, "chat_template", None) is None: tokenizer.chat_template = ( "{% for message in messages %}" "<|im_start|>{{ message['role'] }}\n{{ message['content'] }}<|im_end|>\n" "{% endfor %}" "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}" ) sink(f"[{tag}] tokenizer had no chat_template; set ChatML default.") sink(f"[{tag}] materializing dataset (in-memory)") full_dataset = _materialize_dataset(cfg, Dataset) sink(f"[{tag}] dataset size={len(full_dataset)}") # Optional train/eval split — deterministic via meta.seed so the same # recipe always carves the same held-out rows. When eval_split is # None we pass the full dataset to SFTTrainer (legacy behaviour); # otherwise we split, pass train_dataset + eval_dataset, and turn on # step-based evaluation so eval_loss appears in log_history. train_dataset = full_dataset eval_dataset = None if cfg.data.eval_split is not None: n = len(full_dataset) if n < 4: sink( f"[{tag}] dataset too small ({n} rows) for eval_split=" f"{cfg.data.eval_split}; skipping held-out split", ) else: split = full_dataset.train_test_split( test_size=cfg.data.eval_split, seed=cfg.meta.seed, ) train_dataset, eval_dataset = split["train"], split["test"] sink( f"[{tag}] split: train={len(train_dataset)} " f"eval={len(eval_dataset)} (seed={cfg.meta.seed})", ) sink(f"[{tag}] loading base model on {device.label}") model = AutoModelForCausalLM.from_pretrained( cfg.model.name, torch_dtype=device.torch_dtype, device_map=device.device_map, attn_implementation=device.attn_impl, ) peft_config = _build_lora_config(cfg, LoraConfig) # Estimate max_steps so we can floor logging_steps + eval_steps for # short runs. HF Trainer computes max_steps as # num_epochs * (len(train_dataset) // (batch * grad_accum)). When # packing is on, sample count drops post-tokenization — this estimate # is a lower bound but good enough for the cadence floor. per_device_batch = _capped_batch(cfg.train.batch.per_device, device.batch_cap) eff_batch = per_device_batch * cfg.train.batch.grad_accum est_max_steps = max(1, (len(train_dataset) // eff_batch) * cfg.train.schedule.epochs) effective_logging_steps = max(1, min(cfg.train.logging_steps, max(1, est_max_steps // 4))) effective_eval_steps = max(1, est_max_steps // 4) sink( f"[{tag}] est_max_steps={est_max_steps} " f"logging_steps={effective_logging_steps} " f"eval_steps={effective_eval_steps if eval_dataset is not None else 'off'}", ) sft_kwargs: dict[str, Any] = dict( output_dir=str(checkpoint_dir), num_train_epochs=cfg.train.schedule.epochs, per_device_train_batch_size=per_device_batch, gradient_accumulation_steps=cfg.train.batch.grad_accum, learning_rate=cfg.train.optimizer.lr, warmup_ratio=cfg.train.schedule.warmup_ratio, max_length=cfg.data.seq_len, packing=cfg.data.packing, logging_steps=effective_logging_steps, save_strategy="epoch", report_to="none", bf16=device.bf16, fp16=device.fp16, gradient_checkpointing=device.grad_checkpointing, seed=cfg.meta.seed, ) if eval_dataset is not None: sft_kwargs["eval_strategy"] = "steps" sft_kwargs["eval_steps"] = effective_eval_steps sft_kwargs["per_device_eval_batch_size"] = per_device_batch sft_args = SFTConfig(**sft_kwargs) trainer_kwargs: dict[str, Any] = dict( model=model, args=sft_args, train_dataset=train_dataset, peft_config=peft_config, processing_class=tokenizer, ) if eval_dataset is not None: trainer_kwargs["eval_dataset"] = eval_dataset trainer = SFTTrainer(**trainer_kwargs) # Wire the structured event callback if the caller wants per-step # telemetry (Coach UI's loss chart). When `on_event` is None this is # a no-op — the CLI path doesn't need it. if on_event is not None: trainer.add_callback(_build_event_callback(on_event, sink)) sink(f"[{tag}] starting trainer.train()") trainer.train() sink(f"[{tag}] training complete, saving checkpoint") trainer.save_model(str(checkpoint_dir)) tokenizer.save_pretrained(str(checkpoint_dir)) sink(f"[{tag}] checkpoint at {checkpoint_dir}") return checkpoint_dir __all__ = ["run_trl_cpu", "run_trl_local"]