Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / environ_meter_callback.py: 66%
183 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Training and environment metric collection callback."""
17from __future__ import annotations
19import time
20from collections.abc import Mapping, Sequence
21from typing import Any
23from hyper_parallel.trainer.runtime.distributed import get_world_size_safe
24from hyper_parallel.trainer.runtime.distributed import all_reduce
25from hyper_parallel.data.constants import IGNORE_INDEX
26from hyper_parallel.trainer.runtime.device import get_device_type, get_torch_device
28from .base import Callback, TrainerState
31class EnvironMeterCallback(Callback):
32 """Collect structured training, throughput, and memory metrics.
34 The callback is the single producer of ``trainer.step_train_metrics`` and
35 ``trainer.step_env_metrics``. Presentation and remote logging callbacks
36 consume those dictionaries without recalculating or reducing metrics.
37 """
39 def __init__(self, trainer: Any) -> None:
40 """Initialize per-step and cumulative counters.
42 Args:
43 trainer: Trainer that owns the callback lifecycle.
44 """
45 super().__init__(trainer)
46 self._step_start_time = 0.0
47 self._local_step_tokens: Any = None
48 self._local_step_samples = 0
49 self._consumed_tokens = 0
50 self._consumed_samples = 0
51 self.trainer.step_train_metrics = {}
52 self.trainer.step_env_metrics = {}
54 @staticmethod
55 def _scalar(value: Any, name: str) -> float:
56 """Convert a scalar or scalar tensor-like value to ``float``.
58 Args:
59 value: Scalar value to convert.
60 name: Metric name used in validation errors.
62 Returns:
63 Converted floating-point value.
65 Raises:
66 ValueError: If the value cannot be converted to a scalar float.
67 """
68 item = getattr(value, "item", None)
69 if callable(item):
70 value = item()
71 try:
72 return float(value)
73 except (TypeError, ValueError) as exc:
74 raise ValueError(f"Metric {name!r} must be scalar, but got {value!r}") from exc
76 @staticmethod
77 def _tensor_numel(value: Any) -> int | None:
78 """Return ``value.numel()`` when it exposes a tensor-like interface."""
79 numel = getattr(value, "numel", None)
80 if not callable(numel):
81 return None
82 return int(numel())
84 @classmethod
85 def _batch_tokens(cls, batch: Mapping[str, Any]) -> Any:
86 """Count text tokens without synchronizing a device scalar to the host."""
87 token_count = batch.get("token_count")
88 if token_count is not None:
89 return token_count
91 labels = batch.get("labels")
92 if labels is not None and callable(getattr(labels, "sum", None)):
93 return (labels != IGNORE_INDEX).sum()
95 attention_mask = batch.get("attention_mask")
96 attention_mask_shape = getattr(attention_mask, "shape", ())
97 if (
98 len(attention_mask_shape) <= 2
99 and attention_mask is not None
100 and callable(getattr(attention_mask, "sum", None))
101 ):
102 return attention_mask.sum()
104 input_ids = batch.get("input_ids")
105 input_numel = cls._tensor_numel(input_ids)
106 if input_numel is not None:
107 return input_numel
108 return 0
110 @staticmethod
111 def _batch_samples(batch: Mapping[str, Any]) -> int:
112 """Count logical samples in one micro-batch."""
113 value = batch.get("input_ids")
114 if value is None:
115 value = batch.get("labels")
116 shape = getattr(value, "shape", None)
117 if shape is None or len(shape) == 0:
118 return 0
119 if len(shape) == 1:
120 return 1
121 return int(shape[0])
123 @staticmethod
124 def _batch_mapping(value: Any) -> Mapping[str, Any] | None:
125 """Return metric inputs for a mapping or prepared runtime batch."""
126 if isinstance(value, Mapping):
127 return value
128 loss_count_inputs = getattr(value, "loss_count_inputs", None)
129 if not callable(loss_count_inputs):
130 return None
131 metric_inputs = loss_count_inputs()
132 return metric_inputs if isinstance(metric_inputs, Mapping) else None
134 @classmethod
135 def _micro_batches(cls, value: Any) -> list[Mapping[str, Any]]:
136 """Normalize callback input into a list of mapping micro-batches."""
137 if value is None:
138 return []
139 batch = cls._batch_mapping(value)
140 if batch is not None:
141 return [batch]
142 if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
143 batches = []
144 for item in value:
145 batch = cls._batch_mapping(item)
146 if batch is not None:
147 batches.append(batch)
148 return batches
149 return []
151 def _metric_group(self) -> Any:
152 """Return the DP+CP process group used by loss normalization."""
153 dp_cp_mesh = getattr(self.trainer.mesh, "dp_cp_mesh", None)
154 if dp_cp_mesh is None:
155 return None
156 return dp_cp_mesh.get_group()
158 def _reduce(self, value: Any, op: str) -> float:
159 """Reduce one scalar metric, with a single-process no-op fallback."""
160 if get_world_size_safe() <= 1:
161 return float(value)
162 reduced = all_reduce(value, op=op, group=self._metric_group())
163 return float(reduced)
165 def _accumulate_batches(self, value: Any) -> None:
166 """Accumulate batch metrics without retaining input tensor references."""
167 for batch in self._micro_batches(value):
168 token_count = self._batch_tokens(batch)
169 if callable(getattr(token_count, "detach", None)):
170 token_count = token_count.detach()
171 if self._local_step_tokens is None:
172 if callable(getattr(token_count, "clone", None)):
173 token_count = token_count.clone()
174 self._local_step_tokens = token_count
175 else:
176 self._local_step_tokens = self._local_step_tokens + token_count
177 self._local_step_samples += self._batch_samples(batch)
179 def _global_samples(self) -> int:
180 """Reduce samples across DP+CP while removing CP replicas."""
181 cp_size = int(getattr(self.trainer.mesh, "cp_size", 1))
182 if cp_size < 1:
183 raise ValueError(f"mesh.cp_size must be positive, but got {cp_size}")
184 reduced_samples = self._reduce(self._local_step_samples, op="sum")
185 global_samples = reduced_samples / cp_size
186 if not global_samples.is_integer():
187 raise ValueError(
188 "Reduced sample count must be divisible by cp_size, "
189 f"but got reduced_samples={reduced_samples} and cp_size={cp_size}"
190 )
191 return int(global_samples)
193 def _current_lr(self) -> float:
194 """Return the maximum learning rate across scheduler or optimizer groups."""
195 schedulers = self.trainer.lr_scheduler
196 if schedulers is not None:
197 scheduler_list = schedulers if isinstance(schedulers, list) else [schedulers]
198 learning_rates = []
199 for scheduler in scheduler_list:
200 for learning_rate in scheduler.get_last_lr():
201 learning_rates.append(float(learning_rate))
202 if learning_rates:
203 return max(learning_rates)
205 optimizers = self.trainer.optimizer
206 optimizer_list = optimizers if isinstance(optimizers, list) else [optimizers]
207 learning_rates = []
208 for optimizer in optimizer_list:
209 if optimizer is None:
210 continue
211 for param_group in optimizer.param_groups:
212 learning_rates.append(float(param_group["lr"]))
213 return max(learning_rates, default=0.0)
215 def _memory_metrics(self) -> dict[str, float]:
216 """Collect maximum accelerator memory metrics, if available."""
217 if get_device_type() == "cpu":
218 return {}
219 device = get_torch_device()
220 allocated = self._reduce(device.max_memory_allocated(), op="max")
221 reserved = self._reduce(device.max_memory_reserved(), op="max")
222 gibibyte = 1024 ** 3
223 return {
224 "memory/device_max_allocated_gb": allocated / gibibyte,
225 "memory/device_max_reserved_gb": reserved / gibibyte,
226 }
228 def state_dict(self) -> dict[str, int]:
229 """Return cumulative metric state for future checkpoint integration."""
230 return {
231 "consumed_tokens": self._consumed_tokens,
232 "consumed_samples": self._consumed_samples,
233 }
235 def load_state_dict(self, state_dict: dict[str, int]) -> None:
236 """Restore cumulative metric state.
238 Args:
239 state_dict: Mapping produced by :meth:`state_dict`.
241 Raises:
242 ValueError: If required counters are missing or negative.
243 """
244 try:
245 consumed_tokens = int(state_dict["consumed_tokens"])
246 consumed_samples = int(state_dict["consumed_samples"])
247 except (KeyError, TypeError, ValueError) as exc:
248 raise ValueError(
249 "EnvironMeterCallback state must contain integer consumed_tokens and consumed_samples"
250 ) from exc
251 if consumed_tokens < 0 or consumed_samples < 0:
252 raise ValueError("EnvironMeterCallback cumulative counters must be non-negative")
253 self._consumed_tokens = consumed_tokens
254 self._consumed_samples = consumed_samples
256 def on_step_begin(
257 self,
258 state: TrainerState,
259 micro_batches: list[dict[str, Any]] | None = None,
260 **kwargs: Any,
261 ) -> None:
262 """Start timing and count local input tokens and samples.
264 Args:
265 state: Current training progress.
266 micro_batches: Batches available before the step starts.
267 **kwargs: Unused callback context.
268 """
269 del state, kwargs
270 self._local_step_tokens = None
271 self._local_step_samples = 0
272 self._step_start_time = time.perf_counter()
273 self._accumulate_batches(micro_batches)
275 def on_micro_step_begin(
276 self,
277 state: TrainerState,
278 micro_batch: dict[str, Any],
279 **kwargs: Any,
280 ) -> None:
281 """Accumulate metrics for one prepared micro-batch.
283 Args:
284 state: Current training progress.
285 micro_batch: Prepared inputs and lightweight metric metadata.
286 **kwargs: Unused callback context.
287 """
288 del state, kwargs
289 self._accumulate_batches(micro_batch)
291 def on_step_end(
292 self,
293 state: TrainerState,
294 loss: float,
295 loss_dict: dict[str, float] | None,
296 grad_norm: float,
297 **kwargs: Any,
298 ) -> None:
299 """Reduce and publish metrics for one completed optimizer step.
301 Args:
302 state: Current training progress.
303 loss: Aggregated loss for the optimizer step.
304 loss_dict: Named loss values for the optimizer step.
305 grad_norm: Gradient norm measured before the optimizer update.
306 **kwargs: Unused callback context.
307 """
308 del state, kwargs
309 step_time = max(time.perf_counter() - self._step_start_time, 0.0)
310 global_step_time = self._reduce(step_time, op="max")
311 local_step_tokens = 0 if self._local_step_tokens is None else self._local_step_tokens
312 global_tokens = int(self._reduce(local_step_tokens, op="sum"))
313 global_samples = self._global_samples()
314 self._local_step_tokens = None
315 self._local_step_samples = 0
316 self._consumed_tokens += global_tokens
317 self._consumed_samples += global_samples
319 train_metrics = {
320 "training/total_loss": self._reduce(self._scalar(loss, "total_loss"), op="mean"),
321 "training/grad_norm": self._reduce(self._scalar(grad_norm, "grad_norm"), op="mean"),
322 "training/lr": self._current_lr(),
323 }
324 for name, value in sorted((loss_dict or {}).items()):
325 metric_name = name if name.startswith("training/") else f"training/{name}"
326 train_metrics[metric_name] = self._reduce(self._scalar(value, name), op="mean")
328 tokens_per_second = global_tokens / global_step_time if global_step_time > 0 else 0.0
329 env_metrics = {
330 **train_metrics,
331 "performance/step_time": global_step_time,
332 "performance/tokens_per_second": tokens_per_second,
333 "data/step_tokens": float(global_tokens),
334 "data/consumed_tokens": float(self._consumed_tokens),
335 "data/step_samples": float(global_samples),
336 "data/consumed_samples": float(self._consumed_samples),
337 **self._memory_metrics(),
338 }
339 self.trainer.step_train_metrics = train_metrics
340 self.trainer.step_env_metrics = env_metrics