diff --git a/src/llamafactory/train/callbacks.py b/src/llamafactory/train/callbacks.py index b3826a7a9db..04459a65f27 100644 --- a/src/llamafactory/train/callbacks.py +++ b/src/llamafactory/train/callbacks.py @@ -13,6 +13,9 @@ # limitations under the License. import fnmatch +import csv +import io +import subprocess import json import os import signal @@ -51,6 +54,111 @@ logger = logging.get_logger(__name__) +_XPU_TELEMETRY_FIELDS: tuple[tuple[tuple[str, ...], str], ...] = ( + # Compute utilization: use metric 31 (Compute engine group utilization) which + # measures compute engines only, providing the most accurate view of AI/ML workload. + (("Compute engine group utilization (%)",), "xpu_compute_util_pct"), + (("GPU Memory Utilization (%)",), "xpu_mem_bandwidth_pct"), + (("GPU Memory Used (MiB)",), "xpu_mem_in_use_mib"), +) + + +def _sample_xpu_device_metrics(device_id: int = 0) -> dict[str, float]: + r"""Sample per-device XPU telemetry via the ``xpu-smi dump`` interface. + + ``xpu-smi dump`` emits CSV (the ``-j`` flag is only honored by ``stats``/``discovery``), + so we parse the header row to map column names onto our metric keys. Returns an + empty dict if ``xpu-smi`` is missing, times out, or returns malformed output. + A single warning is logged per failure mode. + + Metric IDs requested: 5=Mem Util, 18=Mem Used (MiB), 31=Compute Engine Group Util. + Metric 31 provides the most accurate compute utilization for AI/ML workloads. + """ + metrics: dict[str, float] = {} + cmd = ["sudo", "xpu-smi", "dump", "-d", str(device_id), "-m", "5,18,31", "-n", "1"] + + try: + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=3) + except FileNotFoundError: + logger.warning_rank0_once( + "xpu-smi binary not found on PATH; XPU telemetry will be skipped. " + "Install Intel XPU Manager or unset RECORD_XPU to silence this warning." + ) + return metrics + except subprocess.TimeoutExpired: + logger.warning_rank0_once("xpu-smi query timed out; XPU telemetry will be skipped for this step.") + return metrics + except Exception as err: + logger.warning_rank0_once(f"xpu-smi query failed ({err!r}); XPU telemetry will be skipped.") + return metrics + + if proc.returncode != 0: + logger.warning_rank0_once( + f"xpu-smi exited with code {proc.returncode}; XPU telemetry will be skipped. " + f"stderr: {proc.stderr.strip()[:200]}" + ) + return metrics + + rows = list(csv.reader(io.StringIO(proc.stdout))) + # Expect a header line plus at least one data row. + data_rows = [r for r in rows if r and r[0].strip()] + if len(data_rows) < 2: + logger.warning_rank0_once("xpu-smi returned no data rows; XPU telemetry will be skipped.") + return metrics + + header = [col.strip() for col in data_rows[0]] + values = [col.strip() for col in data_rows[1]] + row = dict(zip(header, values)) + + for src_keys, dst_key in _XPU_TELEMETRY_FIELDS: + for src_key in src_keys: + raw = row.get(src_key) + if raw is None or raw == "" or raw.lower() == "n/a": + continue + try: + metrics[dst_key] = round(float(raw), 2) + break + except (ValueError, TypeError): + continue + + return metrics + + +def _sample_host_resource_usage() -> dict[str, float]: + r"""Sample host-side CPU and RAM utilization via ``psutil``. + + Returns an empty dict when ``psutil`` is not installed (warned once). RAM metrics + are always included when available. The first call to ``psutil.cpu_percent(interval=None)`` + has no prior sample to diff against and would report 0.0, so ``host_cpu_busy_pct`` + is omitted on the first invocation; the counter is primed and reported from the + second invocation onward. + """ + try: + import psutil + except ImportError: + logger.warning_rank0_once( + "psutil is not installed; host CPU/RAM telemetry will be skipped. " + "Run `pip install psutil` or unset RECORD_CPU to silence this warning." + ) + return {} + + try: + first_call = not getattr(_sample_host_resource_usage, "_primed", False) + cpu_pct = psutil.cpu_percent(interval=None) + _sample_host_resource_usage._primed = True + + vmem = psutil.virtual_memory() + result: dict[str, float] = { + "host_ram_in_use_gib": round(vmem.used / (1024**3), 2), + "host_ram_full_pct": round(vmem.percent, 2), + } + + if not first_call: + result["host_cpu_busy_pct"] = round(cpu_pct, 2) + return result + except Exception as err: + logger.warning_rank0_once(f"psutil host query failed ({err!r}); host telemetry will be skipped.") + return {} def fix_valuehead_checkpoint( model: "AutoModelForCausalLMWithValueHead", output_dir: str, safe_serialization: bool @@ -216,6 +324,22 @@ def _write_log(self, output_dir: str, logs: dict[str, Any]) -> None: with open(os.path.join(output_dir, TRAINER_LOG), "a", encoding="utf-8") as f: f.write(json.dumps(logs) + "\n") + def _sample_and_write_log(self, output_dir: str, logs: dict[str, Any], device_id: int) -> None: + r"""Sample device/host metrics on the worker thread and append to the log file. + + ``_sample_xpu_device_metrics`` shells out to ``xpu-smi`` (subprocess + IO) and + ``_sample_host_resource_usage`` calls ``psutil``; both are blocking. Running them + here keeps the main training thread free, which matters when ``logging_steps`` is + small. Metrics are not used by the Web UI's inline log line, so deferring them + only affects the persisted ``trainer_log.jsonl`` content. + """ + if is_env_enabled("RECORD_XPU"): + logs.update(_sample_xpu_device_metrics(device_id)) + if is_env_enabled("RECORD_CPU"): + logs.update(_sample_host_resource_usage()) + logs = {k: v for k, v in logs.items() if v is not None} + self._write_log(output_dir, logs) + def _create_thread_pool(self, output_dir: str) -> None: os.makedirs(output_dir, exist_ok=True) self.thread_pool = ThreadPoolExecutor(max_workers=1) @@ -297,7 +421,16 @@ def on_log(self, args: "TrainingArguments", state: "TrainerState", control: "Tra logs["vram_allocated"] = round(vram_allocated / (1024**3), 2) logs["vram_reserved"] = round(vram_reserved / (1024**3), 2) - logs = {k: v for k, v in logs.items() if v is not None} + # Resolve the active XPU device on the main thread (cheap torch query) so the + # worker can sample without touching torch state. Actual xpu-smi / psutil + # sampling is offloaded below to avoid blocking training on logging steps. + xpu_device_id = 0 + if is_env_enabled("RECORD_XPU") and hasattr(torch, "xpu") and torch.xpu.is_available(): + try: + xpu_device_id = torch.xpu.current_device() + except Exception: + xpu_device_id = 0 + if self.webui_mode and all(key in logs for key in ("loss", "lr", "epoch")): log_str = f"'loss': {logs['loss']:.4f}, 'learning_rate': {logs['lr']:2.4e}, 'epoch': {logs['epoch']:.2f}" for extra_key in ("reward", "accuracy", "throughput"): @@ -307,7 +440,9 @@ def on_log(self, args: "TrainingArguments", state: "TrainerState", control: "Tra logger.info_rank0("{" + log_str + "}") if self.thread_pool is not None: - self.thread_pool.submit(self._write_log, args.output_dir, logs) + # Offload XPU/host sampling (xpu-smi subprocess + psutil) plus the file write + # onto the single-worker pool so they cannot stall the training step. + self.thread_pool.submit(self._sample_and_write_log, args.output_dir, logs, xpu_device_id) @override def on_prediction_step(