Metrics for training

examples/training/qwen3/finetune_qwen3.sh

@@ -10,10 +10,12 @@ export MALLOC_ARENA_MAX=64 # Host OOM mitigation
 PROCESSES_PER_NODE=32
 NUM_EPOCHS=3
 TP_DEGREE=8
+PP_DEGREE=4
 BS=1
 GRADIENT_ACCUMULATION_STEPS=8
 LOGGING_STEPS=2
-MODEL_NAME="Qwen/Qwen3-8B" # Change this to the desired model name
+# MODEL_NAME="Qwen/Qwen3-8B" # Change this to the desired model name
+MODEL_NAME="Qwen/Qwen3-0.6B" # Change this to the desired model name
 OUTPUT_DIR="$(echo $MODEL_NAME | cut -d'/' -f2)-finetuned"
 DISTRIBUTED_ARGS="--nproc_per_node $PROCESSES_PER_NODE"
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
@@ -34,6 +36,7 @@ torchrun --nproc_per_node $PROCESSES_PER_NODE finetune_qwen3.py \
   --learning_rate 8e-4 \
   --bf16 \
   --tensor_parallel_size $TP_DEGREE \
+  --pipeline_parallel_size $PP_DEGREE \
   --zero_1 \
   --async_save \
   --logging_steps $LOGGING_STEPS \

optimum/neuron/trainers/metrics/__init__.py

@@ -0,0 +1,17 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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 .collector import TrainingMetricsCollector
+from .window import MovingAverageWindow

optimum/neuron/trainers/metrics/base.py

@@ -0,0 +1,74 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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 abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..training_args import NeuronTrainingArguments
+
+
+if TYPE_CHECKING:
+    from .collector import TrainingMetricsCollector
+
+
+class MetricUnit:
+    SECONDS = "s"
+    MILLISECONDS = "ms"
+    TOKENS_PER_SECOND = "tokens/s"
+    SAMPLES_PER_SECOND = "samples/s"
+    PERCENT = "%"
+    COUNT = "count"
+    TFLOPS = "TFLOP/s"
+    RATIO = "ratio"
+    NONE = ""
+
+
+@dataclass
+class MetricPlugin(ABC):
+    """Base class for metrics plugins. Each plugin calculates one type of metric."""
+
+    name: str
+    requires_accumulation: bool = False
+    depends_on: list[str] | None = None
+
+    @abstractmethod
+    def is_enabled(self, args: NeuronTrainingArguments) -> bool:
+        """Check if this plugin should be active."""
+        pass
+
+    @abstractmethod
+    def calculate_realtime(self, window_stats: dict, collector: "TrainingMetricsCollector") -> dict[str, float]:
+        """Calculate train/ metrics from current window data."""
+        pass
+
+    @abstractmethod
+    def calculate_summary(self, summary_data: dict, collector: "TrainingMetricsCollector") -> dict[str, float]:
+        """Calculate summary/ metrics from all collected data."""
+        pass
+
+    def get_metric_names(self) -> list[str]:
+        """Get the metrics this plugin provides. Override for multi-metric plugins."""
+        return [self.name]
+
+    def handles_metric(self, metric_name: str) -> bool:
+        """Check if this plugin handles the given metric."""
+        return metric_name in self.get_metric_names()
+
+    def get_metric_units(self) -> dict[str, str]:
+        """Get units for each metric this plugin produces. Override in subclasses."""
+        return dict.fromkeys(self.get_metric_names(), MetricUnit.NONE)