Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / model_integration.py: 25%
250 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"""Trainer-owned model-integration validation lifecycle."""
17from __future__ import annotations
19import os
20import time
21from pathlib import Path
22from typing import Any, Mapping, Optional, Protocol
24from hyper_parallel.tools.model_integration.checkpoint_coverage import build_checkpoint_coverage
25from hyper_parallel.tools.model_integration.data_contract import (
26 modality_gradient_parameters,
27 validate_data_contract,
28 validate_observed_forward_fields,
29)
30from hyper_parallel.tools.model_integration.evidence_store import EvidenceStore
31from hyper_parallel.tools.model_integration.optimizer_layout import (
32 flatten_tensor_layouts,
33 snapshot_training_layout,
34 validate_optimizer_layout,
35)
36from hyper_parallel.tools.model_integration.tensor_probes import (
37 InputIdentityRecorder,
38 ParameterProbeRecorder,
39)
40from hyper_parallel.tools.model_integration.schemas import IntegrationState
41from hyper_parallel.tools.model_integration.schemas import ModelIntegrationReport
42from hyper_parallel.tools.model_integration.integration_validator import (
43 collect_activation_checkpoint_wrappers,
44 inventory_final_model,
45 resolve_model_validation_spec,
46 validate_final_model,
47 validate_state_invariants,
48)
49from hyper_parallel.trainer.runtime.memory import peak_device_memory, reset_peak_device_memory
52_OUTPUT_DIR_ENV = "HYPER_PARALLEL_MODEL_INTEGRATION_OUTPUT_DIR"
53_SMALL_PARAMETER_NUMEL = 8192
56def _checkpoint_event_name(event: str, checkpoint_path: str) -> str:
57 """Build a filesystem-safe event name that preserves checkpoint identity."""
58 checkpoint_name = Path(checkpoint_path).name
59 safe_name = "".join(
60 character if character.isalnum() or character in "._-" else "_"
61 for character in checkpoint_name
62 )
63 return f"{event}_{safe_name}" if safe_name else event
66class ModelIntegrationSession(Protocol):
67 """Trainer-facing lifecycle shared by disabled and enabled sessions."""
69 runtime_enabled: bool
71 def attach_optimizer(self, optimizer: Any) -> None:
72 """Attach the final optimizer."""
74 def attach_data_pipeline(self, **pipeline: Any) -> None:
75 """Attach resolved data-pipeline objects."""
77 def begin_step(self, step: int) -> None:
78 """Begin one optimizer step."""
80 def record_batch(
81 self,
82 model_inputs: Mapping[str, Any],
83 loss_inputs: Mapping[str, Any],
84 ) -> None:
85 """Observe one micro-batch."""
87 def after_backward_before_clip(self) -> None:
88 """Observe gradients before clipping."""
90 def after_clip(self, grad_norm: Any) -> None:
91 """Observe gradients and their known norm after clipping."""
93 def after_optimizer(self) -> None:
94 """Observe parameters after the optimizer step."""
96 def end_step(self, metrics: Mapping[str, Any]) -> None:
97 """Finish one optimizer step."""
99 def record_fsdp_trace(self, trace: Mapping[str, Any]) -> None:
100 """Record optional FSDP runtime evidence."""
102 def capture_checkpoint_payload(
103 self,
104 event: str,
105 checkpoint_path: str,
106 payload: Mapping[str, Any],
107 ) -> None:
108 """Record optional checkpoint layout evidence."""
111def _analyze_shared_state_events(
112 events: tuple[dict[str, Any], ...],
113 producer_keys: tuple[str, ...],
114 consumer_keys: tuple[str, ...],
115) -> dict[str, Any]:
116 """Normalize trace identities and classify shared-state contract failures."""
117 state_indices: dict[Any, int] = {}
118 normalized_events = []
119 for event in events:
120 state_index = state_indices.setdefault(event["state_id"], len(state_indices))
121 normalized_events.append({**event, "state_id": state_index})
123 published = set()
124 duplicate_publications = []
125 missing_sources = []
126 for event in normalized_events:
127 identity = (event["state_id"], event["key"], event["source_layer"])
128 if event["action"] == "publish":
129 if identity in published:
130 duplicate_publications.append(identity)
131 published.add(identity)
132 elif identity not in published:
133 missing_sources.append(event)
135 produced_keys = {
136 event["key"] for event in normalized_events if event["action"] == "publish"
137 }
138 consumed_keys = {
139 event["key"] for event in normalized_events if event["action"] == "consume"
140 }
141 return {
142 "events": normalized_events,
143 "missing_producer_keys": sorted(set(producer_keys) - produced_keys),
144 "missing_consumer_keys": sorted(set(consumer_keys) - consumed_keys),
145 "missing_sources": missing_sources,
146 "duplicate_publications": duplicate_publications,
147 }
150class DisabledModelIntegrationSession:
151 """No-op session used to keep validation branches out of training code."""
153 runtime_enabled = False
155 def attach_optimizer(self, optimizer: Any) -> None:
156 """Ignore optimizer construction when validation is disabled."""
157 del optimizer
159 def attach_data_pipeline(self, **pipeline: Any) -> None:
160 """Ignore data-pipeline construction when validation is disabled."""
161 del pipeline
163 def begin_step(self, step: int) -> None:
164 """Ignore the optimizer-step boundary when validation is disabled."""
165 del step
167 def record_batch(
168 self,
169 model_inputs: Mapping[str, Any],
170 loss_inputs: Mapping[str, Any],
171 ) -> None:
172 """Ignore model inputs when validation is disabled."""
173 del model_inputs, loss_inputs
175 def after_backward_before_clip(self) -> None:
176 """Ignore the pre-clip boundary when validation is disabled."""
178 def after_clip(self, grad_norm: Any) -> None:
179 """Ignore the post-clip boundary when validation is disabled."""
180 del grad_norm
182 def after_optimizer(self) -> None:
183 """Ignore the optimizer boundary when validation is disabled."""
185 def end_step(self, metrics: Mapping[str, Any]) -> None:
186 """Ignore step metrics when validation is disabled."""
187 del metrics
189 def record_fsdp_trace(self, trace: Mapping[str, Any]) -> None:
190 """Ignore FSDP traces when validation is disabled."""
191 del trace
193 def capture_checkpoint_payload(
194 self,
195 event: str,
196 checkpoint_path: str,
197 payload: Mapping[str, Any],
198 ) -> None:
199 """Ignore checkpoint payloads when validation is disabled."""
200 del event, checkpoint_path, payload
203class TrainerModelIntegrationSession:
204 """Coordinate build checks and runtime evidence without changing execution."""
206 def __init__(
207 self,
208 *,
209 model: Any,
210 config: Any,
211 global_rank: int,
212 mesh_context: Any = None,
213 ) -> None:
214 """Validate the final model and persist integration evidence."""
215 self.model = model
216 self.config = config
217 self.global_rank = global_rank
218 self.mode = config.model_integration.mode
219 configured_output = os.environ.get(_OUTPUT_DIR_ENV)
220 if not configured_output:
221 configured_output = str(
222 Path(config.checkpoint.checkpoint_dir) / "model_integration"
223 )
224 self.evidence_store = EvidenceStore(configured_output)
225 self.parameter_probes: Optional[ParameterProbeRecorder] = None
226 self.optimizer: Any = None
227 self.validation_spec = resolve_model_validation_spec(self.model)
228 data_spec = self.validation_spec.data if self.validation_spec is not None else None
229 self.input_identity = InputIdentityRecorder(
230 mesh_context,
231 cp_replicated_forward_fields=(
232 () if data_spec is None else data_spec.cp_replicated_forward_fields
233 ),
234 )
235 self._current_step = 0
236 self._step_start_time = 0.0
237 self._grad_norm_post_clip: Optional[float] = None
238 self._shared_trace_token: Any = None
239 self._observed_forward_fields = False
240 self._modality_fields_observed = False
241 self._field_ownership: dict[str, str] = {}
242 self._validate_final_model()
244 @property
245 def runtime_enabled(self) -> bool:
246 """Whether step-level precision and execution probes are enabled."""
247 return self.mode == "runtime"
249 def _validate_final_model(self) -> None:
250 """Run final-tree validation and record the structured report."""
251 report, ownership, fsdp_units = validate_final_model(
252 self.model,
253 context={"trainer_config": self.config.to_dict()},
254 )
255 self.parameter_ownership = ownership
256 if self.global_rank == 0:
257 self.evidence_store.write_json("check/findings.json", report.to_dict())
258 self.evidence_store.write_json("check/fsdp_units.json", fsdp_units)
259 self.evidence_store.write_csv(
260 "check/parameter_ownership.csv",
261 (entry.to_dict() for entry in ownership),
262 fieldnames=(
263 "parameter_fqn",
264 "owner_fqn",
265 "gradient_domain",
266 "source_mesh",
267 "source_placements",
268 "global_shape",
269 "local_shape",
270 ),
271 )
272 self.evidence_store.write_json("trainer_config.resolved.json", self.config.to_dict())
273 self.evidence_store.capture_environment()
274 self.evidence_store.advance_if_before(
275 IntegrationState.DISCOVERED,
276 ("environment.json",),
277 )
278 if not report.errors:
279 self.evidence_store.advance_if_before(
280 IntegrationState.STRUCTURE_VALIDATED,
281 (
282 "check/findings.json",
283 "check/parameter_ownership.csv",
284 "check/fsdp_units.json",
285 ),
286 )
287 self.evidence_store.write_json(
288 "inventory/candidate.json",
289 inventory_final_model(self.model),
290 )
291 checkpoint_spec = (
292 self.validation_spec.checkpoint
293 if self.validation_spec is not None
294 else None
295 )
296 checkpoint_coverage, _ = build_checkpoint_coverage(
297 self.model,
298 checkpoint_spec,
299 )
300 self.evidence_store.write_json(
301 "checkpoint/coverage.json",
302 checkpoint_coverage,
303 )
304 self.evidence_store.write_json(
305 "check/activation_checkpoint.json",
306 collect_activation_checkpoint_wrappers(self.model),
307 )
308 report.raise_for_errors()
310 def attach_optimizer(self, optimizer: Any) -> None:
311 """Resolve the semantic parameter watchlist after optimizer construction."""
312 self.optimizer = optimizer
313 findings = validate_optimizer_layout(self.model, optimizer)
314 if self.global_rank == 0:
315 self.evidence_store.write_json(
316 "check/optimizer_layout_initial.json",
317 {
318 "status": "FAIL" if findings else "PASS",
319 "findings": [finding.to_dict() for finding in findings],
320 "layout": snapshot_training_layout(self.model, optimizer),
321 },
322 )
323 if findings:
324 ModelIntegrationReport(findings=findings).raise_for_errors()
325 if not self.runtime_enabled:
326 return
327 self.parameter_probes = ParameterProbeRecorder(
328 self.model,
329 optimizer,
330 self.evidence_store,
331 rank=self.global_rank,
332 validation_spec=self.validation_spec,
333 include_small_parameters_numel_le=_SMALL_PARAMETER_NUMEL,
334 representative_parameter_names=self._representative_parameter_names(),
335 )
336 if self.global_rank == 0:
337 self.evidence_store.write_json(
338 "check/parameter_watchlist.json",
339 {"parameters": list(self.parameter_probes.watched_parameter_names)},
340 )
342 def _representative_parameter_names(self) -> tuple[str, ...]:
343 """Select at least one parameter from every owner/domain/placement group."""
344 representatives = {}
345 for entry in self.parameter_ownership:
346 key = (
347 entry.owner_fqn,
348 entry.gradient_domain,
349 entry.source_mesh,
350 entry.source_placements,
351 )
352 representatives.setdefault(key, entry.parameter_fqn)
353 return tuple(sorted(representatives.values()))
355 def attach_data_pipeline(
356 self,
357 *,
358 runtime_adapter: Any,
359 get_batch: Any,
360 model_assets: tuple[Any, ...],
361 ) -> None:
362 """Validate and record the resolved data/forward-input ownership."""
363 ownership, findings = validate_data_contract(
364 self.validation_spec,
365 runtime_adapter,
366 labels_are_shifted=getattr(get_batch, "labels_are_shifted", None),
367 source_type=getattr(get_batch, "source_type", None),
368 assets=model_assets,
369 )
370 self._field_ownership = dict(ownership["fields"])
371 if self.global_rank == 0:
372 self.evidence_store.write_json("data/field_ownership.json", ownership)
373 self.evidence_store.write_json(
374 "data/findings.json",
375 {
376 "status": "FAIL" if findings else "PASS",
377 "findings": [finding.to_dict() for finding in findings],
378 },
379 )
380 if findings:
381 ModelIntegrationReport(findings=findings).raise_for_errors()
383 def begin_step(self, step: int) -> None:
384 """Begin input and parameter evidence for one optimizer step."""
385 if not self.runtime_enabled:
386 return
387 self._current_step = step
388 self._step_start_time = time.perf_counter()
389 self._grad_norm_post_clip = None
390 reset_peak_device_memory()
391 shared_state_spec = (
392 self.validation_spec.shared_state
393 if self.validation_spec is not None
394 else None
395 )
396 if shared_state_spec is not None and shared_state_spec.begin_trace is not None:
397 self._shared_trace_token = shared_state_spec.begin_trace()
398 self.input_identity.begin_step()
399 if self.parameter_probes is not None:
400 self.parameter_probes.begin_step(step)
402 def record_batch(
403 self,
404 model_inputs: Mapping[str, Any],
405 loss_inputs: Mapping[str, Any],
406 ) -> None:
407 """Record exact tensor fields reaching the model/loss boundary.
409 Args:
410 model_inputs: Keyword inputs passed to the model forward call.
411 loss_inputs: Labels and masks consumed by the loss path.
412 """
413 if not self.runtime_enabled:
414 return
415 if not self._observed_forward_fields:
416 findings = validate_observed_forward_fields(
417 self.validation_spec,
418 model_inputs,
419 )
420 data_spec = self.validation_spec.data if self.validation_spec is not None else None
421 if data_spec is not None:
422 self._modality_fields_observed = bool(
423 set(data_spec.modality_fields).intersection(model_inputs)
424 )
425 if self.global_rank == 0:
426 actual_ownership = {
427 field_name: self._field_ownership.get(
428 field_name,
429 "source_or_get_batch",
430 )
431 for field_name in sorted(model_inputs)
432 }
433 self.evidence_store.write_json(
434 "data/observed_forward_fields.json",
435 {
436 "status": "FAIL" if findings else "PASS",
437 "model_fields": sorted(model_inputs),
438 "loss_fields": sorted(loss_inputs),
439 "field_ownership": actual_ownership,
440 "findings": [finding.to_dict() for finding in findings],
441 },
442 )
443 self._observed_forward_fields = True
444 if findings:
445 ModelIntegrationReport(findings=findings).raise_for_errors()
446 self.input_identity.record(model_inputs, loss_inputs)
448 def after_backward_before_clip(self) -> None:
449 """Capture un-clipped main gradients."""
450 if not self.runtime_enabled:
451 return
452 runtime_findings = validate_state_invariants(
453 self.model,
454 self.validation_spec,
455 {"trainer": self, "step": self._current_step},
456 phases=("runtime",),
457 )
458 if runtime_findings:
459 ModelIntegrationReport(findings=runtime_findings).raise_for_errors()
460 if self.parameter_probes is not None:
461 self.parameter_probes.capture("main_grad_before_clip")
462 if self.optimizer is not None:
463 findings = validate_optimizer_layout(self.model, self.optimizer)
464 if self.global_rank == 0:
465 self.evidence_store.write_json(
466 f"cases/runtime/optimizer_layout_step{self._current_step}.json",
467 {
468 "status": "FAIL" if findings else "PASS",
469 "findings": [finding.to_dict() for finding in findings],
470 "layout": snapshot_training_layout(self.model, self.optimizer),
471 },
472 )
473 if findings:
474 ModelIntegrationReport(findings=findings).raise_for_errors()
475 if self._modality_fields_observed:
476 modality_gradients = modality_gradient_parameters(
477 self.validation_spec,
478 self.model,
479 )
480 if self.global_rank == 0:
481 self.evidence_store.write_json(
482 f"cases/runtime/modality_gradients_step{self._current_step}.json",
483 modality_gradients,
484 )
485 if modality_gradients and not all(modality_gradients.values()):
486 raise RuntimeError(
487 "multimodal inputs reached the model but one or more "
488 "declared visual parameters have missing or non-finite gradients"
489 )
491 def record_fsdp_trace(self, trace: Mapping[str, Any]) -> None:
492 """Persist rank-local runtime order before the optimizer-boundary gate."""
493 if not self.runtime_enabled:
494 return
495 self.evidence_store.write_json(
496 f"check/fsdp_runtime_rank{self.global_rank}.json",
497 dict(trace),
498 )
500 def after_clip(self, grad_norm: Any) -> None:
501 """Capture clipped main gradients."""
502 if not self.runtime_enabled:
503 return
504 if grad_norm is None:
505 raise ValueError("runtime model integration requires the post-clip gradient norm")
506 if self.parameter_probes is not None:
507 self.parameter_probes.capture("main_grad_after_clip")
508 self._grad_norm_post_clip = float(grad_norm)
510 def after_optimizer(self) -> None:
511 """Capture main/model updates and optimizer moment layouts."""
512 if not self.runtime_enabled:
513 return
514 if self.parameter_probes is None:
515 return
516 self.parameter_probes.capture("main_after_optimizer")
517 self.parameter_probes.capture("model_after_copy_back")
518 self.parameter_probes.capture("optimizer_state")
520 def end_step(self, metrics: Mapping[str, Any]) -> None:
521 """Persist structured scalar and global input identity evidence."""
522 if not self.runtime_enabled:
523 return
524 self._finish_shared_state_trace()
525 input_identity = self.input_identity.finish_step()
526 elapsed = time.perf_counter() - self._step_start_time
527 tokens = metrics.get("tokens")
528 samples = metrics.get("samples")
529 memory = peak_device_memory()
530 row = {
531 "step": self._current_step,
532 "loss": metrics.get("loss"),
533 "grad_norm_pre_clip": metrics.get("grad_norm"),
534 "grad_norm_post_clip": metrics.get(
535 "grad_norm_post_clip",
536 self._grad_norm_post_clip,
537 ),
538 "lr": metrics.get("lr"),
539 "global_input_sha256": input_identity["global_sha256"],
540 "input_identity": input_identity,
541 "step_time_seconds": elapsed,
542 "tokens_per_second": None if tokens is None else float(tokens) / elapsed,
543 "samples_per_second": None if samples is None else float(samples) / elapsed,
544 **memory,
545 }
546 performance_row = {
547 "step": self._current_step,
548 "rank": self.global_rank,
549 "step_time_seconds": elapsed,
550 "tokens_per_second": row["tokens_per_second"],
551 "samples_per_second": row["samples_per_second"],
552 **memory,
553 }
554 self.evidence_store.append_jsonl(
555 f"cases/runtime/performance_rank{self.global_rank}.jsonl",
556 performance_row,
557 )
558 if self.global_rank == 0:
559 self.evidence_store.append_jsonl("cases/runtime/metrics.jsonl", row)
560 self.evidence_store.append_jsonl(
561 "cases/runtime/performance.jsonl",
562 performance_row,
563 )
565 def _finish_shared_state_trace(self) -> None:
566 """Validate model-declared shared-state publication and consumption."""
567 if self._shared_trace_token is None or self.validation_spec is None:
568 return
569 shared_state_spec = self.validation_spec.shared_state
570 if shared_state_spec is None or shared_state_spec.finish_trace is None:
571 return
572 events = shared_state_spec.finish_trace(self._shared_trace_token)
573 self._shared_trace_token = None
574 analysis = _analyze_shared_state_events(
575 events,
576 shared_state_spec.producer_keys,
577 shared_state_spec.consumer_keys,
578 )
579 if self.global_rank == 0:
580 self.evidence_store.write_json(
581 f"cases/runtime/shared_state_step{self._current_step}.json",
582 analysis,
583 )
584 if (
585 analysis["missing_sources"]
586 or analysis["missing_producer_keys"]
587 or analysis["missing_consumer_keys"]
588 ):
589 raise RuntimeError(
590 "shared-state producer/consumer contract failed: "
591 f"missing_producers={analysis['missing_producer_keys']}, "
592 f"missing_consumers={analysis['missing_consumer_keys']}, "
593 f"missing_sources={analysis['missing_sources']}"
594 )
595 if analysis["duplicate_publications"]:
596 raise RuntimeError(
597 "shared-state producer replayed inside one forward: "
598 f"{analysis['duplicate_publications']}"
599 )
601 def capture_checkpoint_payload(
602 self,
603 event: str,
604 checkpoint_path: str,
605 payload: Mapping[str, Any],
606 ) -> None:
607 """Record versioned DCP layouts without rebuilding checkpoint chunks."""
608 if not self.runtime_enabled:
609 return
610 event_name = _checkpoint_event_name(event, checkpoint_path)
611 self.evidence_store.write_json(
612 f"checkpoint/{event_name}_rank{self.global_rank}.json",
613 {
614 "checkpoint_path": checkpoint_path,
615 "tensor_layouts": flatten_tensor_layouts(payload),
616 },
617 )
620def build_model_integration_session(
621 model: Any,
622 config: Any,
623 global_rank: int,
624 mesh_context: Any = None,
625) -> TrainerModelIntegrationSession | DisabledModelIntegrationSession:
626 """Build the configured validation session or a branch-free no-op session."""
627 if getattr(getattr(config, "model_integration", None), "mode", "off") == "off":
628 return DisabledModelIntegrationSession()
629 return TrainerModelIntegrationSession(
630 model=model,
631 config=config,
632 global_rank=global_rank,
633 mesh_context=mesh_context,
634 )
637__all__ = [
638 "DisabledModelIntegrationSession",
639 "ModelIntegrationSession",
640 "TrainerModelIntegrationSession",
641 "build_model_integration_session",
642]