Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / tools / model_integration / tensor_probes.py: 77%
284 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"""Tensor summaries, input identity, and parameter probe recorders."""
16# pylint: disable=forbidden-backend-import
18from __future__ import annotations
20import dataclasses
21import fnmatch
22import hashlib
23import json
24import math
25from typing import Any, Iterable, Mapping, Optional
27import torch
28import torch.distributed as dist
30from hyper_parallel.models.validation_spec import (
31 ModelValidationSpec,
32 ParameterProbeSpec,
33)
34from hyper_parallel.tools.model_integration.evidence_store import EvidenceStore
35from hyper_parallel.tools.model_integration.optimizer_layout import (
36 optimizer_state_for_parameter,
37 tensor_layout,
38)
40_SMALL_TENSOR_NUMEL = 8192
42def _tensor_numel(tensor: Any) -> int:
43 shape = getattr(tensor, "shape", ())
44 return math.prod(int(size) for size in shape)
47def _local_tensor(tensor: Any) -> torch.Tensor:
48 """Return one detached local shard without requesting redistribution."""
49 to_local = getattr(tensor, "to_local", None)
50 value = to_local() if callable(to_local) else tensor
51 return value.detach().float().cpu().contiguous()
54def _shard_identity(tensor: Any) -> tuple[tuple[int, int], ...]:
55 """Identify the logical shard while excluding replicated mesh axes."""
56 layout = getattr(tensor, "layout", None)
57 mesh = getattr(layout, "mesh", getattr(tensor, "device_mesh", None))
58 placements = getattr(layout, "placements", getattr(tensor, "placements", ()))
59 get_coordinate = getattr(mesh, "get_coordinate", None)
60 coordinate = get_coordinate() if callable(get_coordinate) else None
61 if coordinate is None:
62 return ()
63 shard_coordinates = []
64 for mesh_dim, placement in enumerate(placements or ()):
65 is_shard = getattr(placement, "is_shard", None)
66 is_ragged = getattr(placement, "is_ragged_shard", None)
67 if (callable(is_shard) and is_shard()) or (
68 callable(is_ragged) and is_ragged()
69 ):
70 shard_coordinates.append((mesh_dim, int(coordinate[mesh_dim])))
71 return tuple(shard_coordinates)
74def _local_shards_follow_flattened_order(tensor: Any) -> bool:
75 """Return whether rank-ordered shard values reconstruct the flattened tensor."""
76 layout = getattr(tensor, "layout", None)
77 placements = getattr(layout, "placements", getattr(tensor, "placements", ()))
78 shard_placements = []
79 for placement in placements or ():
80 is_shard = getattr(placement, "is_shard", None)
81 is_ragged = getattr(placement, "is_ragged_shard", None)
82 if (callable(is_shard) and is_shard()) or (
83 callable(is_ragged) and is_ragged()
84 ):
85 shard_placements.append(placement)
86 if not shard_placements:
87 return True
88 if len(shard_placements) != 1:
89 return False
90 placement = shard_placements[0]
91 return getattr(placement, "dim", None) == 0 and not hasattr(
92 placement,
93 "split_factor",
94 )
97def _local_tensor_summary(tensor: Any) -> dict[str, Any]:
98 """Summarize one local shard for cross-rank aggregation."""
99 value = _local_tensor(tensor).reshape(-1)
100 digest = hashlib.sha256(value.numpy().tobytes()).hexdigest()
101 return {
102 "rank": dist.get_rank() if dist.is_available() and dist.is_initialized() else 0,
103 "shard": _shard_identity(tensor),
104 "sha256": digest,
105 "numel": value.numel(),
106 "finite": bool(torch.isfinite(value).all()) if value.numel() else True,
107 "sum": float(value.double().sum()),
108 "squared_l2": float(torch.square(value.double()).sum()),
109 "min": float(value.min()) if value.numel() else None,
110 "max": float(value.max()) if value.numel() else None,
111 "values": value.tolist() if value.numel() <= _SMALL_TENSOR_NUMEL else None,
112 "global_shape": [int(size) for size in getattr(tensor, "shape", ())],
113 "dtype": str(getattr(tensor, "dtype", None)),
114 "layout": tensor_layout(tensor),
115 }
118def _gather_shard_summaries(
119 local_summary: Optional[dict[str, Any]],
120) -> tuple[list[dict[str, Any]], list[Optional[dict[str, Any]]]]:
121 """Gather optional metadata records and retain one value per logical shard."""
122 rank_summaries = [local_summary]
123 if dist.is_available() and dist.is_initialized():
124 rank_summaries = [None] * dist.get_world_size()
125 dist.all_gather_object(rank_summaries, local_summary)
126 by_shard: dict[tuple[tuple[int, int], ...], list[dict[str, Any]]] = {}
127 for summary in rank_summaries:
128 if summary is None:
129 continue
130 shard = tuple(tuple(item) for item in summary["shard"])
131 by_shard.setdefault(shard, []).append(summary)
132 shards = [
133 min(replicas, key=lambda item: item["rank"])
134 for replicas in by_shard.values()
135 ]
136 return shards, rank_summaries
139def _replica_digest_groups(
140 rank_summaries: list[Optional[dict[str, Any]]],
141) -> dict[tuple[Any, ...], set[str]]:
142 """Group per-rank payload digests by their logical shard identity."""
143 replica_groups: dict[tuple[Any, ...], set[str]] = {}
144 for rank_summary in rank_summaries:
145 if rank_summary is None:
146 continue
147 shard = tuple(tuple(item) for item in rank_summary["shard"])
148 replica_groups.setdefault(shard, set()).add(rank_summary["sha256"])
149 return replica_groups
152def _canonical_shard_digest(shards: list[dict[str, Any]]) -> str:
153 """Hash the ordered shard digests into one canonical tensor identity."""
154 return hashlib.sha256(
155 json.dumps(
156 [(summary["shard"], summary["sha256"]) for summary in shards],
157 sort_keys=True,
158 ).encode("utf-8")
159 ).hexdigest()
162def _flattened_shard_values(tensor: Any, shards: list[dict[str, Any]]) -> Optional[list[Any]]:
163 """Reconstruct full values only from flattened-order small shard sets."""
164 if not (
165 _local_shards_follow_flattened_order(tensor)
166 and sum(shard["numel"] for shard in shards) <= _SMALL_TENSOR_NUMEL
167 and all(shard["values"] is not None for shard in shards)
168 ):
169 return None
170 return [value for shard in shards for value in shard["values"]]
173def _global_shard_summary(
174 name: str,
175 tensor: Any,
176 shards: list[dict[str, Any]],
177 rank_summaries: list[Optional[dict[str, Any]]],
178) -> dict[str, Any]:
179 """Aggregate sorted shard summaries into one canonical global record."""
180 replica_groups = _replica_digest_groups(rank_summaries)
181 minima = [summary["min"] for summary in shards if summary["min"] is not None]
182 maxima = [summary["max"] for summary in shards if summary["max"] is not None]
183 reference = shards[0]
184 summary = {
185 "name": name,
186 "global_shape": reference["global_shape"],
187 "canonical_shape": reference["global_shape"],
188 "dtype": reference["dtype"],
189 "sha256": _canonical_shard_digest(shards),
190 "finite": all(shard["finite"] for shard in shards),
191 "sum": sum(shard["sum"] for shard in shards),
192 "l2": math.sqrt(sum(shard["squared_l2"] for shard in shards)),
193 "min": min(minima) if minima else None,
194 "max": max(maxima) if maxima else None,
195 "layout": reference["layout"],
196 "replica": {
197 "logical_shards": len(replica_groups),
198 "replicas_equal": all(len(digests) == 1 for digests in replica_groups.values()),
199 },
200 "missing_ranks": [
201 rank for rank, rank_summary in enumerate(rank_summaries) if rank_summary is None
202 ],
203 }
204 values = _flattened_shard_values(tensor, shards)
205 if values is not None:
206 summary["values"] = values
207 return summary
210def _tensor_summary(name: str, tensor: Any) -> Optional[dict[str, Any]]:
211 """Build one global summary while allowing rank-local lazy state absence."""
212 local_summary = None if tensor is None else _local_tensor_summary(tensor)
213 shards, rank_summaries = _gather_shard_summaries(local_summary)
214 if not shards:
215 return None
216 shards.sort(key=lambda item: item["shard"])
217 return _global_shard_summary(name, tensor, shards, rank_summaries)
220def _global_tensor_state_names(local_names: Iterable[str]) -> tuple[str, ...]:
221 """Return a rank-consistent tensor-state key order for collective probes."""
222 local_names = tuple(sorted(local_names))
223 rank_names = [local_names]
224 if dist.is_available() and dist.is_initialized():
225 rank_names = [None] * dist.get_world_size()
226 dist.all_gather_object(rank_names, local_names)
227 return tuple(sorted({name for names in rank_names for name in (names or ())}))
230class InputIdentityRecorder:
231 """Hash global step inputs from deterministic per-rank tensor payloads."""
233 def __init__(
234 self,
235 mesh_context: Any = None,
236 cp_replicated_forward_fields: Iterable[str] = (),
237 ) -> None:
238 """Configure mesh-aware canonicalization for one Trainer run."""
239 self.mesh_context = mesh_context
240 self.cp_replicated_forward_fields = set(cp_replicated_forward_fields)
241 self._entries: list[dict[str, Any]] = []
242 self._micro_steps = 0
244 def begin_step(self) -> None:
245 """Reset the current optimizer-step identity."""
246 self._entries = []
247 self._micro_steps = 0
249 @staticmethod
250 def _tensor_leaves(value: Any) -> list[Any]:
251 """Find tensors inside mappings, sequences, and metadata dataclasses."""
252 if torch.is_tensor(value):
253 return [value]
254 if dataclasses.is_dataclass(value) and not isinstance(value, type):
255 return [
256 tensor
257 for dataclass_field in dataclasses.fields(value)
258 for field_value in (getattr(value, dataclass_field.name),)
259 for tensor in InputIdentityRecorder._tensor_leaves(field_value)
260 ]
261 if isinstance(value, Mapping):
262 return [
263 tensor
264 for name in sorted(value)
265 for tensor in InputIdentityRecorder._tensor_leaves(value[name])
266 ]
267 if isinstance(value, (list, tuple)):
268 return [
269 tensor
270 for child in value
271 for tensor in InputIdentityRecorder._tensor_leaves(child)
272 ]
273 return []
275 def record(self, model_inputs: Mapping[str, Any], loss_inputs: Mapping[str, Any]) -> None:
276 """Add model/loss tensor fields without changing their device or values."""
277 for owner, values in (("model", model_inputs), ("loss", loss_inputs)):
278 for field_name in sorted(values):
279 flattened = self._tensor_leaves(values[field_name])
280 for value_index, value in enumerate(flattened):
281 local = value.detach()
282 to_local = getattr(local, "to_local", None)
283 if callable(to_local):
284 local = to_local()
285 self._entries.append(
286 {
287 "micro_step": self._micro_steps,
288 "owner": owner,
289 "field": field_name,
290 "value_index": value_index,
291 "cp_replicated": (
292 owner == "model"
293 and field_name in self.cp_replicated_forward_fields
294 ),
295 "tensor": local.cpu().contiguous(),
296 }
297 )
298 self._micro_steps += 1
300 def _coordinates(self) -> dict[str, int]:
301 return {
302 name: int(getattr(self.mesh_context, f"{name}_rank", 0))
303 for name in ("dp", "cp", "tp", "pp")
304 }
306 @staticmethod
307 def _hash_tensor(value: torch.Tensor) -> str:
308 hasher = hashlib.sha256()
309 hasher.update(f"{tuple(value.shape)}:{value.dtype}".encode())
310 byte_view = value.contiguous().view(torch.uint8)
311 hasher.update(byte_view.numpy().tobytes())
312 return hasher.hexdigest()
314 @staticmethod
315 def _require_input_payloads(rank_payloads: list[dict[str, Any]]) -> None:
316 """Reject collectives polluted by non-input rank payloads."""
317 malformed_ranks = [
318 rank
319 for rank, payload in enumerate(rank_payloads)
320 if not isinstance(payload, Mapping)
321 or "coordinates" not in payload
322 or "entries" not in payload
323 ]
324 if malformed_ranks:
325 raise RuntimeError(
326 "input-identity collective received non-input payloads from ranks "
327 f"{malformed_ranks}; parameter probes must execute a rank-consistent "
328 "collective sequence"
329 )
331 @staticmethod
332 def _group_input_entries(
333 rank_payloads: list[dict[str, Any]],
334 ) -> dict[tuple[Any, ...], list[tuple[int, torch.Tensor, bool]]]:
335 """Group entry shards by logical field across tp0/pp0 ranks."""
336 selected = [
337 payload
338 for payload in rank_payloads
339 if payload["coordinates"]["tp"] == 0 and payload["coordinates"]["pp"] == 0
340 ]
341 grouped: dict[tuple[Any, ...], list[tuple[int, torch.Tensor, bool]]] = {}
342 for payload in selected:
343 coordinates = payload["coordinates"]
344 for entry in payload["entries"]:
345 key = (
346 entry["micro_step"],
347 coordinates["dp"],
348 entry["owner"],
349 entry["field"],
350 entry["value_index"],
351 )
352 grouped.setdefault(key, []).append(
353 (
354 coordinates["cp"],
355 entry["tensor"],
356 bool(entry.get("cp_replicated", False)),
357 )
358 )
359 return grouped
361 @staticmethod
362 def _merge_entry_shards(
363 shards: list[tuple[int, torch.Tensor, bool]],
364 ) -> torch.Tensor:
365 """Rebuild one logical field tensor from ordered CP shards."""
366 ordered_shards = sorted(shards, key=lambda item: item[0])
367 ordered = [tensor for _, tensor, _ in ordered_shards]
368 tensor = ordered[0]
369 cp_replicated = all(replicated for _, _, replicated in ordered_shards)
370 if cp_replicated and any(
371 not torch.equal(tensor, replica) for replica in ordered[1:]
372 ):
373 cp_replicated = False
374 if len(ordered) > 1 and tensor.ndim > 0 and not cp_replicated:
375 try:
376 tensor = torch.cat(ordered, dim=-1)
377 except RuntimeError:
378 tensor = torch.cat([value.reshape(-1) for value in ordered])
379 return tensor
381 @classmethod
382 def _entry_component_hashes(
383 cls,
384 grouped: dict[tuple[Any, ...], list[tuple[int, torch.Tensor, bool]]],
385 ) -> list[str]:
386 """Hash every logical field occurrence independently of rank order."""
387 component_hashes = []
388 for key, shards in grouped.items():
389 tensor = cls._merge_entry_shards(shards)
390 field_identity = f"{key[2]}:{key[3]}:{key[4]}"
391 if tensor.ndim > 0:
392 values = tensor.reshape(tensor.shape[0], -1)
393 component_hashes.extend(
394 f"{field_identity}:{cls._hash_tensor(value)}" for value in values
395 )
396 else:
397 component_hashes.append(f"{field_identity}:{cls._hash_tensor(tensor)}")
398 return component_hashes
400 @classmethod
401 def _canonical_global_hash(cls, rank_payloads: list[dict[str, Any]]) -> str:
402 cls._require_input_payloads(rank_payloads)
403 grouped = cls._group_input_entries(rank_payloads)
404 component_hashes = cls._entry_component_hashes(grouped)
405 hasher = hashlib.sha256()
406 for component_hash in sorted(component_hashes):
407 hasher.update(component_hash.encode())
408 return hasher.hexdigest()
410 def finish_step(self) -> dict[str, Any]:
411 """Return rank and global digests for the optimizer step."""
412 local_payload = {
413 "coordinates": self._coordinates(),
414 "entries": self._entries,
415 }
416 rank_payloads = [local_payload]
417 if dist.is_available() and dist.is_initialized():
418 rank_payloads = [None] * dist.get_world_size()
419 dist.all_gather_object(rank_payloads, local_payload)
420 return {
421 "global_sha256": self._canonical_global_hash(rank_payloads),
422 "rank_count": len(rank_payloads),
423 "micro_steps": self._micro_steps,
424 }
427class ParameterProbeRecorder:
428 """Capture canonical parameter/main-gradient/update evidence by semantic watchlist."""
430 def __init__(
431 self,
432 model: Any,
433 optimizer: Any,
434 evidence_store: EvidenceStore,
435 *,
436 rank: int,
437 validation_spec: Optional[ModelValidationSpec] = None,
438 include_small_parameters_numel_le: int = _SMALL_TENSOR_NUMEL,
439 representative_parameter_names: Iterable[str] = (),
440 ) -> None:
441 """Resolve a stable watchlist from the final model and adapter declarations."""
442 self.model = model
443 self.optimizer = optimizer
444 self.evidence_store = evidence_store
445 self.rank = rank
446 self.validation_spec = validation_spec
447 self.include_small_parameters_numel_le = include_small_parameters_numel_le
448 self.representative_parameter_names = set(representative_parameter_names)
449 self._named_parameters = dict(model.named_parameters(remove_duplicate=False))
450 self._probes = self._resolve_probes()
451 self._validate_required_patterns()
452 self._step = 0
454 def _resolve_probes(self) -> dict[str, ParameterProbeSpec]:
455 probes = {}
456 declared = self.validation_spec.parameter_probes if self.validation_spec else ()
457 replaced_module_fqns = tuple(
458 getattr(self.model, "_hp_replaced_module_fqns", ()) or ()
459 )
460 for parameter_name, parameter in self._named_parameters.items():
461 main_parameter = getattr(parameter, "main_param", parameter)
462 replacement_parameter = any(
463 parameter_name == module_fqn
464 or parameter_name.startswith(f"{module_fqn}.")
465 for module_fqn in replaced_module_fqns
466 )
467 if (
468 replacement_parameter
469 or parameter_name in self.representative_parameter_names
470 or _tensor_numel(main_parameter) <= self.include_small_parameters_numel_le
471 ):
472 probes[parameter_name] = ParameterProbeSpec(
473 match=parameter_name,
474 required=True,
475 stages=(
476 "model_before",
477 "main_before",
478 "main_grad_before_clip",
479 "main_grad_after_clip",
480 "main_after_optimizer",
481 "model_after_copy_back",
482 "optimizer_state",
483 ),
484 )
485 for probe in declared:
486 if fnmatch.fnmatchcase(parameter_name, probe.match):
487 probes[parameter_name] = probe
488 return dict(sorted(probes.items()))
490 def _validate_required_patterns(self) -> None:
491 for probe in self.validation_spec.parameter_probes if self.validation_spec else ():
492 if probe.required and not any(
493 fnmatch.fnmatchcase(name, probe.match) for name in self._named_parameters
494 ):
495 raise ValueError(
496 f"required parameter probe {probe.match!r} matched no final parameter"
497 )
499 @property
500 def watched_parameter_names(self) -> tuple[str, ...]:
501 """Return stable final FQNs in capture order."""
502 return tuple(self._probes)
504 def begin_step(self, step: int) -> None:
505 """Start evidence for one optimizer step."""
506 self._step = step
507 self.capture("model_before")
508 self.capture("main_before")
510 @staticmethod
511 def _stage_tensor(parameter: Any, stage: str) -> Any:
512 main_parameter = getattr(parameter, "main_param", parameter)
513 if stage in ("model_before", "model_after_copy_back"):
514 return parameter
515 if stage in ("main_before", "main_after_optimizer"):
516 return main_parameter
517 if stage in ("main_grad_before_clip", "main_grad_after_clip"):
518 main_gradient = getattr(parameter, "main_grad", None)
519 return main_gradient if main_gradient is not None else getattr(main_parameter, "grad", None)
520 return None
522 def _optimizer_state(self, parameter: Any) -> dict[str, Any]:
523 main_parameter = (
524 None if parameter is None else getattr(parameter, "main_param", parameter)
525 )
526 state = (
527 {} if main_parameter is None
528 else optimizer_state_for_parameter(self.optimizer, main_parameter)
529 )
530 state_names = _global_tensor_state_names(
531 name for name, value in state.items() if torch.is_tensor(value)
532 )
533 summaries = {}
534 for name in state_names:
535 value = state.get(name)
536 summary = _tensor_summary(name, value if torch.is_tensor(value) else None)
537 if summary is not None:
538 summaries[name] = summary
539 return summaries
541 def capture(self, stage: str) -> None:
542 """Capture global summaries and persist one canonical rank-zero copy."""
543 records = []
544 for parameter_name, probe in self._probes.items():
545 if stage not in probe.stages:
546 continue
547 parameter = self._named_parameters.get(parameter_name)
548 if stage == "optimizer_state":
549 summary = self._optimizer_state(parameter)
550 if not summary:
551 continue
552 record = {"name": parameter_name, "optimizer_state": summary}
553 else:
554 tensor = None if parameter is None else self._stage_tensor(parameter, stage)
555 record = _tensor_summary(parameter_name, tensor)
556 if record is None:
557 continue
558 if not probe.check_replica_consistency:
559 record.pop("replica", None)
560 record.update(
561 {
562 "step": self._step,
563 "stage": stage,
564 "rank": self.rank,
565 "comparison": probe.compare,
566 }
567 )
568 records.append(record)
569 if self.rank != 0:
570 return
571 for record in records:
572 self.evidence_store.append_jsonl(
573 f"parameter_probes/rank{self.rank}.jsonl",
574 record,
575 )
578__all__ = [
579 "InputIdentityRecorder",
580 "ParameterProbeRecorder",
581]