From 4d0d38788cd248440a07171d4f7624b15395783c Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Sat, 20 Jun 2026 17:24:41 +0000 Subject: [PATCH 1/9] add unified profiler configuration --- src/llamafactory/extras/profiler.py | 378 ++++++++++++++++++++ src/llamafactory/hparams/training_args.py | 40 ++- src/llamafactory/train/callbacks.py | 79 ++-- src/llamafactory/train/tuner.py | 2 +- src/llamafactory/v1/accelerator/profiler.py | 29 ++ src/llamafactory/v1/config/training_args.py | 60 ++++ src/llamafactory/v1/core/base_trainer.py | 5 + 7 files changed, 526 insertions(+), 67 deletions(-) create mode 100644 src/llamafactory/extras/profiler.py diff --git a/src/llamafactory/extras/profiler.py b/src/llamafactory/extras/profiler.py new file mode 100644 index 0000000000..267a8ffeea --- /dev/null +++ b/src/llamafactory/extras/profiler.py @@ -0,0 +1,378 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import sys +from contextlib import suppress +from dataclasses import dataclass +from typing import Any, Optional + +import torch +from transformers.utils import is_torch_cuda_available, is_torch_npu_available + + +_SUPPORTED_ACTIVITIES = {"auto", "all", "cpu", "device"} +_SUPPORTED_RANK_MODES = {"all", "rank0"} + + +@dataclass +class ProfilerConfig: + enabled: bool = False + output_dir: Optional[str] = None + skip_first: int = 0 + wait_steps: int = 1 + warmup_steps: int = 1 + active_steps: int = 1 + repeat: int = 1 + record_shapes: bool = False + profile_memory: bool = False + with_stack: bool = False + with_flops: bool = False + with_modules: bool = False + activities: str = "auto" + rank_mode: str = "all" + deprecated_alias_present: bool = False + legacy_defaults_enabled: bool = False + explicit_profile_kwargs: Optional[set[str]] = None + + @classmethod + def from_args(cls, args: Any) -> "ProfilerConfig": + enable_profiler = bool(getattr(args, "enable_profiler", False)) + enable_torch_profiler = bool(getattr(args, "enable_torch_profiler", False)) + legacy_defaults_enabled = enable_torch_profiler and not enable_profiler + record_shapes = getattr(args, "profiler_record_shapes", None) + profile_memory = getattr(args, "profiler_profile_memory", None) + with_stack = getattr(args, "profiler_with_stack", None) + explicit_profile_kwargs = _get_explicit_profile_kwargs( + record_shapes=record_shapes, + profile_memory=profile_memory, + with_stack=with_stack, + ) + if getattr(args, "profiler_with_flops", False): + explicit_profile_kwargs.add("with_flops") + if getattr(args, "profiler_with_modules", False): + explicit_profile_kwargs.add("with_modules") + return cls( + enabled=enable_profiler or enable_torch_profiler, + output_dir=getattr(args, "profiler_output_dir", None), + skip_first=getattr(args, "profiler_skip_first", 0), + wait_steps=getattr(args, "profiler_wait_steps", 1), + warmup_steps=getattr(args, "profiler_warmup_steps", 1), + active_steps=getattr(args, "profiler_active_steps", 1), + repeat=getattr(args, "profiler_repeat", 1), + record_shapes=_resolve_optional_bool(record_shapes, legacy_defaults_enabled), + profile_memory=_resolve_optional_bool(profile_memory, legacy_defaults_enabled), + with_stack=_resolve_optional_bool(with_stack, legacy_defaults_enabled), + with_flops=getattr(args, "profiler_with_flops", False), + with_modules=getattr(args, "profiler_with_modules", False), + activities=getattr(args, "profiler_activities", "auto"), + rank_mode=getattr(args, "profiler_rank_mode", "all"), + deprecated_alias_present=enable_torch_profiler, + legacy_defaults_enabled=legacy_defaults_enabled, + explicit_profile_kwargs=explicit_profile_kwargs, + ) + + def validate(self, backend_name: Optional[str] = None) -> None: + if not self.enabled: + return + + _validate_int_option("profiler_skip_first", self.skip_first, min_value=0) + _validate_int_option("profiler_wait_steps", self.wait_steps, min_value=0) + _validate_int_option("profiler_warmup_steps", self.warmup_steps, min_value=0) + _validate_int_option("profiler_active_steps", self.active_steps, min_value=1) + _validate_int_option("profiler_repeat", self.repeat, min_value=0) + if self.activities not in _SUPPORTED_ACTIVITIES: + raise ValueError(f"`profiler_activities` must be one of {sorted(_SUPPORTED_ACTIVITIES)}.") + if self.rank_mode not in _SUPPORTED_RANK_MODES: + raise ValueError(f"`profiler_rank_mode` must be one of {sorted(_SUPPORTED_RANK_MODES)}.") + + self.schedule_kwargs() + + def schedule_kwargs(self) -> dict[str, int]: + return dict( + wait=self.wait_steps, + warmup=self.warmup_steps, + active=self.active_steps, + repeat=self.repeat, + skip_first=self.skip_first, + ) + + +class _ProfilerBackend: + def __init__(self, name: str, profiler_module: Any, device_activity_name: Optional[str]) -> None: + self.name = name + self.profiler_module = profiler_module + self.device_activity_name = device_activity_name + + def build_schedule(self, kwargs: dict[str, Any]) -> Any: + if self.name == "npu": + return self.profiler_module.schedule( + wait=kwargs["wait"], + active=kwargs["active"], + warmup=kwargs["warmup"], + repeat=kwargs["repeat"], + skip_first=kwargs["skip_first"], + ) + + return self.profiler_module.schedule(**kwargs) + + def build_activities(self, config: ProfilerConfig) -> list[Any]: + activity_cls = self.profiler_module.ProfilerActivity + activities = [] + if config.activities in ("auto", "all", "cpu"): + activities.append(activity_cls.CPU) + if config.activities == "device" and self.device_activity_name is None: + raise ValueError("`profiler_activities: device` requires a CUDA or NPU backend.") + if config.activities in ("auto", "all", "device"): + if self.device_activity_name is not None: + activities.append(getattr(activity_cls, self.device_activity_name)) + return activities + + def build_experimental_config(self, config: ProfilerConfig, logger: Any = None) -> Optional[Any]: + if self.name != "npu" or config.activities == "cpu": + return None + + profiler = self.profiler_module + config_kwargs = dict( + export_type=[profiler.ExportType.Text], + profiler_level=_get_profiler_enum_value(profiler, "ProfilerLevel", "Level0", "`profiler_level`"), + aic_metrics=_get_profiler_enum_value(profiler, "AiCMetrics", "AiCoreNone", "`profiler_aic_metrics`"), + l2_cache=False, + op_attr=False, + data_simplification=True, + record_op_args=False, + gc_detect_threshold=None, + host_sys=[], + sys_io=False, + sys_interconnection=False, + mstx=False, + mstx_domain_include=[], + mstx_domain_exclude=[], + ) + return profiler._ExperimentalConfig( + **_filter_supported_kwargs( + profiler._ExperimentalConfig, + config_kwargs, + logger=logger, + ) + ) + + +def _get_rank() -> int: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + + return int(os.getenv("RANK", "0")) + + +def _get_current_accelerator_type() -> str: + if is_torch_npu_available(): + return "npu" + if is_torch_cuda_available(): + return "cuda" + return "cpu" + + +def _resolve_optional_bool(value: Any, default: bool) -> bool: + if value is None: + return default + + return bool(value) + + +def _validate_int_option(name: str, value: Any, min_value: int) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value < min_value: + if min_value == 1: + raise ValueError(f"`{name}` must be a positive integer.") + + raise ValueError(f"`{name}` must be an integer greater than or equal to {min_value}.") + + +def _get_explicit_profile_kwargs(**kwargs: Any) -> set[str]: + name_mapping = { + "record_shapes": "record_shapes", + "profile_memory": "profile_memory", + "with_stack": "with_stack", + } + return {profile_key for arg_key, profile_key in name_mapping.items() if kwargs[arg_key] is not None} + + +def _get_backend() -> _ProfilerBackend: + accelerator_type = _get_current_accelerator_type() + if accelerator_type == "cpu": + return _ProfilerBackend("cpu", torch.profiler, None) + + if accelerator_type == "cuda": + return _ProfilerBackend("cuda", torch.profiler, "CUDA") + + if accelerator_type == "npu": + try: + import torch_npu # type: ignore + except ImportError as exc: + raise RuntimeError("NPU profiler requires `torch_npu` to be installed.") from exc + + return _ProfilerBackend("npu", torch_npu.profiler, "NPU") + + raise RuntimeError(f"Profiler only supports CPU, CUDA and NPU devices, got current accelerator: {accelerator_type}.") + + +def _get_profiler_enum_value(profiler: Any, enum_name: str, member_name: str, option_name: str) -> Any: + enum_cls = getattr(profiler, enum_name, None) + if enum_cls is not None and hasattr(enum_cls, member_name): + return getattr(enum_cls, member_name) + + supported = [] if enum_cls is None else [name for name in dir(enum_cls) if not name.startswith("_")] + package_name = str(getattr(profiler, "__name__", "")).split(".", maxsplit=1)[0] + torch_npu_version = getattr(sys.modules.get(package_name), "__version__", None) + raise RuntimeError( + f"The installed torch_npu profiler does not support {option_name}={member_name}. " + f"Supported values: {supported}. torch_npu version: {torch_npu_version or 'unknown'}." + ) + + +def _filter_supported_kwargs( + fn: Any, + kwargs: dict[str, Any], + logger: Any = None, + explicit_keys: Optional[set[str]] = None, +) -> dict[str, Any]: + parameters = inspect.signature(fn).parameters + if any(param.kind is inspect.Parameter.VAR_KEYWORD for param in parameters.values()): + return kwargs + + supported = set(parameters) + dropped = set(kwargs) - supported + if dropped: + explicit_dropped = dropped & (explicit_keys or set()) + message_level = "warning" if explicit_dropped else "debug" + if explicit_dropped or logger is not None: + _log( + logger, + message_level, + f"Profiler backend ignored unsupported arguments: {sorted(dropped)}.", + ) + + return {key: value for key, value in kwargs.items() if key in supported} + + +def _log(logger: Any, level: str, message: str) -> None: + if logger is None: + return + + rank_method = f"{level}_rank0" + if hasattr(logger, rank_method): + getattr(logger, rank_method)(message) + else: + getattr(logger, level)(message) + + +class ProfilerController: + def __init__(self, args: Any, logger: Any = None) -> None: + self.config = ProfilerConfig.from_args(args) + self.logger = logger + self.profiler = None + self.backend_name: Optional[str] = None + self.trace_dir: Optional[str] = None + + @property + def enabled(self) -> bool: + return self.config.enabled + + def start(self, output_dir: str, initial_step: int = 0) -> None: + self.stop() + if not self.config.enabled: + return + + if self.config.deprecated_alias_present: + message = "`enable_torch_profiler` is deprecated; use `enable_profiler` instead." + if self.config.legacy_defaults_enabled: + message += " Legacy shape, memory, and stack profiling defaults are preserved for this alias." + _log( + self.logger, + "warning", + message, + ) + + backend = _get_backend() + self.backend_name = backend.name + self.config.validate(backend.name) + if backend.name == "npu" and os.getenv("PROF_CONFIG_PATH"): + _log( + self.logger, + "warning", + "`PROF_CONFIG_PATH` is set. Do not enable dynamic_profile together with `enable_profiler`.", + ) + schedule = self.config.schedule_kwargs() + + rank = _get_rank() + if self.config.rank_mode == "rank0" and rank != 0: + return + + trace_root = self.config.output_dir or os.path.join(output_dir, "profiler") + self.trace_dir = os.path.realpath(os.path.join(trace_root, f"rank_{rank}")) + os.makedirs(self.trace_dir, exist_ok=True) + + profile_kwargs = dict( + activities=backend.build_activities(self.config), + schedule=backend.build_schedule(schedule), + on_trace_ready=backend.profiler_module.tensorboard_trace_handler(self.trace_dir), + record_shapes=self.config.record_shapes, + profile_memory=self.config.profile_memory, + with_stack=self.config.with_stack, + with_flops=self.config.with_flops, + with_modules=self.config.with_modules, + experimental_config=backend.build_experimental_config(self.config, logger=self.logger), + ) + profiler = backend.profiler_module.profile( + **_filter_supported_kwargs( + backend.profiler_module.profile, + profile_kwargs, + logger=self.logger, + explicit_keys=set(self.config.explicit_profile_kwargs or set()), + ) + ) + try: + profiler.start() + except Exception: + with suppress(Exception): + profiler.stop() + raise + + self.profiler = profiler + + first_active_step = initial_step + schedule["skip_first"] + schedule["wait"] + schedule["warmup"] + 1 + schedule_message = ( + f"skip_first={schedule['skip_first']}, wait={schedule['wait']}, warmup={schedule['warmup']}, " + f"active={schedule['active']}, repeat={schedule['repeat']}, first_active_step={first_active_step}" + ) + _log( + self.logger, + "info", + f"Profiler started on {backend.name}: {schedule_message}, initial_step={initial_step}. Traces -> {trace_root}", + ) + + def step(self) -> None: + if self.profiler is not None: + self.profiler.step() + + def stop(self) -> None: + profiler, self.profiler = self.profiler, None + if profiler is None: + return + + profiler.stop() + _log(self.logger, "info", "Profiler stopped.") diff --git a/src/llamafactory/hparams/training_args.py b/src/llamafactory/hparams/training_args.py index 84f4d8e6cc..6eee680b30 100644 --- a/src/llamafactory/hparams/training_args.py +++ b/src/llamafactory/hparams/training_args.py @@ -68,14 +68,22 @@ def __post_init__(self): class ProfilerArguments: r"""Arguments for torch profiler configuration.""" + enable_profiler: bool = field( + default=False, + metadata={"help": "Whether to enable profiler for collecting CPU/CUDA/NPU performance traces."}, + ) enable_torch_profiler: bool = field( default=False, - metadata={"help": "Whether to enable torch profiler for collecting performance traces."}, + metadata={"help": "Deprecated. Use `enable_profiler` instead."}, ) profiler_output_dir: Optional[str] = field( default=None, metadata={"help": "Directory to write profiler traces. Defaults to /profiler if not set."}, ) + profiler_skip_first: int = field( + default=0, + metadata={"help": "Number of steps to skip before the first profiler wait/warmup/active cycle."}, + ) profiler_wait_steps: int = field( default=1, metadata={"help": "Number of steps to skip at the start of each profiling cycle."}, @@ -92,23 +100,39 @@ class ProfilerArguments: default=1, metadata={"help": "Number of profiling cycles. Set to 0 for continuous profiling."}, ) - profiler_record_shapes: bool = field( - default=True, + profiler_record_shapes: Optional[bool] = field( + default=None, metadata={"help": "Whether to record tensor shapes during profiling."}, ) - profiler_profile_memory: bool = field( - default=True, + profiler_profile_memory: Optional[bool] = field( + default=None, metadata={"help": "Whether to profile memory usage."}, ) - profiler_with_stack: bool = field( - default=True, + profiler_with_stack: Optional[bool] = field( + default=None, metadata={"help": "Whether to record stack traces during profiling."}, ) + profiler_with_flops: bool = field( + default=False, + metadata={"help": "Whether to estimate FLOPs where supported by the profiler backend."}, + ) + profiler_with_modules: bool = field( + default=False, + metadata={"help": "Whether to record module hierarchy where supported by the profiler backend."}, + ) + profiler_activities: str = field( + default="auto", + metadata={"help": "Profiler activities to collect: auto, all, cpu, or device."}, + ) + profiler_rank_mode: str = field( + default="all", + metadata={"help": "Profiler rank collection mode: all or rank0."}, + ) profile_modules: Optional[str] = field( default=None, metadata={ "help": ( - "Comma-separated list of module name patterns to profile with CUDA events. " + "Comma-separated list of module name patterns to profile with accelerator events. " "Supports fnmatch wildcards (e.g. 'model.layers.0.self_attn,model.layers.*.mlp'). " "Reports per-module forward/backward timing statistics at each logging step." ) diff --git a/src/llamafactory/train/callbacks.py b/src/llamafactory/train/callbacks.py index b3826a7a9d..4d8bcd7eef 100644 --- a/src/llamafactory/train/callbacks.py +++ b/src/llamafactory/train/callbacks.py @@ -35,6 +35,7 @@ from ..extras.constants import TRAINER_LOG, V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME from ..extras.misc import get_peak_memory, is_env_enabled, is_torch_cuda_available, is_torch_npu_available, use_ray from ..extras.packages import is_safetensors_available +from ..extras.profiler import ProfilerController if is_safetensors_available(): @@ -341,93 +342,55 @@ def on_prediction_step( class TorchProfilerCallback(TrainerCallback): - r"""A callback for collecting torch.profiler traces during training. + r"""A callback for collecting CPU/CUDA/NPU profiler traces during training. - Activated by setting ``enable_torch_profiler: true`` in the YAML config. + Activated by setting ``enable_profiler: true`` in the YAML config. + ``enable_torch_profiler`` is still accepted for backward compatibility. Configuration fields (in YAML): profiler_output_dir – where to write traces (default: /profiler) + profiler_skip_first – steps to skip before the first cycle (default: 0) profiler_wait_steps – steps to skip at start of each cycle (default: 1) profiler_warmup_steps – profiler warm-up steps per cycle (default: 1) profiler_active_steps – steps to record per cycle (default: 1) profiler_repeat – number of cycles; 0 = forever (default: 1) - profiler_record_shapes – record tensor shapes (default: true) - profiler_profile_memory – profile memory usage (default: true) - profiler_with_stack – record stack traces (default: true) + profiler_record_shapes – record tensor shapes (default: false; true for deprecated alias) + profiler_profile_memory – profile memory usage (default: false; true for deprecated alias) + profiler_with_stack – record stack traces (default: false; true for deprecated alias) + profiler_activities – auto, all, cpu, or device (default: auto) + profiler_rank_mode – all ranks or rank0 only (default: all) Trace files (one per rank, Chrome / TensorBoard JSON format) are written to ``/rank_/``. + + Example: + profiler_skip_first: 8 # skip 8 optimizer steps before wait/warmup/active + profiler_wait_steps: 0 + profiler_warmup_steps: 1 + profiler_active_steps: 3 + profiler_rank_mode: rank0 """ def __init__(self, training_args: "TrainingArguments") -> None: - self.profiler = None - self.profiler_args = training_args - - @staticmethod - def _get_rank() -> int: - import torch.distributed as dist - - if dist.is_available() and dist.is_initialized(): - return dist.get_rank() - return 0 + self.profiler = ProfilerController(training_args, logger=logger) @override def on_train_begin( self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs ) -> None: - if self.profiler is not None: - self.profiler.stop() - self.profiler = None - - pa = self.profiler_args - output_dir = pa.profiler_output_dir or os.path.join(args.output_dir, "profiler") - rank = self._get_rank() - trace_dir = os.path.join(output_dir, f"rank_{rank}") - os.makedirs(trace_dir, exist_ok=True) - - activities = [torch.profiler.ProfilerActivity.CPU] - try: - if is_torch_cuda_available(): - activities.append(torch.profiler.ProfilerActivity.CUDA) - if is_torch_npu_available(): - activities.append(torch.profiler.ProfilerActivity.NPU) - except Exception: - pass - - self.profiler = torch.profiler.profile( - activities=activities, - schedule=torch.profiler.schedule( - wait=pa.profiler_wait_steps, - warmup=pa.profiler_warmup_steps, - active=pa.profiler_active_steps, - repeat=pa.profiler_repeat, - ), - on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir), - record_shapes=pa.profiler_record_shapes, - profile_memory=pa.profiler_profile_memory, - with_stack=pa.profiler_with_stack, - ) - self.profiler.start() - logger.info_rank0( - f"TorchProfiler started — schedule: wait={pa.profiler_wait_steps}, warmup={pa.profiler_warmup_steps}, " - f"active={pa.profiler_active_steps}, repeat={pa.profiler_repeat}. Traces → {output_dir}" - ) + self.profiler.start(args.output_dir, initial_step=state.global_step) @override def on_step_end( self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs ) -> None: - if self.profiler is not None: - self.profiler.step() + self.profiler.step() @override def on_train_end( self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs ) -> None: - if self.profiler is not None: - self.profiler.stop() - self.profiler = None - logger.info_rank0("TorchProfiler stopped.") + self.profiler.stop() class ReporterCallback(TrainerCallback): diff --git a/src/llamafactory/train/tuner.py b/src/llamafactory/train/tuner.py index 22f7dbf87a..7421a1ed10 100644 --- a/src/llamafactory/train/tuner.py +++ b/src/llamafactory/train/tuner.py @@ -80,7 +80,7 @@ def _training_function(config: dict[str, Any]) -> None: if finetuning_args.early_stopping_steps is not None: callbacks.append(EarlyStoppingCallback(early_stopping_patience=finetuning_args.early_stopping_steps)) - if getattr(training_args, "enable_torch_profiler", False): + if getattr(training_args, "enable_profiler", False) or getattr(training_args, "enable_torch_profiler", False): callbacks.append(TorchProfilerCallback(training_args)) if getattr(training_args, "profile_modules", None): diff --git a/src/llamafactory/v1/accelerator/profiler.py b/src/llamafactory/v1/accelerator/profiler.py index e69de29bb2..ba75573c58 100644 --- a/src/llamafactory/v1/accelerator/profiler.py +++ b/src/llamafactory/v1/accelerator/profiler.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from llamafactory.extras.profiler import ProfilerController + +from ..utils import logging +from ..utils.callbacks import TrainerCallback, TrainerState + + +if TYPE_CHECKING: + from ..config import TrainingArguments + + +logger = logging.get_logger(__name__) + + +class ProfilerCallback(TrainerCallback): + def __init__(self, args: TrainingArguments) -> None: + self.profiler = ProfilerController(args, logger=logger) + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + self.profiler.start(args.output_dir, initial_step=state.global_step) + + def on_step_end(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + self.profiler.step() + + def on_train_end(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + self.profiler.stop() diff --git a/src/llamafactory/v1/config/training_args.py b/src/llamafactory/v1/config/training_args.py index 200938b9cd..4cfd75c83a 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/src/llamafactory/v1/config/training_args.py @@ -69,6 +69,66 @@ class TrainingArguments: default=True, metadata={"help": "Enable activation checkpointing for training."}, ) + enable_profiler: bool = field( + default=False, + metadata={"help": "Whether to enable profiler for collecting CPU/CUDA/NPU performance traces."}, + ) + enable_torch_profiler: bool = field( + default=False, + metadata={"help": "Deprecated. Use `enable_profiler` instead."}, + ) + profiler_output_dir: str | None = field( + default=None, + metadata={"help": "Directory to write profiler traces. Defaults to /profiler if not set."}, + ) + profiler_skip_first: int = field( + default=0, + metadata={"help": "Number of steps to skip before the first profiler wait/warmup/active cycle."}, + ) + profiler_wait_steps: int = field( + default=1, + metadata={"help": "Number of steps to skip at the start of each profiling cycle."}, + ) + profiler_warmup_steps: int = field( + default=1, + metadata={"help": "Number of profiler warm-up steps per cycle."}, + ) + profiler_active_steps: int = field( + default=1, + metadata={"help": "Number of steps to actively record per cycle."}, + ) + profiler_repeat: int = field( + default=1, + metadata={"help": "Number of profiling cycles. Set to 0 for continuous profiling."}, + ) + profiler_record_shapes: bool | None = field( + default=None, + metadata={"help": "Whether to record tensor shapes during profiling."}, + ) + profiler_profile_memory: bool | None = field( + default=None, + metadata={"help": "Whether to profile memory usage."}, + ) + profiler_with_stack: bool | None = field( + default=None, + metadata={"help": "Whether to record stack traces during profiling."}, + ) + profiler_with_flops: bool = field( + default=False, + metadata={"help": "Whether to estimate FLOPs where supported by the profiler backend."}, + ) + profiler_with_modules: bool = field( + default=False, + metadata={"help": "Whether to record module hierarchy where supported by the profiler backend."}, + ) + profiler_activities: str = field( + default="auto", + metadata={"help": "Profiler activities to collect: auto, all, cpu, or device."}, + ) + profiler_rank_mode: str = field( + default="all", + metadata={"help": "Profiler rank collection mode: all or rank0."}, + ) dist_config: PluginConfig | None = field( default=None, metadata={"help": "Distribution configuration for training."}, diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index c2eb2bebd3..b08dbd6d5b 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -125,6 +125,11 @@ def __init__( # Callbacks self.callback_handler = CallbackHandler([LoggingCallback()], trainer=self) + if getattr(self.args, "enable_profiler", False) or getattr(self.args, "enable_torch_profiler", False): + from ..accelerator.profiler import ProfilerCallback + + self.callback_handler.add_callback(ProfilerCallback(self.args)) + for cb in callbacks or []: self.callback_handler.add_callback(cb) From 9f20b2adf55050c6ef55358a315540b2b12bbd16 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Sat, 20 Jun 2026 17:25:10 +0000 Subject: [PATCH 2/9] add npu profiler options --- src/llamafactory/extras/profiler.py | 176 ++++++++++++++++++-- src/llamafactory/hparams/training_args.py | 26 ++- src/llamafactory/train/callbacks.py | 15 +- src/llamafactory/v1/config/training_args.py | 26 ++- 4 files changed, 226 insertions(+), 17 deletions(-) diff --git a/src/llamafactory/extras/profiler.py b/src/llamafactory/extras/profiler.py index 267a8ffeea..186f28574b 100644 --- a/src/llamafactory/extras/profiler.py +++ b/src/llamafactory/extras/profiler.py @@ -25,6 +25,37 @@ _SUPPORTED_ACTIVITIES = {"auto", "all", "cpu", "device"} _SUPPORTED_RANK_MODES = {"all", "rank0"} +_NPU_PROFILER_LEVELS = { + "none": "Level_none", + "level0": "Level0", + "level1": "Level1", + "level2": "Level2", +} +_NPU_AIC_METRICS = { + "none": "AiCoreNone", + "pipe_utilization": "PipeUtilization", + "arithmetic_utilization": "ArithmeticUtilization", + "memory": "Memory", + "memory_l0": "MemoryL0", + "memory_ub": "MemoryUB", + "l2_cache": "L2Cache", + "memory_access": "MemoryAccess", + "resource_conflict_ratio": "ResourceConflictRatio", +} +_NPU_HOST_SYS = { + "cpu": "CPU", + "mem": "MEM", + "disk": "DISK", + "network": "NETWORK", + "osrt": "OSRT", +} +_NPU_BACKEND_OPTIONS = { + "data_simplification", + "host_sys", + "sys_io", + "sys_interconnection", + "gc_detect_threshold", +} @dataclass @@ -43,6 +74,9 @@ class ProfilerConfig: with_modules: bool = False activities: str = "auto" rank_mode: str = "all" + level: str = "level0" + aic_metrics: str = "auto" + backend_options: Optional[dict[str, Any]] = None deprecated_alias_present: bool = False legacy_defaults_enabled: bool = False explicit_profile_kwargs: Optional[set[str]] = None @@ -79,6 +113,9 @@ def from_args(cls, args: Any) -> "ProfilerConfig": with_modules=getattr(args, "profiler_with_modules", False), activities=getattr(args, "profiler_activities", "auto"), rank_mode=getattr(args, "profiler_rank_mode", "all"), + level=getattr(args, "profiler_level", "level0"), + aic_metrics=getattr(args, "profiler_aic_metrics", "auto"), + backend_options=_parse_backend_options(getattr(args, "profiler_backend_options", None)), deprecated_alias_present=enable_torch_profiler, legacy_defaults_enabled=legacy_defaults_enabled, explicit_profile_kwargs=explicit_profile_kwargs, @@ -93,13 +130,85 @@ def validate(self, backend_name: Optional[str] = None) -> None: _validate_int_option("profiler_warmup_steps", self.warmup_steps, min_value=0) _validate_int_option("profiler_active_steps", self.active_steps, min_value=1) _validate_int_option("profiler_repeat", self.repeat, min_value=0) - if self.activities not in _SUPPORTED_ACTIVITIES: - raise ValueError(f"`profiler_activities` must be one of {sorted(_SUPPORTED_ACTIVITIES)}.") - if self.rank_mode not in _SUPPORTED_RANK_MODES: - raise ValueError(f"`profiler_rank_mode` must be one of {sorted(_SUPPORTED_RANK_MODES)}.") + _validate_choice_option("profiler_activities", self.activities, _SUPPORTED_ACTIVITIES) + _validate_choice_option("profiler_rank_mode", self.rank_mode, _SUPPORTED_RANK_MODES) + if backend_name == "npu" and self.activities != "cpu": + self.npu_profiler_level_name() + self.npu_aic_metrics_name() + self.npu_backend_options() self.schedule_kwargs() + def npu_profiler_level_name(self) -> str: + key = _validate_choice_option("profiler_level", self.level, _NPU_PROFILER_LEVELS) + return _NPU_PROFILER_LEVELS[key] + + def npu_aic_metrics_name(self) -> str: + key = _validate_choice_option("profiler_aic_metrics", self.aic_metrics, {"auto", *_NPU_AIC_METRICS}) + if key == "auto": + if self.npu_profiler_level_name() in ("Level1", "Level2"): + return "PipeUtilization" + return "AiCoreNone" + + metric_name = _NPU_AIC_METRICS[key] + if metric_name != "AiCoreNone" and self.npu_profiler_level_name() not in ("Level1", "Level2"): + raise ValueError( + "`profiler_aic_metrics` requires `profiler_level` to be `level1` or `level2` on NPU." + ) + + return metric_name + + def npu_backend_options(self) -> dict[str, Any]: + if self.backend_options is None: + return {} + + unsupported_backends = set(self.backend_options) - {"npu"} + if unsupported_backends: + raise ValueError(f"`profiler_backend_options` only supports `npu`, got {sorted(unsupported_backends)}.") + + if "npu" not in self.backend_options or self.backend_options["npu"] is None: + return {} + + npu_options = self.backend_options["npu"] + if not isinstance(npu_options, dict): + raise ValueError("`profiler_backend_options.npu` must be a mapping.") + + unsupported_options = set(npu_options) - _NPU_BACKEND_OPTIONS + if unsupported_options: + raise ValueError( + f"`profiler_backend_options.npu` only supports {sorted(_NPU_BACKEND_OPTIONS)}, " + f"got {sorted(unsupported_options)}." + ) + + normalized_options = dict(npu_options) + if "data_simplification" in normalized_options and not isinstance( + normalized_options["data_simplification"], bool + ): + raise ValueError("`profiler_backend_options.npu.data_simplification` must be a boolean.") + if "sys_io" in normalized_options and not isinstance(normalized_options["sys_io"], bool): + raise ValueError("`profiler_backend_options.npu.sys_io` must be a boolean.") + if "sys_interconnection" in normalized_options and not isinstance( + normalized_options["sys_interconnection"], bool + ): + raise ValueError("`profiler_backend_options.npu.sys_interconnection` must be a boolean.") + if normalized_options.get("gc_detect_threshold") is not None: + threshold = normalized_options["gc_detect_threshold"] + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)) or threshold < 0: + raise ValueError("`profiler_backend_options.npu.gc_detect_threshold` must be null or >= 0.") + + normalized_options["host_sys"] = _normalize_host_sys(normalized_options.get("host_sys", [])) + return normalized_options + + def explicit_npu_backend_options(self) -> set[str]: + if self.backend_options is None: + return set() + + npu_options = self.backend_options.get("npu") + if not isinstance(npu_options, dict): + return set() + + return set(npu_options) + def schedule_kwargs(self) -> dict[str, int]: return dict( wait=self.wait_steps, @@ -145,27 +254,40 @@ def build_experimental_config(self, config: ProfilerConfig, logger: Any = None) return None profiler = self.profiler_module + npu_options = config.npu_backend_options() config_kwargs = dict( export_type=[profiler.ExportType.Text], - profiler_level=_get_profiler_enum_value(profiler, "ProfilerLevel", "Level0", "`profiler_level`"), - aic_metrics=_get_profiler_enum_value(profiler, "AiCMetrics", "AiCoreNone", "`profiler_aic_metrics`"), + profiler_level=_get_profiler_enum_value( + profiler, "ProfilerLevel", config.npu_profiler_level_name(), "`profiler_level`" + ), + aic_metrics=_get_profiler_enum_value( + profiler, "AiCMetrics", config.npu_aic_metrics_name(), "`profiler_aic_metrics`" + ), l2_cache=False, op_attr=False, - data_simplification=True, + data_simplification=npu_options.get("data_simplification", True), record_op_args=False, - gc_detect_threshold=None, - host_sys=[], - sys_io=False, - sys_interconnection=False, + gc_detect_threshold=npu_options.get("gc_detect_threshold"), + host_sys=[ + _get_profiler_enum_value( + profiler, "HostSystem", item, "`profiler_backend_options.npu.host_sys`" + ) + for item in npu_options.get("host_sys", []) + ], + sys_io=npu_options.get("sys_io", False), + sys_interconnection=npu_options.get("sys_interconnection", False), mstx=False, mstx_domain_include=[], mstx_domain_exclude=[], ) + explicit_keys = {"profiler_level", "aic_metrics"} + explicit_keys.update(config.explicit_npu_backend_options()) return profiler._ExperimentalConfig( **_filter_supported_kwargs( profiler._ExperimentalConfig, config_kwargs, logger=logger, + explicit_keys=explicit_keys, ) ) @@ -187,6 +309,16 @@ def _get_current_accelerator_type() -> str: return "cpu" +def _parse_backend_options(options: Any) -> Optional[dict[str, Any]]: + if options is None: + return None + + if not isinstance(options, dict): + raise ValueError("`profiler_backend_options` must be a mapping.") + + return options + + def _resolve_optional_bool(value: Any, default: bool) -> bool: if value is None: return default @@ -202,6 +334,13 @@ def _validate_int_option(name: str, value: Any, min_value: int) -> None: raise ValueError(f"`{name}` must be an integer greater than or equal to {min_value}.") +def _validate_choice_option(name: str, value: Any, supported: set[str] | dict[str, Any]) -> str: + if not isinstance(value, str) or value not in supported: + raise ValueError(f"`{name}` must be one of {sorted(supported)}.") + + return value + + def _get_explicit_profile_kwargs(**kwargs: Any) -> set[str]: name_mapping = { "record_shapes": "record_shapes", @@ -211,6 +350,21 @@ def _get_explicit_profile_kwargs(**kwargs: Any) -> set[str]: return {profile_key for arg_key, profile_key in name_mapping.items() if kwargs[arg_key] is not None} +def _normalize_host_sys(host_sys: Any) -> list[str]: + if host_sys is None: + return [] + + if not isinstance(host_sys, list): + raise ValueError("`profiler_backend_options.npu.host_sys` must be a list.") + + normalized = [] + for item in host_sys: + key = _validate_choice_option("profiler_backend_options.npu.host_sys", item, _NPU_HOST_SYS) + normalized.append(_NPU_HOST_SYS[key]) + + return normalized + + def _get_backend() -> _ProfilerBackend: accelerator_type = _get_current_accelerator_type() if accelerator_type == "cpu": diff --git a/src/llamafactory/hparams/training_args.py b/src/llamafactory/hparams/training_args.py index 6eee680b30..52bab007c1 100644 --- a/src/llamafactory/hparams/training_args.py +++ b/src/llamafactory/hparams/training_args.py @@ -122,11 +122,33 @@ class ProfilerArguments: ) profiler_activities: str = field( default="auto", - metadata={"help": "Profiler activities to collect: auto, all, cpu, or device."}, + metadata={"help": "Profiler activities to collect. Choices: auto, all, cpu, device."}, ) profiler_rank_mode: str = field( default="all", - metadata={"help": "Profiler rank collection mode: all or rank0."}, + metadata={"help": "Profiler rank collection mode. Choices: all, rank0."}, + ) + profiler_level: str = field( + default="level0", + metadata={"help": "NPU profiler collection level. Choices: none, level0, level1, level2."}, + ) + profiler_aic_metrics: str = field( + default="auto", + metadata={ + "help": ( + "NPU AI Core metric: auto, none, pipe_utilization, arithmetic_utilization, memory, " + "memory_l0, memory_ub, l2_cache, memory_access, or resource_conflict_ratio." + ) + }, + ) + profiler_backend_options: Optional[dict] = field( + default=None, + metadata={ + "help": ( + "Backend-specific profiler mapping. Currently supports npu.data_simplification, npu.host_sys, " + "npu.sys_io, npu.sys_interconnection, and npu.gc_detect_threshold." + ) + }, ) profile_modules: Optional[str] = field( default=None, diff --git a/src/llamafactory/train/callbacks.py b/src/llamafactory/train/callbacks.py index 4d8bcd7eef..8d1135d669 100644 --- a/src/llamafactory/train/callbacks.py +++ b/src/llamafactory/train/callbacks.py @@ -357,8 +357,12 @@ class TorchProfilerCallback(TrainerCallback): profiler_record_shapes – record tensor shapes (default: false; true for deprecated alias) profiler_profile_memory – profile memory usage (default: false; true for deprecated alias) profiler_with_stack – record stack traces (default: false; true for deprecated alias) - profiler_activities – auto, all, cpu, or device (default: auto) - profiler_rank_mode – all ranks or rank0 only (default: all) + profiler_activities – choices: auto, all, cpu, device (default: auto) + profiler_rank_mode – choices: all, rank0 (default: all) + profiler_level – choices: none, level0, level1, level2 + profiler_aic_metrics – choices: auto, none, pipe_utilization, arithmetic_utilization, + memory, memory_l0, memory_ub, l2_cache, memory_access, + resource_conflict_ratio Trace files (one per rank, Chrome / TensorBoard JSON format) are written to ``/rank_/``. @@ -369,6 +373,13 @@ class TorchProfilerCallback(TrainerCallback): profiler_warmup_steps: 1 profiler_active_steps: 3 profiler_rank_mode: rank0 + profiler_level: level1 + profiler_aic_metrics: pipe_utilization + profiler_backend_options: + npu: + host_sys: [cpu, mem] + sys_io: true + data_simplification: false """ def __init__(self, training_args: "TrainingArguments") -> None: diff --git a/src/llamafactory/v1/config/training_args.py b/src/llamafactory/v1/config/training_args.py index 4cfd75c83a..32af576cb9 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/src/llamafactory/v1/config/training_args.py @@ -123,11 +123,33 @@ class TrainingArguments: ) profiler_activities: str = field( default="auto", - metadata={"help": "Profiler activities to collect: auto, all, cpu, or device."}, + metadata={"help": "Profiler activities to collect. Choices: auto, all, cpu, device."}, ) profiler_rank_mode: str = field( default="all", - metadata={"help": "Profiler rank collection mode: all or rank0."}, + metadata={"help": "Profiler rank collection mode. Choices: all, rank0."}, + ) + profiler_level: str = field( + default="level0", + metadata={"help": "NPU profiler collection level. Choices: none, level0, level1, level2."}, + ) + profiler_aic_metrics: str = field( + default="auto", + metadata={ + "help": ( + "NPU AI Core metric: auto, none, pipe_utilization, arithmetic_utilization, memory, " + "memory_l0, memory_ub, l2_cache, memory_access, or resource_conflict_ratio." + ) + }, + ) + profiler_backend_options: dict | None = field( + default=None, + metadata={ + "help": ( + "Backend-specific profiler mapping. Currently supports npu.data_simplification, npu.host_sys, " + "npu.sys_io, npu.sys_interconnection, and npu.gc_detect_threshold." + ) + }, ) dist_config: PluginConfig | None = field( default=None, From 367f14e2fd35274027a6ad86bc6baafce35eee1d Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Sat, 20 Jun 2026 17:25:27 +0000 Subject: [PATCH 3/9] document profiler training arguments --- docs/en/hyperparameters/training-argument.md | 39 ++++++++++++++++++- docs/zh/hyperparameters/training-argument.md | 40 ++++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/docs/en/hyperparameters/training-argument.md b/docs/en/hyperparameters/training-argument.md index 2a9bea5d85..1af9d0c19b 100644 --- a/docs/en/hyperparameters/training-argument.md +++ b/docs/en/hyperparameters/training-argument.md @@ -1,3 +1,40 @@ # Training Argument -This page is not yet available in English. Use the language switcher to view Simplified Chinese. +## Profiler + +LLaMA Factory can collect PyTorch profiler traces during training on CPU, CUDA GPU, and Ascend NPU devices. Enable it in the training YAML: + +```yaml +enable_profiler: true +profiler_output_dir: ./saves/profile +profiler_skip_first: 8 +profiler_wait_steps: 0 +profiler_warmup_steps: 1 +profiler_active_steps: 3 +profiler_repeat: 1 +profiler_rank_mode: rank0 +``` + +The schedule follows the official `torch.profiler.schedule` / `torch_npu.profiler.schedule` semantics: first skip `profiler_skip_first` steps, then each cycle runs `wait -> warmup -> active`, repeated by `profiler_repeat`. The callback calls `prof.step()` once after each optimizer step. + +For Ascend NPU, use `profiler_level` and `profiler_aic_metrics` to control collection depth: + +```yaml +profiler_level: level1 +profiler_aic_metrics: pipe_utilization +profiler_backend_options: + npu: + data_simplification: true + host_sys: [cpu, mem] + sys_io: false + sys_interconnection: false +``` + +String enum values are fixed short values. `profiler_activities` supports `auto`, `all`, `cpu`, `device`; `profiler_rank_mode` supports `all`, `rank0`; `profiler_level` supports `none`, `level0`, `level1`, `level2`; `profiler_aic_metrics` supports `auto`, `none`, `pipe_utilization`, `arithmetic_utilization`, `memory`, `memory_l0`, `memory_ub`, `l2_cache`, `memory_access`, `resource_conflict_ratio`; `profiler_backend_options.npu.host_sys` supports `cpu`, `mem`, `disk`, `network`, `osrt`. Values outside these lists fail validation. `profiler_backend_options` must be a YAML mapping, not a JSON string. + +Do not enable Ascend `dynamic_profile` through `PROF_CONFIG_PATH` at the same time as `enable_profiler`. + +Official references: + +- [PyTorch Profiler](https://docs.pytorch.org/docs/stable/profiler.html) +- [Ascend CANN 9.0 PyTorch Profiler](https://www.hiascend.com/document/detail/zh/canncommercial/900/devaids/Profiling/atlasprofiling_16_0033.html) diff --git a/docs/zh/hyperparameters/training-argument.md b/docs/zh/hyperparameters/training-argument.md index e69de29bb2..8d5e1dbb1a 100644 --- a/docs/zh/hyperparameters/training-argument.md +++ b/docs/zh/hyperparameters/training-argument.md @@ -0,0 +1,40 @@ +# 训练参数 + +## Profiler + +LLaMA Factory 支持在训练过程中采集 PyTorch profiler trace，覆盖 CPU、CUDA GPU 和昇腾 NPU。可以在训练 YAML 中开启： + +```yaml +enable_profiler: true +profiler_output_dir: ./saves/profile +profiler_skip_first: 8 +profiler_wait_steps: 0 +profiler_warmup_steps: 1 +profiler_active_steps: 3 +profiler_repeat: 1 +profiler_rank_mode: rank0 +``` + +调度语义对齐官方 `torch.profiler.schedule` / `torch_npu.profiler.schedule`：先跳过 `profiler_skip_first` 个 step，之后每个周期按 `wait -> warmup -> active` 执行，并由 `profiler_repeat` 控制重复次数。callback 在每个 optimizer step 结束后调用一次 `prof.step()`。 + +昇腾 NPU 上可以通过 `profiler_level` 和 `profiler_aic_metrics` 控制采集深度： + +```yaml +profiler_level: level1 +profiler_aic_metrics: pipe_utilization +profiler_backend_options: + npu: + data_simplification: true + host_sys: [cpu, mem] + sys_io: false + sys_interconnection: false +``` + +字符串枚举参数固定为短值：`profiler_activities` 支持 `auto`、`all`、`cpu`、`device`；`profiler_rank_mode` 支持 `all`、`rank0`；`profiler_level` 支持 `none`、`level0`、`level1`、`level2`；`profiler_aic_metrics` 支持 `auto`、`none`、`pipe_utilization`、`arithmetic_utilization`、`memory`、`memory_l0`、`memory_ub`、`l2_cache`、`memory_access`、`resource_conflict_ratio`；`profiler_backend_options.npu.host_sys` 支持 `cpu`、`mem`、`disk`、`network`、`osrt`。传入不在列表内的值会直接报错。`profiler_backend_options` 必须是 YAML mapping，不支持 JSON 字符串。 + +不要同时通过 `PROF_CONFIG_PATH` 开启昇腾 `dynamic_profile` 和 LLaMA Factory 的 `enable_profiler`。 + +官方参考： + +- [PyTorch Profiler](https://docs.pytorch.org/docs/stable/profiler.html) +- [昇腾 CANN 9.0 PyTorch Profiler](https://www.hiascend.com/document/detail/zh/canncommercial/900/devaids/Profiling/atlasprofiling_16_0033.html) From e863a2ebb20d661f4b08447167838ddbe0d14f61 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Sun, 21 Jun 2026 08:52:46 +0000 Subject: [PATCH 4/9] reuse torch profiler enable flag --- docs/en/hyperparameters/training-argument.md | 4 +-- docs/zh/hyperparameters/training-argument.md | 4 +-- src/llamafactory/extras/profiler.py | 27 +++++--------------- src/llamafactory/hparams/training_args.py | 6 +---- src/llamafactory/train/callbacks.py | 9 +++---- src/llamafactory/train/tuner.py | 2 +- src/llamafactory/v1/config/training_args.py | 6 +---- src/llamafactory/v1/core/base_trainer.py | 2 +- 8 files changed, 18 insertions(+), 42 deletions(-) diff --git a/docs/en/hyperparameters/training-argument.md b/docs/en/hyperparameters/training-argument.md index 1af9d0c19b..70546dc307 100644 --- a/docs/en/hyperparameters/training-argument.md +++ b/docs/en/hyperparameters/training-argument.md @@ -5,7 +5,7 @@ LLaMA Factory can collect PyTorch profiler traces during training on CPU, CUDA GPU, and Ascend NPU devices. Enable it in the training YAML: ```yaml -enable_profiler: true +enable_torch_profiler: true profiler_output_dir: ./saves/profile profiler_skip_first: 8 profiler_wait_steps: 0 @@ -32,7 +32,7 @@ profiler_backend_options: String enum values are fixed short values. `profiler_activities` supports `auto`, `all`, `cpu`, `device`; `profiler_rank_mode` supports `all`, `rank0`; `profiler_level` supports `none`, `level0`, `level1`, `level2`; `profiler_aic_metrics` supports `auto`, `none`, `pipe_utilization`, `arithmetic_utilization`, `memory`, `memory_l0`, `memory_ub`, `l2_cache`, `memory_access`, `resource_conflict_ratio`; `profiler_backend_options.npu.host_sys` supports `cpu`, `mem`, `disk`, `network`, `osrt`. Values outside these lists fail validation. `profiler_backend_options` must be a YAML mapping, not a JSON string. -Do not enable Ascend `dynamic_profile` through `PROF_CONFIG_PATH` at the same time as `enable_profiler`. +Do not enable Ascend `dynamic_profile` through `PROF_CONFIG_PATH` at the same time as `enable_torch_profiler`. Official references: diff --git a/docs/zh/hyperparameters/training-argument.md b/docs/zh/hyperparameters/training-argument.md index 8d5e1dbb1a..394c30ab91 100644 --- a/docs/zh/hyperparameters/training-argument.md +++ b/docs/zh/hyperparameters/training-argument.md @@ -5,7 +5,7 @@ LLaMA Factory 支持在训练过程中采集 PyTorch profiler trace，覆盖 CPU、CUDA GPU 和昇腾 NPU。可以在训练 YAML 中开启： ```yaml -enable_profiler: true +enable_torch_profiler: true profiler_output_dir: ./saves/profile profiler_skip_first: 8 profiler_wait_steps: 0 @@ -32,7 +32,7 @@ profiler_backend_options: 字符串枚举参数固定为短值：`profiler_activities` 支持 `auto`、`all`、`cpu`、`device`；`profiler_rank_mode` 支持 `all`、`rank0`；`profiler_level` 支持 `none`、`level0`、`level1`、`level2`；`profiler_aic_metrics` 支持 `auto`、`none`、`pipe_utilization`、`arithmetic_utilization`、`memory`、`memory_l0`、`memory_ub`、`l2_cache`、`memory_access`、`resource_conflict_ratio`；`profiler_backend_options.npu.host_sys` 支持 `cpu`、`mem`、`disk`、`network`、`osrt`。传入不在列表内的值会直接报错。`profiler_backend_options` 必须是 YAML mapping，不支持 JSON 字符串。 -不要同时通过 `PROF_CONFIG_PATH` 开启昇腾 `dynamic_profile` 和 LLaMA Factory 的 `enable_profiler`。 +不要同时通过 `PROF_CONFIG_PATH` 开启昇腾 `dynamic_profile` 和 LLaMA Factory 的 `enable_torch_profiler`。 官方参考： diff --git a/src/llamafactory/extras/profiler.py b/src/llamafactory/extras/profiler.py index 186f28574b..1bbfdf0e42 100644 --- a/src/llamafactory/extras/profiler.py +++ b/src/llamafactory/extras/profiler.py @@ -77,15 +77,11 @@ class ProfilerConfig: level: str = "level0" aic_metrics: str = "auto" backend_options: Optional[dict[str, Any]] = None - deprecated_alias_present: bool = False - legacy_defaults_enabled: bool = False explicit_profile_kwargs: Optional[set[str]] = None @classmethod def from_args(cls, args: Any) -> "ProfilerConfig": - enable_profiler = bool(getattr(args, "enable_profiler", False)) enable_torch_profiler = bool(getattr(args, "enable_torch_profiler", False)) - legacy_defaults_enabled = enable_torch_profiler and not enable_profiler record_shapes = getattr(args, "profiler_record_shapes", None) profile_memory = getattr(args, "profiler_profile_memory", None) with_stack = getattr(args, "profiler_with_stack", None) @@ -99,16 +95,16 @@ def from_args(cls, args: Any) -> "ProfilerConfig": if getattr(args, "profiler_with_modules", False): explicit_profile_kwargs.add("with_modules") return cls( - enabled=enable_profiler or enable_torch_profiler, + enabled=enable_torch_profiler, output_dir=getattr(args, "profiler_output_dir", None), skip_first=getattr(args, "profiler_skip_first", 0), wait_steps=getattr(args, "profiler_wait_steps", 1), warmup_steps=getattr(args, "profiler_warmup_steps", 1), active_steps=getattr(args, "profiler_active_steps", 1), repeat=getattr(args, "profiler_repeat", 1), - record_shapes=_resolve_optional_bool(record_shapes, legacy_defaults_enabled), - profile_memory=_resolve_optional_bool(profile_memory, legacy_defaults_enabled), - with_stack=_resolve_optional_bool(with_stack, legacy_defaults_enabled), + record_shapes=_resolve_optional_bool(record_shapes, enable_torch_profiler), + profile_memory=_resolve_optional_bool(profile_memory, enable_torch_profiler), + with_stack=_resolve_optional_bool(with_stack, enable_torch_profiler), with_flops=getattr(args, "profiler_with_flops", False), with_modules=getattr(args, "profiler_with_modules", False), activities=getattr(args, "profiler_activities", "auto"), @@ -116,8 +112,6 @@ def from_args(cls, args: Any) -> "ProfilerConfig": level=getattr(args, "profiler_level", "level0"), aic_metrics=getattr(args, "profiler_aic_metrics", "auto"), backend_options=_parse_backend_options(getattr(args, "profiler_backend_options", None)), - deprecated_alias_present=enable_torch_profiler, - legacy_defaults_enabled=legacy_defaults_enabled, explicit_profile_kwargs=explicit_profile_kwargs, ) @@ -451,16 +445,6 @@ def start(self, output_dir: str, initial_step: int = 0) -> None: if not self.config.enabled: return - if self.config.deprecated_alias_present: - message = "`enable_torch_profiler` is deprecated; use `enable_profiler` instead." - if self.config.legacy_defaults_enabled: - message += " Legacy shape, memory, and stack profiling defaults are preserved for this alias." - _log( - self.logger, - "warning", - message, - ) - backend = _get_backend() self.backend_name = backend.name self.config.validate(backend.name) @@ -468,7 +452,8 @@ def start(self, output_dir: str, initial_step: int = 0) -> None: _log( self.logger, "warning", - "`PROF_CONFIG_PATH` is set. Do not enable dynamic_profile together with `enable_profiler`.", + "`PROF_CONFIG_PATH` is set. Do not enable dynamic_profile together with " + "`enable_torch_profiler`.", ) schedule = self.config.schedule_kwargs() diff --git a/src/llamafactory/hparams/training_args.py b/src/llamafactory/hparams/training_args.py index 52bab007c1..ed566111e3 100644 --- a/src/llamafactory/hparams/training_args.py +++ b/src/llamafactory/hparams/training_args.py @@ -68,13 +68,9 @@ def __post_init__(self): class ProfilerArguments: r"""Arguments for torch profiler configuration.""" - enable_profiler: bool = field( - default=False, - metadata={"help": "Whether to enable profiler for collecting CPU/CUDA/NPU performance traces."}, - ) enable_torch_profiler: bool = field( default=False, - metadata={"help": "Deprecated. Use `enable_profiler` instead."}, + metadata={"help": "Whether to enable torch profiler for collecting CPU/CUDA/NPU performance traces."}, ) profiler_output_dir: Optional[str] = field( default=None, diff --git a/src/llamafactory/train/callbacks.py b/src/llamafactory/train/callbacks.py index 8d1135d669..cae67e8d5b 100644 --- a/src/llamafactory/train/callbacks.py +++ b/src/llamafactory/train/callbacks.py @@ -344,8 +344,7 @@ def on_prediction_step( class TorchProfilerCallback(TrainerCallback): r"""A callback for collecting CPU/CUDA/NPU profiler traces during training. - Activated by setting ``enable_profiler: true`` in the YAML config. - ``enable_torch_profiler`` is still accepted for backward compatibility. + Activated by setting ``enable_torch_profiler: true`` in the YAML config. Configuration fields (in YAML): profiler_output_dir – where to write traces (default: /profiler) @@ -354,9 +353,9 @@ class TorchProfilerCallback(TrainerCallback): profiler_warmup_steps – profiler warm-up steps per cycle (default: 1) profiler_active_steps – steps to record per cycle (default: 1) profiler_repeat – number of cycles; 0 = forever (default: 1) - profiler_record_shapes – record tensor shapes (default: false; true for deprecated alias) - profiler_profile_memory – profile memory usage (default: false; true for deprecated alias) - profiler_with_stack – record stack traces (default: false; true for deprecated alias) + profiler_record_shapes – record tensor shapes (default: true) + profiler_profile_memory – profile memory usage (default: true) + profiler_with_stack – record stack traces (default: true) profiler_activities – choices: auto, all, cpu, device (default: auto) profiler_rank_mode – choices: all, rank0 (default: all) profiler_level – choices: none, level0, level1, level2 diff --git a/src/llamafactory/train/tuner.py b/src/llamafactory/train/tuner.py index 7421a1ed10..22f7dbf87a 100644 --- a/src/llamafactory/train/tuner.py +++ b/src/llamafactory/train/tuner.py @@ -80,7 +80,7 @@ def _training_function(config: dict[str, Any]) -> None: if finetuning_args.early_stopping_steps is not None: callbacks.append(EarlyStoppingCallback(early_stopping_patience=finetuning_args.early_stopping_steps)) - if getattr(training_args, "enable_profiler", False) or getattr(training_args, "enable_torch_profiler", False): + if getattr(training_args, "enable_torch_profiler", False): callbacks.append(TorchProfilerCallback(training_args)) if getattr(training_args, "profile_modules", None): diff --git a/src/llamafactory/v1/config/training_args.py b/src/llamafactory/v1/config/training_args.py index 32af576cb9..9963d861f4 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/src/llamafactory/v1/config/training_args.py @@ -69,13 +69,9 @@ class TrainingArguments: default=True, metadata={"help": "Enable activation checkpointing for training."}, ) - enable_profiler: bool = field( - default=False, - metadata={"help": "Whether to enable profiler for collecting CPU/CUDA/NPU performance traces."}, - ) enable_torch_profiler: bool = field( default=False, - metadata={"help": "Deprecated. Use `enable_profiler` instead."}, + metadata={"help": "Whether to enable torch profiler for collecting CPU/CUDA/NPU performance traces."}, ) profiler_output_dir: str | None = field( default=None, diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index b08dbd6d5b..e45fdcb681 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -125,7 +125,7 @@ def __init__( # Callbacks self.callback_handler = CallbackHandler([LoggingCallback()], trainer=self) - if getattr(self.args, "enable_profiler", False) or getattr(self.args, "enable_torch_profiler", False): + if getattr(self.args, "enable_torch_profiler", False): from ..accelerator.profiler import ProfilerCallback self.callback_handler.add_callback(ProfilerCallback(self.args)) From b1844bedc1f5778533010fb2911a2b94e28f13f5 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Mon, 22 Jun 2026 02:14:37 +0000 Subject: [PATCH 5/9] isolate profiler support to v1 --- docs/en/hyperparameters/training-argument.md | 39 +- docs/zh/hyperparameters/training-argument.md | 40 -- src/llamafactory/extras/profiler.py | 517 ------------------ src/llamafactory/hparams/training_args.py | 58 +- src/llamafactory/train/callbacks.py | 81 ++- src/llamafactory/v1/accelerator/profiler.py | 526 ++++++++++++++++++- 6 files changed, 582 insertions(+), 679 deletions(-) delete mode 100644 src/llamafactory/extras/profiler.py diff --git a/docs/en/hyperparameters/training-argument.md b/docs/en/hyperparameters/training-argument.md index 70546dc307..2a9bea5d85 100644 --- a/docs/en/hyperparameters/training-argument.md +++ b/docs/en/hyperparameters/training-argument.md @@ -1,40 +1,3 @@ # Training Argument -## Profiler - -LLaMA Factory can collect PyTorch profiler traces during training on CPU, CUDA GPU, and Ascend NPU devices. Enable it in the training YAML: - -```yaml -enable_torch_profiler: true -profiler_output_dir: ./saves/profile -profiler_skip_first: 8 -profiler_wait_steps: 0 -profiler_warmup_steps: 1 -profiler_active_steps: 3 -profiler_repeat: 1 -profiler_rank_mode: rank0 -``` - -The schedule follows the official `torch.profiler.schedule` / `torch_npu.profiler.schedule` semantics: first skip `profiler_skip_first` steps, then each cycle runs `wait -> warmup -> active`, repeated by `profiler_repeat`. The callback calls `prof.step()` once after each optimizer step. - -For Ascend NPU, use `profiler_level` and `profiler_aic_metrics` to control collection depth: - -```yaml -profiler_level: level1 -profiler_aic_metrics: pipe_utilization -profiler_backend_options: - npu: - data_simplification: true - host_sys: [cpu, mem] - sys_io: false - sys_interconnection: false -``` - -String enum values are fixed short values. `profiler_activities` supports `auto`, `all`, `cpu`, `device`; `profiler_rank_mode` supports `all`, `rank0`; `profiler_level` supports `none`, `level0`, `level1`, `level2`; `profiler_aic_metrics` supports `auto`, `none`, `pipe_utilization`, `arithmetic_utilization`, `memory`, `memory_l0`, `memory_ub`, `l2_cache`, `memory_access`, `resource_conflict_ratio`; `profiler_backend_options.npu.host_sys` supports `cpu`, `mem`, `disk`, `network`, `osrt`. Values outside these lists fail validation. `profiler_backend_options` must be a YAML mapping, not a JSON string. - -Do not enable Ascend `dynamic_profile` through `PROF_CONFIG_PATH` at the same time as `enable_torch_profiler`. - -Official references: - -- [PyTorch Profiler](https://docs.pytorch.org/docs/stable/profiler.html) -- [Ascend CANN 9.0 PyTorch Profiler](https://www.hiascend.com/document/detail/zh/canncommercial/900/devaids/Profiling/atlasprofiling_16_0033.html) +This page is not yet available in English. Use the language switcher to view Simplified Chinese. diff --git a/docs/zh/hyperparameters/training-argument.md b/docs/zh/hyperparameters/training-argument.md index 394c30ab91..e69de29bb2 100644 --- a/docs/zh/hyperparameters/training-argument.md +++ b/docs/zh/hyperparameters/training-argument.md @@ -1,40 +0,0 @@ -# 训练参数 - -## Profiler - -LLaMA Factory 支持在训练过程中采集 PyTorch profiler trace，覆盖 CPU、CUDA GPU 和昇腾 NPU。可以在训练 YAML 中开启： - -```yaml -enable_torch_profiler: true -profiler_output_dir: ./saves/profile -profiler_skip_first: 8 -profiler_wait_steps: 0 -profiler_warmup_steps: 1 -profiler_active_steps: 3 -profiler_repeat: 1 -profiler_rank_mode: rank0 -``` - -调度语义对齐官方 `torch.profiler.schedule` / `torch_npu.profiler.schedule`：先跳过 `profiler_skip_first` 个 step，之后每个周期按 `wait -> warmup -> active` 执行，并由 `profiler_repeat` 控制重复次数。callback 在每个 optimizer step 结束后调用一次 `prof.step()`。 - -昇腾 NPU 上可以通过 `profiler_level` 和 `profiler_aic_metrics` 控制采集深度： - -```yaml -profiler_level: level1 -profiler_aic_metrics: pipe_utilization -profiler_backend_options: - npu: - data_simplification: true - host_sys: [cpu, mem] - sys_io: false - sys_interconnection: false -``` - -字符串枚举参数固定为短值：`profiler_activities` 支持 `auto`、`all`、`cpu`、`device`；`profiler_rank_mode` 支持 `all`、`rank0`；`profiler_level` 支持 `none`、`level0`、`level1`、`level2`；`profiler_aic_metrics` 支持 `auto`、`none`、`pipe_utilization`、`arithmetic_utilization`、`memory`、`memory_l0`、`memory_ub`、`l2_cache`、`memory_access`、`resource_conflict_ratio`；`profiler_backend_options.npu.host_sys` 支持 `cpu`、`mem`、`disk`、`network`、`osrt`。传入不在列表内的值会直接报错。`profiler_backend_options` 必须是 YAML mapping，不支持 JSON 字符串。 - -不要同时通过 `PROF_CONFIG_PATH` 开启昇腾 `dynamic_profile` 和 LLaMA Factory 的 `enable_torch_profiler`。 - -官方参考： - -- [PyTorch Profiler](https://docs.pytorch.org/docs/stable/profiler.html) -- [昇腾 CANN 9.0 PyTorch Profiler](https://www.hiascend.com/document/detail/zh/canncommercial/900/devaids/Profiling/atlasprofiling_16_0033.html) diff --git a/src/llamafactory/extras/profiler.py b/src/llamafactory/extras/profiler.py deleted file mode 100644 index 1bbfdf0e42..0000000000 --- a/src/llamafactory/extras/profiler.py +++ /dev/null @@ -1,517 +0,0 @@ -# Copyright 2025 the LlamaFactory team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import inspect -import os -import sys -from contextlib import suppress -from dataclasses import dataclass -from typing import Any, Optional - -import torch -from transformers.utils import is_torch_cuda_available, is_torch_npu_available - - -_SUPPORTED_ACTIVITIES = {"auto", "all", "cpu", "device"} -_SUPPORTED_RANK_MODES = {"all", "rank0"} -_NPU_PROFILER_LEVELS = { - "none": "Level_none", - "level0": "Level0", - "level1": "Level1", - "level2": "Level2", -} -_NPU_AIC_METRICS = { - "none": "AiCoreNone", - "pipe_utilization": "PipeUtilization", - "arithmetic_utilization": "ArithmeticUtilization", - "memory": "Memory", - "memory_l0": "MemoryL0", - "memory_ub": "MemoryUB", - "l2_cache": "L2Cache", - "memory_access": "MemoryAccess", - "resource_conflict_ratio": "ResourceConflictRatio", -} -_NPU_HOST_SYS = { - "cpu": "CPU", - "mem": "MEM", - "disk": "DISK", - "network": "NETWORK", - "osrt": "OSRT", -} -_NPU_BACKEND_OPTIONS = { - "data_simplification", - "host_sys", - "sys_io", - "sys_interconnection", - "gc_detect_threshold", -} - - -@dataclass -class ProfilerConfig: - enabled: bool = False - output_dir: Optional[str] = None - skip_first: int = 0 - wait_steps: int = 1 - warmup_steps: int = 1 - active_steps: int = 1 - repeat: int = 1 - record_shapes: bool = False - profile_memory: bool = False - with_stack: bool = False - with_flops: bool = False - with_modules: bool = False - activities: str = "auto" - rank_mode: str = "all" - level: str = "level0" - aic_metrics: str = "auto" - backend_options: Optional[dict[str, Any]] = None - explicit_profile_kwargs: Optional[set[str]] = None - - @classmethod - def from_args(cls, args: Any) -> "ProfilerConfig": - enable_torch_profiler = bool(getattr(args, "enable_torch_profiler", False)) - record_shapes = getattr(args, "profiler_record_shapes", None) - profile_memory = getattr(args, "profiler_profile_memory", None) - with_stack = getattr(args, "profiler_with_stack", None) - explicit_profile_kwargs = _get_explicit_profile_kwargs( - record_shapes=record_shapes, - profile_memory=profile_memory, - with_stack=with_stack, - ) - if getattr(args, "profiler_with_flops", False): - explicit_profile_kwargs.add("with_flops") - if getattr(args, "profiler_with_modules", False): - explicit_profile_kwargs.add("with_modules") - return cls( - enabled=enable_torch_profiler, - output_dir=getattr(args, "profiler_output_dir", None), - skip_first=getattr(args, "profiler_skip_first", 0), - wait_steps=getattr(args, "profiler_wait_steps", 1), - warmup_steps=getattr(args, "profiler_warmup_steps", 1), - active_steps=getattr(args, "profiler_active_steps", 1), - repeat=getattr(args, "profiler_repeat", 1), - record_shapes=_resolve_optional_bool(record_shapes, enable_torch_profiler), - profile_memory=_resolve_optional_bool(profile_memory, enable_torch_profiler), - with_stack=_resolve_optional_bool(with_stack, enable_torch_profiler), - with_flops=getattr(args, "profiler_with_flops", False), - with_modules=getattr(args, "profiler_with_modules", False), - activities=getattr(args, "profiler_activities", "auto"), - rank_mode=getattr(args, "profiler_rank_mode", "all"), - level=getattr(args, "profiler_level", "level0"), - aic_metrics=getattr(args, "profiler_aic_metrics", "auto"), - backend_options=_parse_backend_options(getattr(args, "profiler_backend_options", None)), - explicit_profile_kwargs=explicit_profile_kwargs, - ) - - def validate(self, backend_name: Optional[str] = None) -> None: - if not self.enabled: - return - - _validate_int_option("profiler_skip_first", self.skip_first, min_value=0) - _validate_int_option("profiler_wait_steps", self.wait_steps, min_value=0) - _validate_int_option("profiler_warmup_steps", self.warmup_steps, min_value=0) - _validate_int_option("profiler_active_steps", self.active_steps, min_value=1) - _validate_int_option("profiler_repeat", self.repeat, min_value=0) - _validate_choice_option("profiler_activities", self.activities, _SUPPORTED_ACTIVITIES) - _validate_choice_option("profiler_rank_mode", self.rank_mode, _SUPPORTED_RANK_MODES) - - if backend_name == "npu" and self.activities != "cpu": - self.npu_profiler_level_name() - self.npu_aic_metrics_name() - self.npu_backend_options() - self.schedule_kwargs() - - def npu_profiler_level_name(self) -> str: - key = _validate_choice_option("profiler_level", self.level, _NPU_PROFILER_LEVELS) - return _NPU_PROFILER_LEVELS[key] - - def npu_aic_metrics_name(self) -> str: - key = _validate_choice_option("profiler_aic_metrics", self.aic_metrics, {"auto", *_NPU_AIC_METRICS}) - if key == "auto": - if self.npu_profiler_level_name() in ("Level1", "Level2"): - return "PipeUtilization" - return "AiCoreNone" - - metric_name = _NPU_AIC_METRICS[key] - if metric_name != "AiCoreNone" and self.npu_profiler_level_name() not in ("Level1", "Level2"): - raise ValueError( - "`profiler_aic_metrics` requires `profiler_level` to be `level1` or `level2` on NPU." - ) - - return metric_name - - def npu_backend_options(self) -> dict[str, Any]: - if self.backend_options is None: - return {} - - unsupported_backends = set(self.backend_options) - {"npu"} - if unsupported_backends: - raise ValueError(f"`profiler_backend_options` only supports `npu`, got {sorted(unsupported_backends)}.") - - if "npu" not in self.backend_options or self.backend_options["npu"] is None: - return {} - - npu_options = self.backend_options["npu"] - if not isinstance(npu_options, dict): - raise ValueError("`profiler_backend_options.npu` must be a mapping.") - - unsupported_options = set(npu_options) - _NPU_BACKEND_OPTIONS - if unsupported_options: - raise ValueError( - f"`profiler_backend_options.npu` only supports {sorted(_NPU_BACKEND_OPTIONS)}, " - f"got {sorted(unsupported_options)}." - ) - - normalized_options = dict(npu_options) - if "data_simplification" in normalized_options and not isinstance( - normalized_options["data_simplification"], bool - ): - raise ValueError("`profiler_backend_options.npu.data_simplification` must be a boolean.") - if "sys_io" in normalized_options and not isinstance(normalized_options["sys_io"], bool): - raise ValueError("`profiler_backend_options.npu.sys_io` must be a boolean.") - if "sys_interconnection" in normalized_options and not isinstance( - normalized_options["sys_interconnection"], bool - ): - raise ValueError("`profiler_backend_options.npu.sys_interconnection` must be a boolean.") - if normalized_options.get("gc_detect_threshold") is not None: - threshold = normalized_options["gc_detect_threshold"] - if isinstance(threshold, bool) or not isinstance(threshold, (int, float)) or threshold < 0: - raise ValueError("`profiler_backend_options.npu.gc_detect_threshold` must be null or >= 0.") - - normalized_options["host_sys"] = _normalize_host_sys(normalized_options.get("host_sys", [])) - return normalized_options - - def explicit_npu_backend_options(self) -> set[str]: - if self.backend_options is None: - return set() - - npu_options = self.backend_options.get("npu") - if not isinstance(npu_options, dict): - return set() - - return set(npu_options) - - def schedule_kwargs(self) -> dict[str, int]: - return dict( - wait=self.wait_steps, - warmup=self.warmup_steps, - active=self.active_steps, - repeat=self.repeat, - skip_first=self.skip_first, - ) - - -class _ProfilerBackend: - def __init__(self, name: str, profiler_module: Any, device_activity_name: Optional[str]) -> None: - self.name = name - self.profiler_module = profiler_module - self.device_activity_name = device_activity_name - - def build_schedule(self, kwargs: dict[str, Any]) -> Any: - if self.name == "npu": - return self.profiler_module.schedule( - wait=kwargs["wait"], - active=kwargs["active"], - warmup=kwargs["warmup"], - repeat=kwargs["repeat"], - skip_first=kwargs["skip_first"], - ) - - return self.profiler_module.schedule(**kwargs) - - def build_activities(self, config: ProfilerConfig) -> list[Any]: - activity_cls = self.profiler_module.ProfilerActivity - activities = [] - if config.activities in ("auto", "all", "cpu"): - activities.append(activity_cls.CPU) - if config.activities == "device" and self.device_activity_name is None: - raise ValueError("`profiler_activities: device` requires a CUDA or NPU backend.") - if config.activities in ("auto", "all", "device"): - if self.device_activity_name is not None: - activities.append(getattr(activity_cls, self.device_activity_name)) - return activities - - def build_experimental_config(self, config: ProfilerConfig, logger: Any = None) -> Optional[Any]: - if self.name != "npu" or config.activities == "cpu": - return None - - profiler = self.profiler_module - npu_options = config.npu_backend_options() - config_kwargs = dict( - export_type=[profiler.ExportType.Text], - profiler_level=_get_profiler_enum_value( - profiler, "ProfilerLevel", config.npu_profiler_level_name(), "`profiler_level`" - ), - aic_metrics=_get_profiler_enum_value( - profiler, "AiCMetrics", config.npu_aic_metrics_name(), "`profiler_aic_metrics`" - ), - l2_cache=False, - op_attr=False, - data_simplification=npu_options.get("data_simplification", True), - record_op_args=False, - gc_detect_threshold=npu_options.get("gc_detect_threshold"), - host_sys=[ - _get_profiler_enum_value( - profiler, "HostSystem", item, "`profiler_backend_options.npu.host_sys`" - ) - for item in npu_options.get("host_sys", []) - ], - sys_io=npu_options.get("sys_io", False), - sys_interconnection=npu_options.get("sys_interconnection", False), - mstx=False, - mstx_domain_include=[], - mstx_domain_exclude=[], - ) - explicit_keys = {"profiler_level", "aic_metrics"} - explicit_keys.update(config.explicit_npu_backend_options()) - return profiler._ExperimentalConfig( - **_filter_supported_kwargs( - profiler._ExperimentalConfig, - config_kwargs, - logger=logger, - explicit_keys=explicit_keys, - ) - ) - - -def _get_rank() -> int: - import torch.distributed as dist - - if dist.is_available() and dist.is_initialized(): - return dist.get_rank() - - return int(os.getenv("RANK", "0")) - - -def _get_current_accelerator_type() -> str: - if is_torch_npu_available(): - return "npu" - if is_torch_cuda_available(): - return "cuda" - return "cpu" - - -def _parse_backend_options(options: Any) -> Optional[dict[str, Any]]: - if options is None: - return None - - if not isinstance(options, dict): - raise ValueError("`profiler_backend_options` must be a mapping.") - - return options - - -def _resolve_optional_bool(value: Any, default: bool) -> bool: - if value is None: - return default - - return bool(value) - - -def _validate_int_option(name: str, value: Any, min_value: int) -> None: - if isinstance(value, bool) or not isinstance(value, int) or value < min_value: - if min_value == 1: - raise ValueError(f"`{name}` must be a positive integer.") - - raise ValueError(f"`{name}` must be an integer greater than or equal to {min_value}.") - - -def _validate_choice_option(name: str, value: Any, supported: set[str] | dict[str, Any]) -> str: - if not isinstance(value, str) or value not in supported: - raise ValueError(f"`{name}` must be one of {sorted(supported)}.") - - return value - - -def _get_explicit_profile_kwargs(**kwargs: Any) -> set[str]: - name_mapping = { - "record_shapes": "record_shapes", - "profile_memory": "profile_memory", - "with_stack": "with_stack", - } - return {profile_key for arg_key, profile_key in name_mapping.items() if kwargs[arg_key] is not None} - - -def _normalize_host_sys(host_sys: Any) -> list[str]: - if host_sys is None: - return [] - - if not isinstance(host_sys, list): - raise ValueError("`profiler_backend_options.npu.host_sys` must be a list.") - - normalized = [] - for item in host_sys: - key = _validate_choice_option("profiler_backend_options.npu.host_sys", item, _NPU_HOST_SYS) - normalized.append(_NPU_HOST_SYS[key]) - - return normalized - - -def _get_backend() -> _ProfilerBackend: - accelerator_type = _get_current_accelerator_type() - if accelerator_type == "cpu": - return _ProfilerBackend("cpu", torch.profiler, None) - - if accelerator_type == "cuda": - return _ProfilerBackend("cuda", torch.profiler, "CUDA") - - if accelerator_type == "npu": - try: - import torch_npu # type: ignore - except ImportError as exc: - raise RuntimeError("NPU profiler requires `torch_npu` to be installed.") from exc - - return _ProfilerBackend("npu", torch_npu.profiler, "NPU") - - raise RuntimeError(f"Profiler only supports CPU, CUDA and NPU devices, got current accelerator: {accelerator_type}.") - - -def _get_profiler_enum_value(profiler: Any, enum_name: str, member_name: str, option_name: str) -> Any: - enum_cls = getattr(profiler, enum_name, None) - if enum_cls is not None and hasattr(enum_cls, member_name): - return getattr(enum_cls, member_name) - - supported = [] if enum_cls is None else [name for name in dir(enum_cls) if not name.startswith("_")] - package_name = str(getattr(profiler, "__name__", "")).split(".", maxsplit=1)[0] - torch_npu_version = getattr(sys.modules.get(package_name), "__version__", None) - raise RuntimeError( - f"The installed torch_npu profiler does not support {option_name}={member_name}. " - f"Supported values: {supported}. torch_npu version: {torch_npu_version or 'unknown'}." - ) - - -def _filter_supported_kwargs( - fn: Any, - kwargs: dict[str, Any], - logger: Any = None, - explicit_keys: Optional[set[str]] = None, -) -> dict[str, Any]: - parameters = inspect.signature(fn).parameters - if any(param.kind is inspect.Parameter.VAR_KEYWORD for param in parameters.values()): - return kwargs - - supported = set(parameters) - dropped = set(kwargs) - supported - if dropped: - explicit_dropped = dropped & (explicit_keys or set()) - message_level = "warning" if explicit_dropped else "debug" - if explicit_dropped or logger is not None: - _log( - logger, - message_level, - f"Profiler backend ignored unsupported arguments: {sorted(dropped)}.", - ) - - return {key: value for key, value in kwargs.items() if key in supported} - - -def _log(logger: Any, level: str, message: str) -> None: - if logger is None: - return - - rank_method = f"{level}_rank0" - if hasattr(logger, rank_method): - getattr(logger, rank_method)(message) - else: - getattr(logger, level)(message) - - -class ProfilerController: - def __init__(self, args: Any, logger: Any = None) -> None: - self.config = ProfilerConfig.from_args(args) - self.logger = logger - self.profiler = None - self.backend_name: Optional[str] = None - self.trace_dir: Optional[str] = None - - @property - def enabled(self) -> bool: - return self.config.enabled - - def start(self, output_dir: str, initial_step: int = 0) -> None: - self.stop() - if not self.config.enabled: - return - - backend = _get_backend() - self.backend_name = backend.name - self.config.validate(backend.name) - if backend.name == "npu" and os.getenv("PROF_CONFIG_PATH"): - _log( - self.logger, - "warning", - "`PROF_CONFIG_PATH` is set. Do not enable dynamic_profile together with " - "`enable_torch_profiler`.", - ) - schedule = self.config.schedule_kwargs() - - rank = _get_rank() - if self.config.rank_mode == "rank0" and rank != 0: - return - - trace_root = self.config.output_dir or os.path.join(output_dir, "profiler") - self.trace_dir = os.path.realpath(os.path.join(trace_root, f"rank_{rank}")) - os.makedirs(self.trace_dir, exist_ok=True) - - profile_kwargs = dict( - activities=backend.build_activities(self.config), - schedule=backend.build_schedule(schedule), - on_trace_ready=backend.profiler_module.tensorboard_trace_handler(self.trace_dir), - record_shapes=self.config.record_shapes, - profile_memory=self.config.profile_memory, - with_stack=self.config.with_stack, - with_flops=self.config.with_flops, - with_modules=self.config.with_modules, - experimental_config=backend.build_experimental_config(self.config, logger=self.logger), - ) - profiler = backend.profiler_module.profile( - **_filter_supported_kwargs( - backend.profiler_module.profile, - profile_kwargs, - logger=self.logger, - explicit_keys=set(self.config.explicit_profile_kwargs or set()), - ) - ) - try: - profiler.start() - except Exception: - with suppress(Exception): - profiler.stop() - raise - - self.profiler = profiler - - first_active_step = initial_step + schedule["skip_first"] + schedule["wait"] + schedule["warmup"] + 1 - schedule_message = ( - f"skip_first={schedule['skip_first']}, wait={schedule['wait']}, warmup={schedule['warmup']}, " - f"active={schedule['active']}, repeat={schedule['repeat']}, first_active_step={first_active_step}" - ) - _log( - self.logger, - "info", - f"Profiler started on {backend.name}: {schedule_message}, initial_step={initial_step}. Traces -> {trace_root}", - ) - - def step(self) -> None: - if self.profiler is not None: - self.profiler.step() - - def stop(self) -> None: - profiler, self.profiler = self.profiler, None - if profiler is None: - return - - profiler.stop() - _log(self.logger, "info", "Profiler stopped.") diff --git a/src/llamafactory/hparams/training_args.py b/src/llamafactory/hparams/training_args.py index ed566111e3..84f4d8e6cc 100644 --- a/src/llamafactory/hparams/training_args.py +++ b/src/llamafactory/hparams/training_args.py @@ -70,16 +70,12 @@ class ProfilerArguments: enable_torch_profiler: bool = field( default=False, - metadata={"help": "Whether to enable torch profiler for collecting CPU/CUDA/NPU performance traces."}, + metadata={"help": "Whether to enable torch profiler for collecting performance traces."}, ) profiler_output_dir: Optional[str] = field( default=None, metadata={"help": "Directory to write profiler traces. Defaults to /profiler if not set."}, ) - profiler_skip_first: int = field( - default=0, - metadata={"help": "Number of steps to skip before the first profiler wait/warmup/active cycle."}, - ) profiler_wait_steps: int = field( default=1, metadata={"help": "Number of steps to skip at the start of each profiling cycle."}, @@ -96,61 +92,23 @@ class ProfilerArguments: default=1, metadata={"help": "Number of profiling cycles. Set to 0 for continuous profiling."}, ) - profiler_record_shapes: Optional[bool] = field( - default=None, + profiler_record_shapes: bool = field( + default=True, metadata={"help": "Whether to record tensor shapes during profiling."}, ) - profiler_profile_memory: Optional[bool] = field( - default=None, + profiler_profile_memory: bool = field( + default=True, metadata={"help": "Whether to profile memory usage."}, ) - profiler_with_stack: Optional[bool] = field( - default=None, + profiler_with_stack: bool = field( + default=True, metadata={"help": "Whether to record stack traces during profiling."}, ) - profiler_with_flops: bool = field( - default=False, - metadata={"help": "Whether to estimate FLOPs where supported by the profiler backend."}, - ) - profiler_with_modules: bool = field( - default=False, - metadata={"help": "Whether to record module hierarchy where supported by the profiler backend."}, - ) - profiler_activities: str = field( - default="auto", - metadata={"help": "Profiler activities to collect. Choices: auto, all, cpu, device."}, - ) - profiler_rank_mode: str = field( - default="all", - metadata={"help": "Profiler rank collection mode. Choices: all, rank0."}, - ) - profiler_level: str = field( - default="level0", - metadata={"help": "NPU profiler collection level. Choices: none, level0, level1, level2."}, - ) - profiler_aic_metrics: str = field( - default="auto", - metadata={ - "help": ( - "NPU AI Core metric: auto, none, pipe_utilization, arithmetic_utilization, memory, " - "memory_l0, memory_ub, l2_cache, memory_access, or resource_conflict_ratio." - ) - }, - ) - profiler_backend_options: Optional[dict] = field( - default=None, - metadata={ - "help": ( - "Backend-specific profiler mapping. Currently supports npu.data_simplification, npu.host_sys, " - "npu.sys_io, npu.sys_interconnection, and npu.gc_detect_threshold." - ) - }, - ) profile_modules: Optional[str] = field( default=None, metadata={ "help": ( - "Comma-separated list of module name patterns to profile with accelerator events. " + "Comma-separated list of module name patterns to profile with CUDA events. " "Supports fnmatch wildcards (e.g. 'model.layers.0.self_attn,model.layers.*.mlp'). " "Reports per-module forward/backward timing statistics at each logging step." ) diff --git a/src/llamafactory/train/callbacks.py b/src/llamafactory/train/callbacks.py index cae67e8d5b..b3826a7a9d 100644 --- a/src/llamafactory/train/callbacks.py +++ b/src/llamafactory/train/callbacks.py @@ -35,7 +35,6 @@ from ..extras.constants import TRAINER_LOG, V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME from ..extras.misc import get_peak_memory, is_env_enabled, is_torch_cuda_available, is_torch_npu_available, use_ray from ..extras.packages import is_safetensors_available -from ..extras.profiler import ProfilerController if is_safetensors_available(): @@ -342,13 +341,12 @@ def on_prediction_step( class TorchProfilerCallback(TrainerCallback): - r"""A callback for collecting CPU/CUDA/NPU profiler traces during training. + r"""A callback for collecting torch.profiler traces during training. Activated by setting ``enable_torch_profiler: true`` in the YAML config. Configuration fields (in YAML): profiler_output_dir – where to write traces (default: /profiler) - profiler_skip_first – steps to skip before the first cycle (default: 0) profiler_wait_steps – steps to skip at start of each cycle (default: 1) profiler_warmup_steps – profiler warm-up steps per cycle (default: 1) profiler_active_steps – steps to record per cycle (default: 1) @@ -356,51 +354,80 @@ class TorchProfilerCallback(TrainerCallback): profiler_record_shapes – record tensor shapes (default: true) profiler_profile_memory – profile memory usage (default: true) profiler_with_stack – record stack traces (default: true) - profiler_activities – choices: auto, all, cpu, device (default: auto) - profiler_rank_mode – choices: all, rank0 (default: all) - profiler_level – choices: none, level0, level1, level2 - profiler_aic_metrics – choices: auto, none, pipe_utilization, arithmetic_utilization, - memory, memory_l0, memory_ub, l2_cache, memory_access, - resource_conflict_ratio Trace files (one per rank, Chrome / TensorBoard JSON format) are written to ``/rank_/``. - - Example: - profiler_skip_first: 8 # skip 8 optimizer steps before wait/warmup/active - profiler_wait_steps: 0 - profiler_warmup_steps: 1 - profiler_active_steps: 3 - profiler_rank_mode: rank0 - profiler_level: level1 - profiler_aic_metrics: pipe_utilization - profiler_backend_options: - npu: - host_sys: [cpu, mem] - sys_io: true - data_simplification: false """ def __init__(self, training_args: "TrainingArguments") -> None: - self.profiler = ProfilerController(training_args, logger=logger) + self.profiler = None + self.profiler_args = training_args + + @staticmethod + def _get_rank() -> int: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + return 0 @override def on_train_begin( self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs ) -> None: - self.profiler.start(args.output_dir, initial_step=state.global_step) + if self.profiler is not None: + self.profiler.stop() + self.profiler = None + + pa = self.profiler_args + output_dir = pa.profiler_output_dir or os.path.join(args.output_dir, "profiler") + rank = self._get_rank() + trace_dir = os.path.join(output_dir, f"rank_{rank}") + os.makedirs(trace_dir, exist_ok=True) + + activities = [torch.profiler.ProfilerActivity.CPU] + try: + if is_torch_cuda_available(): + activities.append(torch.profiler.ProfilerActivity.CUDA) + if is_torch_npu_available(): + activities.append(torch.profiler.ProfilerActivity.NPU) + except Exception: + pass + + self.profiler = torch.profiler.profile( + activities=activities, + schedule=torch.profiler.schedule( + wait=pa.profiler_wait_steps, + warmup=pa.profiler_warmup_steps, + active=pa.profiler_active_steps, + repeat=pa.profiler_repeat, + ), + on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir), + record_shapes=pa.profiler_record_shapes, + profile_memory=pa.profiler_profile_memory, + with_stack=pa.profiler_with_stack, + ) + self.profiler.start() + logger.info_rank0( + f"TorchProfiler started — schedule: wait={pa.profiler_wait_steps}, warmup={pa.profiler_warmup_steps}, " + f"active={pa.profiler_active_steps}, repeat={pa.profiler_repeat}. Traces → {output_dir}" + ) @override def on_step_end( self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs ) -> None: - self.profiler.step() + if self.profiler is not None: + self.profiler.step() @override def on_train_end( self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs ) -> None: - self.profiler.stop() + if self.profiler is not None: + self.profiler.stop() + self.profiler = None + logger.info_rank0("TorchProfiler stopped.") class ReporterCallback(TrainerCallback): diff --git a/src/llamafactory/v1/accelerator/profiler.py b/src/llamafactory/v1/accelerator/profiler.py index ba75573c58..bc43a8885e 100644 --- a/src/llamafactory/v1/accelerator/profiler.py +++ b/src/llamafactory/v1/accelerator/profiler.py @@ -1,8 +1,26 @@ -from __future__ import annotations +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. -from typing import TYPE_CHECKING, Any +import inspect +import os +import sys +from contextlib import suppress +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional -from llamafactory.extras.profiler import ProfilerController +import torch +from transformers.utils import is_torch_cuda_available, is_torch_npu_available from ..utils import logging from ..utils.callbacks import TrainerCallback, TrainerState @@ -15,15 +33,509 @@ logger = logging.get_logger(__name__) +_SUPPORTED_ACTIVITIES = {"auto", "all", "cpu", "device"} +_SUPPORTED_RANK_MODES = {"all", "rank0"} +_NPU_PROFILER_LEVELS = { + "none": "Level_none", + "level0": "Level0", + "level1": "Level1", + "level2": "Level2", +} +_NPU_AIC_METRICS = { + "none": "AiCoreNone", + "pipe_utilization": "PipeUtilization", + "arithmetic_utilization": "ArithmeticUtilization", + "memory": "Memory", + "memory_l0": "MemoryL0", + "memory_ub": "MemoryUB", + "l2_cache": "L2Cache", + "memory_access": "MemoryAccess", + "resource_conflict_ratio": "ResourceConflictRatio", +} +_NPU_HOST_SYS = { + "cpu": "CPU", + "mem": "MEM", + "disk": "DISK", + "network": "NETWORK", + "osrt": "OSRT", +} +_NPU_BACKEND_OPTIONS = { + "data_simplification", + "host_sys", + "sys_io", + "sys_interconnection", + "gc_detect_threshold", +} + + +@dataclass +class ProfilerConfig: + enabled: bool = False + output_dir: Optional[str] = None + skip_first: int = 0 + wait_steps: int = 1 + warmup_steps: int = 1 + active_steps: int = 1 + repeat: int = 1 + record_shapes: bool = False + profile_memory: bool = False + with_stack: bool = False + with_flops: bool = False + with_modules: bool = False + activities: str = "auto" + rank_mode: str = "all" + level: str = "level0" + aic_metrics: str = "auto" + backend_options: Optional[dict[str, Any]] = None + explicit_profile_kwargs: Optional[set[str]] = None + + @classmethod + def from_args(cls, args: Any) -> "ProfilerConfig": + enable_torch_profiler = bool(getattr(args, "enable_torch_profiler", False)) + record_shapes = getattr(args, "profiler_record_shapes", None) + profile_memory = getattr(args, "profiler_profile_memory", None) + with_stack = getattr(args, "profiler_with_stack", None) + explicit_profile_kwargs = _get_explicit_profile_kwargs( + record_shapes=record_shapes, + profile_memory=profile_memory, + with_stack=with_stack, + ) + if getattr(args, "profiler_with_flops", False): + explicit_profile_kwargs.add("with_flops") + if getattr(args, "profiler_with_modules", False): + explicit_profile_kwargs.add("with_modules") + return cls( + enabled=enable_torch_profiler, + output_dir=getattr(args, "profiler_output_dir", None), + skip_first=getattr(args, "profiler_skip_first", 0), + wait_steps=getattr(args, "profiler_wait_steps", 1), + warmup_steps=getattr(args, "profiler_warmup_steps", 1), + active_steps=getattr(args, "profiler_active_steps", 1), + repeat=getattr(args, "profiler_repeat", 1), + record_shapes=_resolve_optional_bool(record_shapes, enable_torch_profiler), + profile_memory=_resolve_optional_bool(profile_memory, enable_torch_profiler), + with_stack=_resolve_optional_bool(with_stack, enable_torch_profiler), + with_flops=getattr(args, "profiler_with_flops", False), + with_modules=getattr(args, "profiler_with_modules", False), + activities=getattr(args, "profiler_activities", "auto"), + rank_mode=getattr(args, "profiler_rank_mode", "all"), + level=getattr(args, "profiler_level", "level0"), + aic_metrics=getattr(args, "profiler_aic_metrics", "auto"), + backend_options=_parse_backend_options(getattr(args, "profiler_backend_options", None)), + explicit_profile_kwargs=explicit_profile_kwargs, + ) + + def validate(self, backend_name: Optional[str] = None) -> None: + if not self.enabled: + return + + _validate_int_option("profiler_skip_first", self.skip_first, min_value=0) + _validate_int_option("profiler_wait_steps", self.wait_steps, min_value=0) + _validate_int_option("profiler_warmup_steps", self.warmup_steps, min_value=0) + _validate_int_option("profiler_active_steps", self.active_steps, min_value=1) + _validate_int_option("profiler_repeat", self.repeat, min_value=0) + _validate_choice_option("profiler_activities", self.activities, _SUPPORTED_ACTIVITIES) + _validate_choice_option("profiler_rank_mode", self.rank_mode, _SUPPORTED_RANK_MODES) + + if backend_name == "npu" and self.activities != "cpu": + self.npu_profiler_level_name() + self.npu_aic_metrics_name() + self.npu_backend_options() + self.schedule_kwargs() + + def npu_profiler_level_name(self) -> str: + key = _validate_choice_option("profiler_level", self.level, _NPU_PROFILER_LEVELS) + return _NPU_PROFILER_LEVELS[key] + + def npu_aic_metrics_name(self) -> str: + key = _validate_choice_option("profiler_aic_metrics", self.aic_metrics, {"auto", *_NPU_AIC_METRICS}) + if key == "auto": + if self.npu_profiler_level_name() in ("Level1", "Level2"): + return "PipeUtilization" + return "AiCoreNone" + + metric_name = _NPU_AIC_METRICS[key] + if metric_name != "AiCoreNone" and self.npu_profiler_level_name() not in ("Level1", "Level2"): + raise ValueError( + "`profiler_aic_metrics` requires `profiler_level` to be `level1` or `level2` on NPU." + ) + + return metric_name + + def npu_backend_options(self) -> dict[str, Any]: + if self.backend_options is None: + return {} + + unsupported_backends = set(self.backend_options) - {"npu"} + if unsupported_backends: + raise ValueError(f"`profiler_backend_options` only supports `npu`, got {sorted(unsupported_backends)}.") + + if "npu" not in self.backend_options or self.backend_options["npu"] is None: + return {} + + npu_options = self.backend_options["npu"] + if not isinstance(npu_options, dict): + raise ValueError("`profiler_backend_options.npu` must be a mapping.") + + unsupported_options = set(npu_options) - _NPU_BACKEND_OPTIONS + if unsupported_options: + raise ValueError( + f"`profiler_backend_options.npu` only supports {sorted(_NPU_BACKEND_OPTIONS)}, " + f"got {sorted(unsupported_options)}." + ) + + normalized_options = dict(npu_options) + if "data_simplification" in normalized_options and not isinstance( + normalized_options["data_simplification"], bool + ): + raise ValueError("`profiler_backend_options.npu.data_simplification` must be a boolean.") + if "sys_io" in normalized_options and not isinstance(normalized_options["sys_io"], bool): + raise ValueError("`profiler_backend_options.npu.sys_io` must be a boolean.") + if "sys_interconnection" in normalized_options and not isinstance( + normalized_options["sys_interconnection"], bool + ): + raise ValueError("`profiler_backend_options.npu.sys_interconnection` must be a boolean.") + if normalized_options.get("gc_detect_threshold") is not None: + threshold = normalized_options["gc_detect_threshold"] + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)) or threshold < 0: + raise ValueError("`profiler_backend_options.npu.gc_detect_threshold` must be null or >= 0.") + + normalized_options["host_sys"] = _normalize_host_sys(normalized_options.get("host_sys", [])) + return normalized_options + + def explicit_npu_backend_options(self) -> set[str]: + if self.backend_options is None: + return set() + + npu_options = self.backend_options.get("npu") + if not isinstance(npu_options, dict): + return set() + + return set(npu_options) + + def schedule_kwargs(self) -> dict[str, int]: + return dict( + wait=self.wait_steps, + warmup=self.warmup_steps, + active=self.active_steps, + repeat=self.repeat, + skip_first=self.skip_first, + ) + + +class _ProfilerBackend: + def __init__(self, name: str, profiler_module: Any, device_activity_name: Optional[str]) -> None: + self.name = name + self.profiler_module = profiler_module + self.device_activity_name = device_activity_name + + def build_schedule(self, kwargs: dict[str, Any]) -> Any: + if self.name == "npu": + return self.profiler_module.schedule( + wait=kwargs["wait"], + active=kwargs["active"], + warmup=kwargs["warmup"], + repeat=kwargs["repeat"], + skip_first=kwargs["skip_first"], + ) + + return self.profiler_module.schedule(**kwargs) + + def build_activities(self, config: ProfilerConfig) -> list[Any]: + activity_cls = self.profiler_module.ProfilerActivity + activities = [] + if config.activities in ("auto", "all", "cpu"): + activities.append(activity_cls.CPU) + if config.activities == "device" and self.device_activity_name is None: + raise ValueError("`profiler_activities: device` requires a CUDA or NPU backend.") + if config.activities in ("auto", "all", "device"): + if self.device_activity_name is not None: + activities.append(getattr(activity_cls, self.device_activity_name)) + return activities + + def build_experimental_config(self, config: ProfilerConfig, logger: Any = None) -> Optional[Any]: + if self.name != "npu" or config.activities == "cpu": + return None + + profiler = self.profiler_module + npu_options = config.npu_backend_options() + config_kwargs = dict( + export_type=[profiler.ExportType.Text], + profiler_level=_get_profiler_enum_value( + profiler, "ProfilerLevel", config.npu_profiler_level_name(), "`profiler_level`" + ), + aic_metrics=_get_profiler_enum_value( + profiler, "AiCMetrics", config.npu_aic_metrics_name(), "`profiler_aic_metrics`" + ), + l2_cache=False, + op_attr=False, + data_simplification=npu_options.get("data_simplification", True), + record_op_args=False, + gc_detect_threshold=npu_options.get("gc_detect_threshold"), + host_sys=[ + _get_profiler_enum_value( + profiler, "HostSystem", item, "`profiler_backend_options.npu.host_sys`" + ) + for item in npu_options.get("host_sys", []) + ], + sys_io=npu_options.get("sys_io", False), + sys_interconnection=npu_options.get("sys_interconnection", False), + mstx=False, + mstx_domain_include=[], + mstx_domain_exclude=[], + ) + explicit_keys = {"profiler_level", "aic_metrics"} + explicit_keys.update(config.explicit_npu_backend_options()) + return profiler._ExperimentalConfig( + **_filter_supported_kwargs( + profiler._ExperimentalConfig, + config_kwargs, + logger=logger, + explicit_keys=explicit_keys, + ) + ) + + +def _get_rank() -> int: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + + return int(os.getenv("RANK", "0")) + + +def _get_current_accelerator_type() -> str: + if is_torch_npu_available(): + return "npu" + if is_torch_cuda_available(): + return "cuda" + return "cpu" + + +def _parse_backend_options(options: Any) -> Optional[dict[str, Any]]: + if options is None: + return None + + if not isinstance(options, dict): + raise ValueError("`profiler_backend_options` must be a mapping.") + + return options + + +def _resolve_optional_bool(value: Any, default: bool) -> bool: + if value is None: + return default + + return bool(value) + + +def _validate_int_option(name: str, value: Any, min_value: int) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value < min_value: + if min_value == 1: + raise ValueError(f"`{name}` must be a positive integer.") + + raise ValueError(f"`{name}` must be an integer greater than or equal to {min_value}.") + + +def _validate_choice_option(name: str, value: Any, supported: set[str] | dict[str, Any]) -> str: + if not isinstance(value, str) or value not in supported: + raise ValueError(f"`{name}` must be one of {sorted(supported)}.") + + return value + + +def _get_explicit_profile_kwargs(**kwargs: Any) -> set[str]: + name_mapping = { + "record_shapes": "record_shapes", + "profile_memory": "profile_memory", + "with_stack": "with_stack", + } + return {profile_key for arg_key, profile_key in name_mapping.items() if kwargs[arg_key] is not None} + + +def _normalize_host_sys(host_sys: Any) -> list[str]: + if host_sys is None: + return [] + + if not isinstance(host_sys, list): + raise ValueError("`profiler_backend_options.npu.host_sys` must be a list.") + + normalized = [] + for item in host_sys: + key = _validate_choice_option("profiler_backend_options.npu.host_sys", item, _NPU_HOST_SYS) + normalized.append(_NPU_HOST_SYS[key]) + + return normalized + + +def _get_backend() -> _ProfilerBackend: + accelerator_type = _get_current_accelerator_type() + if accelerator_type == "cpu": + return _ProfilerBackend("cpu", torch.profiler, None) + + if accelerator_type == "cuda": + return _ProfilerBackend("cuda", torch.profiler, "CUDA") + + if accelerator_type == "npu": + try: + import torch_npu # type: ignore + except ImportError as exc: + raise RuntimeError("NPU profiler requires `torch_npu` to be installed.") from exc + + return _ProfilerBackend("npu", torch_npu.profiler, "NPU") + + raise RuntimeError(f"Profiler only supports CPU, CUDA and NPU devices, got current accelerator: {accelerator_type}.") + + +def _get_profiler_enum_value(profiler: Any, enum_name: str, member_name: str, option_name: str) -> Any: + enum_cls = getattr(profiler, enum_name, None) + if enum_cls is not None and hasattr(enum_cls, member_name): + return getattr(enum_cls, member_name) + + supported = [] if enum_cls is None else [name for name in dir(enum_cls) if not name.startswith("_")] + package_name = str(getattr(profiler, "__name__", "")).split(".", maxsplit=1)[0] + torch_npu_version = getattr(sys.modules.get(package_name), "__version__", None) + raise RuntimeError( + f"The installed torch_npu profiler does not support {option_name}={member_name}. " + f"Supported values: {supported}. torch_npu version: {torch_npu_version or 'unknown'}." + ) + + +def _filter_supported_kwargs( + fn: Any, + kwargs: dict[str, Any], + logger: Any = None, + explicit_keys: Optional[set[str]] = None, +) -> dict[str, Any]: + parameters = inspect.signature(fn).parameters + if any(param.kind is inspect.Parameter.VAR_KEYWORD for param in parameters.values()): + return kwargs + + supported = set(parameters) + dropped = set(kwargs) - supported + if dropped: + explicit_dropped = dropped & (explicit_keys or set()) + message_level = "warning" if explicit_dropped else "debug" + if explicit_dropped or logger is not None: + _log( + logger, + message_level, + f"Profiler backend ignored unsupported arguments: {sorted(dropped)}.", + ) + + return {key: value for key, value in kwargs.items() if key in supported} + + +def _log(logger: Any, level: str, message: str) -> None: + if logger is None: + return + + rank_method = f"{level}_rank0" + if hasattr(logger, rank_method): + getattr(logger, rank_method)(message) + else: + getattr(logger, level)(message) + + +class ProfilerController: + def __init__(self, args: Any, logger: Any = None) -> None: + self.config = ProfilerConfig.from_args(args) + self.logger = logger + self.profiler = None + self.backend_name: Optional[str] = None + self.trace_dir: Optional[str] = None + + @property + def enabled(self) -> bool: + return self.config.enabled + + def start(self, output_dir: str, initial_step: int = 0) -> None: + self.stop() + if not self.config.enabled: + return + + backend = _get_backend() + self.backend_name = backend.name + self.config.validate(backend.name) + if backend.name == "npu" and os.getenv("PROF_CONFIG_PATH"): + _log( + self.logger, + "warning", + "`PROF_CONFIG_PATH` is set. Do not enable dynamic_profile together with " + "`enable_torch_profiler`.", + ) + schedule = self.config.schedule_kwargs() + + rank = _get_rank() + if self.config.rank_mode == "rank0" and rank != 0: + return + + trace_root = self.config.output_dir or os.path.join(output_dir, "profiler") + self.trace_dir = os.path.realpath(os.path.join(trace_root, f"rank_{rank}")) + os.makedirs(self.trace_dir, exist_ok=True) + + profile_kwargs = dict( + activities=backend.build_activities(self.config), + schedule=backend.build_schedule(schedule), + on_trace_ready=backend.profiler_module.tensorboard_trace_handler(self.trace_dir), + record_shapes=self.config.record_shapes, + profile_memory=self.config.profile_memory, + with_stack=self.config.with_stack, + with_flops=self.config.with_flops, + with_modules=self.config.with_modules, + experimental_config=backend.build_experimental_config(self.config, logger=self.logger), + ) + profiler = backend.profiler_module.profile( + **_filter_supported_kwargs( + backend.profiler_module.profile, + profile_kwargs, + logger=self.logger, + explicit_keys=set(self.config.explicit_profile_kwargs or set()), + ) + ) + try: + profiler.start() + except Exception: + with suppress(Exception): + profiler.stop() + raise + + self.profiler = profiler + + first_active_step = initial_step + schedule["skip_first"] + schedule["wait"] + schedule["warmup"] + 1 + schedule_message = ( + f"skip_first={schedule['skip_first']}, wait={schedule['wait']}, warmup={schedule['warmup']}, " + f"active={schedule['active']}, repeat={schedule['repeat']}, first_active_step={first_active_step}" + ) + _log( + self.logger, + "info", + f"Profiler started on {backend.name}: {schedule_message}, initial_step={initial_step}. Traces -> {trace_root}", + ) + + def step(self) -> None: + if self.profiler is not None: + self.profiler.step() + + def stop(self) -> None: + profiler, self.profiler = self.profiler, None + if profiler is None: + return + + profiler.stop() + _log(self.logger, "info", "Profiler stopped.") + + class ProfilerCallback(TrainerCallback): - def __init__(self, args: TrainingArguments) -> None: + def __init__(self, args: "TrainingArguments") -> None: self.profiler = ProfilerController(args, logger=logger) - def on_train_begin(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + def on_train_begin(self, args: "TrainingArguments", state: TrainerState, **kwargs: Any) -> None: self.profiler.start(args.output_dir, initial_step=state.global_step) - def on_step_end(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + def on_step_end(self, args: "TrainingArguments", state: TrainerState, **kwargs: Any) -> None: self.profiler.step() - def on_train_end(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + def on_train_end(self, args: "TrainingArguments", state: TrainerState, **kwargs: Any) -> None: self.profiler.stop() From b6f8634ddc18be3ed25128510a44e6b9fb78c7b4 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Mon, 22 Jun 2026 02:37:34 +0000 Subject: [PATCH 6/9] move v1 profiler callback to callbacks package --- src/llamafactory/v1/accelerator/profiler.py | 21 +--------- src/llamafactory/v1/core/base_trainer.py | 3 +- .../v1/utils/callbacks/__init__.py | 2 + .../v1/utils/callbacks/profiler_callback.py | 42 +++++++++++++++++++ 4 files changed, 46 insertions(+), 22 deletions(-) create mode 100644 src/llamafactory/v1/utils/callbacks/profiler_callback.py diff --git a/src/llamafactory/v1/accelerator/profiler.py b/src/llamafactory/v1/accelerator/profiler.py index bc43a8885e..6f1beaca58 100644 --- a/src/llamafactory/v1/accelerator/profiler.py +++ b/src/llamafactory/v1/accelerator/profiler.py @@ -17,17 +17,12 @@ import sys from contextlib import suppress from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional import torch from transformers.utils import is_torch_cuda_available, is_torch_npu_available from ..utils import logging -from ..utils.callbacks import TrainerCallback, TrainerState - - -if TYPE_CHECKING: - from ..config import TrainingArguments logger = logging.get_logger(__name__) @@ -525,17 +520,3 @@ def stop(self) -> None: profiler.stop() _log(self.logger, "info", "Profiler stopped.") - - -class ProfilerCallback(TrainerCallback): - def __init__(self, args: "TrainingArguments") -> None: - self.profiler = ProfilerController(args, logger=logger) - - def on_train_begin(self, args: "TrainingArguments", state: TrainerState, **kwargs: Any) -> None: - self.profiler.start(args.output_dir, initial_step=state.global_step) - - def on_step_end(self, args: "TrainingArguments", state: TrainerState, **kwargs: Any) -> None: - self.profiler.step() - - def on_train_end(self, args: "TrainingArguments", state: TrainerState, **kwargs: Any) -> None: - self.profiler.stop() diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index e45fdcb681..85cb985fa4 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -39,6 +39,7 @@ from ..utils.callbacks import ( CallbackHandler, LoggingCallback, + ProfilerCallback, TrainerCallback, TrainerState, ) @@ -126,8 +127,6 @@ def __init__( # Callbacks self.callback_handler = CallbackHandler([LoggingCallback()], trainer=self) if getattr(self.args, "enable_torch_profiler", False): - from ..accelerator.profiler import ProfilerCallback - self.callback_handler.add_callback(ProfilerCallback(self.args)) for cb in callbacks or []: diff --git a/src/llamafactory/v1/utils/callbacks/__init__.py b/src/llamafactory/v1/utils/callbacks/__init__.py index 0da31ed3fc..0009c8ca16 100644 --- a/src/llamafactory/v1/utils/callbacks/__init__.py +++ b/src/llamafactory/v1/utils/callbacks/__init__.py @@ -13,12 +13,14 @@ # limitations under the License. from .logging_callback import LoggingCallback +from .profiler_callback import ProfilerCallback from .trainer_callback import CallbackHandler, TrainerCallback, TrainerState __all__ = [ "CallbackHandler", "LoggingCallback", + "ProfilerCallback", "TrainerCallback", "TrainerState", ] diff --git a/src/llamafactory/v1/utils/callbacks/profiler_callback.py b/src/llamafactory/v1/utils/callbacks/profiler_callback.py new file mode 100644 index 0000000000..e507b4de0e --- /dev/null +++ b/src/llamafactory/v1/utils/callbacks/profiler_callback.py @@ -0,0 +1,42 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from ...accelerator.profiler import ProfilerController +from .. import logging +from .trainer_callback import TrainerCallback, TrainerState + + +if TYPE_CHECKING: + from ...config import TrainingArguments + + +logger = logging.get_logger(__name__) + + +class ProfilerCallback(TrainerCallback): + def __init__(self, args: TrainingArguments) -> None: + self.profiler = ProfilerController(args, logger=logger) + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + self.profiler.start(args.output_dir, initial_step=state.global_step) + + def on_step_end(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + self.profiler.step() + + def on_train_end(self, args: TrainingArguments, state: TrainerState, **kwargs: Any) -> None: + self.profiler.stop() From 6200725516f909bed4dac5b68d5686d88c647fdc Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Mon, 22 Jun 2026 06:21:39 +0000 Subject: [PATCH 7/9] move v1 profiler controller to core utils --- src/llamafactory/v1/{accelerator => core/utils}/profiler.py | 2 +- src/llamafactory/v1/utils/callbacks/profiler_callback.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename src/llamafactory/v1/{accelerator => core/utils}/profiler.py (99%) diff --git a/src/llamafactory/v1/accelerator/profiler.py b/src/llamafactory/v1/core/utils/profiler.py similarity index 99% rename from src/llamafactory/v1/accelerator/profiler.py rename to src/llamafactory/v1/core/utils/profiler.py index 6f1beaca58..846e950a47 100644 --- a/src/llamafactory/v1/accelerator/profiler.py +++ b/src/llamafactory/v1/core/utils/profiler.py @@ -22,7 +22,7 @@ import torch from transformers.utils import is_torch_cuda_available, is_torch_npu_available -from ..utils import logging +from ...utils import logging logger = logging.get_logger(__name__) diff --git a/src/llamafactory/v1/utils/callbacks/profiler_callback.py b/src/llamafactory/v1/utils/callbacks/profiler_callback.py index e507b4de0e..484bcca768 100644 --- a/src/llamafactory/v1/utils/callbacks/profiler_callback.py +++ b/src/llamafactory/v1/utils/callbacks/profiler_callback.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Any -from ...accelerator.profiler import ProfilerController +from ...core.utils.profiler import ProfilerController from .. import logging from .trainer_callback import TrainerCallback, TrainerState From 4acc10842b72f64c5f17b2b051308057cca4cdae Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Mon, 22 Jun 2026 06:22:07 +0000 Subject: [PATCH 8/9] restore v1 accelerator profiler placeholder --- src/llamafactory/v1/accelerator/profiler.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 src/llamafactory/v1/accelerator/profiler.py diff --git a/src/llamafactory/v1/accelerator/profiler.py b/src/llamafactory/v1/accelerator/profiler.py new file mode 100644 index 0000000000..e69de29bb2 From 95b3cef35867ff19c5005af01e7d03f92e2af0c0 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Mon, 22 Jun 2026 06:31:15 +0000 Subject: [PATCH 9/9] support bare optional profiler bool flags --- src/llamafactory/v1/config/training_args.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/llamafactory/v1/config/training_args.py b/src/llamafactory/v1/config/training_args.py index 9963d861f4..73d42dee3d 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/src/llamafactory/v1/config/training_args.py @@ -99,15 +99,15 @@ class TrainingArguments: ) profiler_record_shapes: bool | None = field( default=None, - metadata={"help": "Whether to record tensor shapes during profiling."}, + metadata={"help": "Whether to record tensor shapes during profiling.", "nargs": "?", "const": True}, ) profiler_profile_memory: bool | None = field( default=None, - metadata={"help": "Whether to profile memory usage."}, + metadata={"help": "Whether to profile memory usage.", "nargs": "?", "const": True}, ) profiler_with_stack: bool | None = field( default=None, - metadata={"help": "Whether to record stack traces during profiling."}, + metadata={"help": "Whether to record stack traces during profiling.", "nargs": "?", "const": True}, ) profiler_with_flops: bool = field( default=False,