Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / environ_meter_callback.py: 82%
243 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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.trainer.runtime.flops import estimate_step_flops
26from hyper_parallel.data.constants import IGNORE_INDEX
27from hyper_parallel.trainer.runtime.device import get_device_type, get_torch_device
29from .base import Callback, TrainerState
32class EnvironMeterCallback(Callback):
33 """Collect structured training, throughput, and memory metrics.
35 The callback is the single producer of ``trainer.step_train_metrics`` and
36 ``trainer.step_env_metrics``. Presentation and remote logging callbacks
37 consume those dictionaries without recalculating or reducing metrics.
38 """
40 def __init__(self, trainer: Any, *, microbatch_mean_metrics: Sequence[str] = ()) -> None:
41 """Initialize per-step and cumulative counters.
43 Args:
44 trainer: Trainer that owns the callback lifecycle.
45 microbatch_mean_metrics: Raw metric names to average over observed
46 micro-batches before the existing DP+CP mean reduction. Other
47 metrics retain the step values supplied by the Trainer.
48 """
49 super().__init__(trainer)
50 self._step_start_time = 0.0
51 self._local_step_tokens: Any = None
52 self._local_step_input_tokens: Any = None
53 self._local_step_seq_len_sq: Any = None
54 self._local_step_samples = 0
55 self._microbatch_mean_metrics = tuple(microbatch_mean_metrics)
56 self._local_mean_metrics: dict[str, tuple[Any, int]] = {}
57 self._consumed_tokens = 0
58 self._consumed_samples = 0
59 self.trainer.step_train_metrics = {}
60 self.trainer.step_env_metrics = {}
62 @staticmethod
63 def _scalar(value: Any, name: str) -> float:
64 """Convert a scalar or scalar tensor-like value to ``float``.
66 Args:
67 value: Scalar value to convert.
68 name: Metric name used in validation errors.
70 Returns:
71 Converted floating-point value.
73 Raises:
74 ValueError: If the value cannot be converted to a scalar float.
75 """
76 item = getattr(value, "item", None)
77 if callable(item):
78 value = item()
79 try:
80 return float(value)
81 except (TypeError, ValueError) as exc:
82 raise ValueError(f"Metric {name!r} must be scalar, but got {value!r}") from exc
84 @staticmethod
85 def _tensor_numel(value: Any) -> int | None:
86 """Return ``value.numel()`` when it exposes a tensor-like interface."""
87 numel = getattr(value, "numel", None)
88 if not callable(numel):
89 return None
90 return int(numel())
92 @classmethod
93 def _batch_tokens(cls, batch: Mapping[str, Any]) -> Any:
94 """Count text tokens without synchronizing a device scalar to the host."""
95 token_count = batch.get("token_count")
96 if token_count is not None:
97 return token_count
99 labels = batch.get("labels")
100 if labels is not None and callable(getattr(labels, "sum", None)):
101 return (labels != IGNORE_INDEX).sum()
103 attention_mask = batch.get("attention_mask")
104 attention_mask_shape = getattr(attention_mask, "shape", ())
105 if (
106 len(attention_mask_shape) <= 2
107 and attention_mask is not None
108 and callable(getattr(attention_mask, "sum", None))
109 ):
110 return attention_mask.sum()
112 input_ids = batch.get("input_ids")
113 input_numel = cls._tensor_numel(input_ids)
114 if input_numel is not None:
115 return input_numel
116 return 0
118 @classmethod
119 def _batch_input_geometry(cls, batch: Mapping[str, Any]) -> tuple[Any, Any] | None:
120 """Return model-input tokens and the sum of squared sequence lengths.
122 The loss token count drops masked prompt tokens, but the model still
123 runs attention and the FFN over them, so FLOPs need their own count:
124 the tokens that reach attention plus ``sum(S ** 2)`` over the real
125 attention boundaries.
127 Every returned quantity is local to this rank, because
128 :meth:`on_step_end` reduces them over DP+CP itself.
130 Args:
131 batch: One micro-batch mapping.
133 Returns:
134 ``(input_tokens, sum_seq_len_sq)`` for the sequences that reach
135 attention, or ``None`` when the batch exposes no input geometry.
136 """
137 cu_seq_lens = batch.get("cu_seq_lens")
138 if cu_seq_lens is not None and callable(getattr(cu_seq_lens, "diff", None)):
139 local_tokens = cls._tensor_numel(batch.get("input_ids"))
140 lengths = cu_seq_lens.diff()
141 covered = int(lengths.sum())
142 # Packed boundaries are only usable when they describe the tokens
143 # this rank actually holds. Under context parallelism the packed
144 # boundaries may still span the global sequence, and then the local
145 # ``input_ids`` fallbacks below are the honest answer.
146 if covered > 0 and (local_tokens is None or covered == local_tokens):
147 return lengths.sum(), (lengths * lengths).sum()
149 attention_mask = batch.get("attention_mask")
150 if (
151 attention_mask is not None
152 and len(getattr(attention_mask, "shape", ())) == 2
153 and callable(getattr(attention_mask, "sum", None))
154 ):
155 lengths = attention_mask.sum(dim=-1)
156 return lengths.sum(), (lengths * lengths).sum()
158 shape = getattr(batch.get("input_ids"), "shape", None)
159 if shape is None or len(shape) < 2:
160 return None
161 rows = int(shape[0])
162 seq_length = int(shape[-1])
163 if rows <= 0 or seq_length <= 0:
164 return None
165 return rows * seq_length, rows * seq_length * seq_length
167 @staticmethod
168 def _batch_samples(batch: Mapping[str, Any]) -> int:
169 """Count logical samples in one micro-batch."""
170 value = batch.get("input_ids")
171 if value is None:
172 value = batch.get("labels")
173 shape = getattr(value, "shape", None)
174 if shape is None or len(shape) == 0:
175 return 0
176 if len(shape) == 1:
177 return 1
178 return int(shape[0])
180 @staticmethod
181 def _batch_mapping(value: Any) -> Mapping[str, Any] | None:
182 """Return metric inputs for a mapping or prepared runtime batch."""
183 if isinstance(value, Mapping):
184 return value
185 loss_count_inputs = getattr(value, "loss_count_inputs", None)
186 if not callable(loss_count_inputs):
187 return None
188 metric_inputs = loss_count_inputs()
189 return metric_inputs if isinstance(metric_inputs, Mapping) else None
191 @classmethod
192 def _micro_batches(cls, value: Any) -> list[Mapping[str, Any]]:
193 """Normalize callback input into a list of mapping micro-batches."""
194 if value is None:
195 return []
196 batch = cls._batch_mapping(value)
197 if batch is not None:
198 return [batch]
199 if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
200 batches = []
201 for item in value:
202 batch = cls._batch_mapping(item)
203 if batch is not None:
204 batches.append(batch)
205 return batches
206 return []
208 def _metric_group(self) -> Any:
209 """Return the DP+CP process group used by loss normalization."""
210 dp_cp_mesh = getattr(self.trainer.mesh, "dp_cp_mesh", None)
211 if dp_cp_mesh is None:
212 return None
213 return dp_cp_mesh.get_group()
215 def _reduce(self, value: Any, op: str) -> float:
216 """Reduce one scalar metric, with a single-process no-op fallback."""
217 if get_world_size_safe() <= 1:
218 return float(value)
219 reduced = all_reduce(value, op=op, group=self._metric_group())
220 return float(reduced)
222 @staticmethod
223 def _accumulate_scalar(current: Any, value: Any) -> Any:
224 """Add one batch scalar to a running per-step total.
226 Args:
227 current: Running total, or ``None`` before the first batch.
228 value: Batch value to add; may be a device scalar.
230 Returns:
231 Updated running total that never aliases the batch tensor.
232 """
233 detached = value.detach() if callable(getattr(value, "detach", None)) else value
234 if current is None:
235 return detached.clone() if callable(getattr(detached, "clone", None)) else detached
236 return current + detached
238 def _accumulate_batches(self, value: Any) -> None:
239 """Accumulate batch metrics without retaining input tensor references."""
240 for batch in self._micro_batches(value):
241 self._local_step_tokens = self._accumulate_scalar(
242 self._local_step_tokens, self._batch_tokens(batch)
243 )
244 geometry = self._batch_input_geometry(batch)
245 if geometry is not None:
246 input_tokens, seq_len_sq = geometry
247 self._local_step_input_tokens = self._accumulate_scalar(
248 self._local_step_input_tokens, input_tokens
249 )
250 self._local_step_seq_len_sq = self._accumulate_scalar(
251 self._local_step_seq_len_sq, seq_len_sq
252 )
253 self._local_step_samples += self._batch_samples(batch)
255 def _cp_size(self) -> int:
256 """Return the validated context-parallel size of the trainer mesh."""
257 cp_size = int(getattr(self.trainer.mesh, "cp_size", 1))
258 if cp_size < 1:
259 raise ValueError(f"mesh.cp_size must be positive, but got {cp_size}")
260 return cp_size
262 def _global_samples(self) -> int:
263 """Reduce samples across DP+CP while removing CP replicas."""
264 cp_size = self._cp_size()
265 reduced_samples = self._reduce(self._local_step_samples, op="sum")
266 global_samples = reduced_samples / cp_size
267 if not global_samples.is_integer():
268 raise ValueError(
269 "Reduced sample count must be divisible by cp_size, "
270 f"but got reduced_samples={reduced_samples} and cp_size={cp_size}"
271 )
272 return int(global_samples)
274 def _current_lr(self) -> float:
275 """Return the maximum learning rate across scheduler or optimizer groups."""
276 schedulers = self.trainer.lr_scheduler
277 if schedulers is not None:
278 scheduler_list = schedulers if isinstance(schedulers, list) else [schedulers]
279 learning_rates = []
280 for scheduler in scheduler_list:
281 for learning_rate in scheduler.get_last_lr():
282 learning_rates.append(float(learning_rate))
283 if learning_rates:
284 return max(learning_rates)
286 optimizers = self.trainer.optimizer
287 optimizer_list = optimizers if isinstance(optimizers, list) else [optimizers]
288 learning_rates = []
289 for optimizer in optimizer_list:
290 if optimizer is None:
291 continue
292 for param_group in optimizer.param_groups:
293 learning_rates.append(float(param_group["lr"]))
294 return max(learning_rates, default=0.0)
296 def _memory_metrics(self) -> dict[str, float]:
297 """Collect maximum accelerator memory metrics, if available."""
298 if get_device_type() == "cpu":
299 return {}
300 device = get_torch_device()
301 allocated = self._reduce(device.max_memory_allocated(), op="max")
302 reserved = self._reduce(device.max_memory_reserved(), op="max")
303 gibibyte = 1024 ** 3
304 return {
305 "memory/device_max_allocated_gb": allocated / gibibyte,
306 "memory/device_max_reserved_gb": reserved / gibibyte,
307 }
309 def state_dict(self) -> dict[str, int]:
310 """Return cumulative metric state for future checkpoint integration."""
311 return {
312 "consumed_tokens": self._consumed_tokens,
313 "consumed_samples": self._consumed_samples,
314 }
316 def load_state_dict(self, state_dict: dict[str, int]) -> None:
317 """Restore cumulative metric state.
319 Args:
320 state_dict: Mapping produced by :meth:`state_dict`.
322 Raises:
323 ValueError: If required counters are missing or negative.
324 """
325 try:
326 consumed_tokens = int(state_dict["consumed_tokens"])
327 consumed_samples = int(state_dict["consumed_samples"])
328 except (KeyError, TypeError, ValueError) as exc:
329 raise ValueError(
330 "EnvironMeterCallback state must contain integer consumed_tokens and consumed_samples"
331 ) from exc
332 if consumed_tokens < 0 or consumed_samples < 0:
333 raise ValueError("EnvironMeterCallback cumulative counters must be non-negative")
334 self._consumed_tokens = consumed_tokens
335 self._consumed_samples = consumed_samples
337 def on_step_begin(
338 self,
339 state: TrainerState,
340 micro_batches: list[dict[str, Any]] | None = None,
341 **kwargs: Any,
342 ) -> None:
343 """Start timing and count local input tokens and samples.
345 Args:
346 state: Current training progress.
347 micro_batches: Batches available before the step starts.
348 **kwargs: Unused callback context.
349 """
350 del state, kwargs
351 self._local_step_tokens = None
352 self._local_step_input_tokens = None
353 self._local_step_seq_len_sq = None
354 self._local_step_samples = 0
355 self._step_start_time = time.perf_counter()
356 self._local_mean_metrics.clear()
357 self._accumulate_batches(micro_batches)
359 def on_micro_step_begin(
360 self,
361 state: TrainerState,
362 micro_batch: dict[str, Any],
363 **kwargs: Any,
364 ) -> None:
365 """Accumulate metrics for one prepared micro-batch.
367 Args:
368 state: Current training progress.
369 micro_batch: Prepared inputs and lightweight metric metadata.
370 **kwargs: Unused callback context.
371 """
372 del state, kwargs
373 self._accumulate_batches(micro_batch)
375 def on_micro_step_end(self, state: TrainerState, loss_dict: dict[str, Any], **kwargs: Any) -> None:
376 """Accumulate explicitly registered micro-batch means without host reads."""
377 del state, kwargs
378 for name in self._microbatch_mean_metrics:
379 value = loss_dict.get(name)
380 if value is not None:
381 total, count = self._local_mean_metrics.get(name, (None, 0))
382 self._local_mean_metrics[name] = (self._accumulate_scalar(total, value), count + 1)
384 def on_step_end(
385 self,
386 state: TrainerState,
387 loss: float,
388 loss_dict: dict[str, float] | None,
389 grad_norm: float,
390 **kwargs: Any,
391 ) -> None:
392 """Reduce and publish metrics for one completed optimizer step.
394 Args:
395 state: Current training progress.
396 loss: Aggregated loss for the optimizer step.
397 loss_dict: Named loss values for the optimizer step.
398 grad_norm: Gradient norm measured before the optimizer update.
399 **kwargs: Unused callback context.
400 """
401 del state, kwargs
402 step_time = max(time.perf_counter() - self._step_start_time, 0.0)
403 global_step_time = self._reduce(step_time, op="max")
404 local_step_tokens = 0 if self._local_step_tokens is None else self._local_step_tokens
405 global_tokens = int(self._reduce(local_step_tokens, op="sum"))
406 global_samples = self._global_samples()
407 cp_size = self._cp_size()
408 local_input_tokens = (
409 0 if self._local_step_input_tokens is None else self._local_step_input_tokens
410 )
411 local_seq_len_sq = (
412 0 if self._local_step_seq_len_sq is None else self._local_step_seq_len_sq
413 )
414 # Context parallelism slices a sequence along its length, so each rank
415 # only sees ``(S / cp) ** 2`` while the real attention cost is ``S ** 2``.
416 # Summing over DP+CP therefore needs one factor of ``cp_size`` to
417 # restore the per-sequence square. Input tokens need no correction
418 # because the slices add back up to ``S``.
419 global_input_tokens = self._reduce(local_input_tokens, op="sum")
420 global_seq_len_sq = self._reduce(local_seq_len_sq, op="sum") * cp_size
421 self._local_step_tokens = None
422 self._local_step_input_tokens = None
423 self._local_step_seq_len_sq = None
424 self._local_step_samples = 0
425 self._consumed_tokens += global_tokens
426 self._consumed_samples += global_samples
428 train_metrics = {
429 "training/total_loss": self._reduce(self._scalar(loss, "total_loss"), op="mean"),
430 "training/grad_norm": self._reduce(self._scalar(grad_norm, "grad_norm"), op="mean"),
431 "training/lr": self._current_lr(),
432 }
433 for name, value in sorted((loss_dict or {}).items()):
434 if name in self._local_mean_metrics:
435 total, count = self._local_mean_metrics[name]
436 value = total / count
437 metric_name = name if name.startswith("training/") else f"training/{name}"
438 train_metrics[metric_name] = self._reduce(self._scalar(value, name), op="mean")
439 self._local_mean_metrics.clear()
441 tokens_per_second = global_tokens / global_step_time if global_step_time > 0 else 0.0
442 env_metrics = {
443 **train_metrics,
444 "performance/step_time": global_step_time,
445 "performance/tokens_per_second": tokens_per_second,
446 "data/step_tokens": float(global_tokens),
447 "data/consumed_tokens": float(self._consumed_tokens),
448 "data/step_samples": float(global_samples),
449 "data/consumed_samples": float(self._consumed_samples),
450 **self._memory_metrics(),
451 }
453 training_config = getattr(getattr(self.trainer, "config", None), "training", None)
454 model_config = getattr(self.trainer, "model_config", None)
455 if (
456 getattr(training_config, "log_tflops", False)
457 and model_config is not None
458 and global_input_tokens > 0
459 and global_step_time > 0
460 ):
461 step_flops = estimate_step_flops(
462 model_config, global_input_tokens, global_seq_len_sq
463 )
464 if step_flops is not None:
465 env_metrics["performance/tflops_per_device"] = (
466 step_flops / (global_step_time * get_world_size_safe()) / 1e12
467 )
469 self.trainer.step_train_metrics = train_metrics
470 self.trainer.step_env_metrics = env_metrics