Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / tools / model_integration / optimizer_layout.py: 28%
88 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"""Optimizer and checkpoint-layout observations for Trainer validation."""
16# pylint: disable=forbidden-backend-import
18from __future__ import annotations
20import math
21from collections.abc import Mapping
22from typing import Any
24import torch
26from hyper_parallel.tools.model_integration.schemas import (
27 FindingSeverity,
28 ModelIntegrationFinding,
29 SCHEMA_VERSION,
30)
33def tensor_layout(value: Any) -> dict[str, Any]:
34 """Return logical/local shape and distributed layout without custom chunks."""
35 to_local = getattr(value, "to_local", None)
36 local = to_local() if callable(to_local) else value
37 layout = getattr(value, "layout", None)
38 mesh = getattr(value, "device_mesh", None)
39 placements = getattr(value, "placements", None)
40 if layout is not None:
41 mesh = getattr(layout, "mesh", mesh)
42 placements = getattr(layout, "placements", placements)
43 mesh_tensor = getattr(mesh, "mesh", None)
44 return {
45 "global_shape": [int(size) for size in getattr(value, "shape", ())],
46 "local_shape": [int(size) for size in getattr(local, "shape", ())],
47 "dtype": str(getattr(value, "dtype", None)),
48 "device": str(getattr(local, "device", None)),
49 "mesh_dim_names": list(getattr(mesh, "mesh_dim_names", ()) or ()),
50 "mesh_shape": [int(size) for size in getattr(mesh_tensor, "shape", ())],
51 "placements": [repr(placement) for placement in (placements or ())],
52 "is_dtensor": callable(to_local),
53 }
56def _same_logical_layout(first: Any, second: Any) -> bool:
57 first_layout = tensor_layout(first)
58 second_layout = tensor_layout(second)
59 return all(
60 first_layout[name] == second_layout[name]
61 for name in ("global_shape", "mesh_dim_names", "mesh_shape", "placements")
62 )
65def optimizer_leaves(optimizer: Any) -> tuple[Any, ...]:
66 """Return optimizer leaves through list, chain, and wrapper containers."""
67 chained = getattr(optimizer, "chained_optimizers", None)
68 if chained is not None:
69 return tuple(leaf for item in chained for leaf in optimizer_leaves(item))
70 inner = getattr(optimizer, "optimizer", None)
71 if inner is not None and inner is not optimizer:
72 return optimizer_leaves(inner)
73 if isinstance(optimizer, (list, tuple)):
74 return tuple(leaf for item in optimizer for leaf in optimizer_leaves(item))
75 return (optimizer,)
78def optimizer_state_for_parameter(optimizer: Any, parameter: Any) -> Mapping[str, Any]:
79 """Find one parameter's state through supported optimizer wrappers."""
80 for leaf in optimizer_leaves(optimizer):
81 state = getattr(leaf, "state", None)
82 if isinstance(state, Mapping) and parameter in state:
83 value = state[parameter]
84 return value if isinstance(value, Mapping) else {}
85 return {}
88def _layout_mismatch_finding(
89 parameter_fqn: str,
90 owner_name: str,
91 expected: Any,
92 actual: Any,
93) -> ModelIntegrationFinding:
94 return ModelIntegrationFinding(
95 phase="D6",
96 owner_fqn=parameter_fqn,
97 severity=FindingSeverity.ERROR,
98 message=f"{owner_name} does not preserve its parameter's global layout",
99 facts={
100 "parameter": tensor_layout(expected),
101 owner_name: tensor_layout(actual),
102 },
103 why_unsafe=(
104 "optimizer state and gradients must describe the same logical tensor; "
105 "a local-only or differently placed value cannot be resharded safely"
106 ),
107 remediation=(
108 "preserve the parameter DTensor global shape, mesh, and placements when "
109 "creating main gradients, main parameters, and optimizer moments"
110 ),
111 related_config=("optimizer.fp32_main_params", "checkpoint.restore_optimizer"),
112 )
115def validate_optimizer_layout(model: Any, optimizer: Any) -> list[ModelIntegrationFinding]:
116 """Validate param/main-param/main-grad/moment layouts from live objects."""
117 findings = []
118 for parameter_fqn, model_parameter in model.named_parameters(remove_duplicate=False):
119 main_parameter = getattr(model_parameter, "main_param", model_parameter)
120 if not _same_logical_layout(model_parameter, main_parameter):
121 findings.append(
122 _layout_mismatch_finding(
123 parameter_fqn,
124 "main_parameter",
125 model_parameter,
126 main_parameter,
127 )
128 )
129 continue
130 gradient = getattr(model_parameter, "main_grad", None)
131 if gradient is None:
132 gradient = getattr(main_parameter, "grad", None)
133 if gradient is not None and not _same_logical_layout(main_parameter, gradient):
134 findings.append(
135 _layout_mismatch_finding(
136 parameter_fqn,
137 "main_gradient",
138 main_parameter,
139 gradient,
140 )
141 )
142 for state_name, state_value in optimizer_state_for_parameter(
143 optimizer, main_parameter
144 ).items():
145 if torch.is_tensor(state_value) and state_value.ndim > 0:
146 if not _same_logical_layout(main_parameter, state_value):
147 findings.append(
148 _layout_mismatch_finding(
149 parameter_fqn,
150 f"optimizer_state.{state_name}",
151 main_parameter,
152 state_value,
153 )
154 )
155 layout = tensor_layout(main_parameter)
156 if (
157 math.prod(layout["global_shape"])
158 and not math.prod(layout["local_shape"])
159 and not layout["is_dtensor"]
160 ):
161 findings.append(
162 ModelIntegrationFinding(
163 phase="D8",
164 owner_fqn=parameter_fqn,
165 severity=FindingSeverity.ERROR,
166 message="empty optimizer shard has no global DTensor layout",
167 facts={"main_parameter": layout},
168 why_unsafe=(
169 "DCP cannot place an empty local tensor into its logical global tensor "
170 "without mesh and placement metadata"
171 ),
172 remediation="keep uneven FSDP shards as DTensors through optimizer state creation",
173 related_config=("parallel.fsdp", "optimizer.fp32_main_params"),
174 )
175 )
176 return findings
179def snapshot_training_layout(model: Any, optimizer: Any) -> dict[str, Any]:
180 """Build model/main-gradient/optimizer-state layout evidence."""
181 records = {}
182 for parameter_fqn, model_parameter in model.named_parameters(remove_duplicate=False):
183 main_parameter = getattr(model_parameter, "main_param", model_parameter)
184 record = {
185 "model_parameter": tensor_layout(model_parameter),
186 "main_parameter": tensor_layout(main_parameter),
187 }
188 gradient = getattr(model_parameter, "main_grad", None)
189 if gradient is None:
190 gradient = getattr(main_parameter, "grad", None)
191 if gradient is not None:
192 record["main_gradient"] = tensor_layout(gradient)
193 state = optimizer_state_for_parameter(optimizer, main_parameter)
194 record["optimizer_state"] = {
195 name: tensor_layout(value)
196 for name, value in sorted(state.items())
197 if torch.is_tensor(value)
198 }
199 records[parameter_fqn] = record
200 return {"schema_version": SCHEMA_VERSION, "parameters": records}
203def flatten_tensor_layouts(value: Any, prefix: str = "") -> dict[str, dict[str, Any]]:
204 """Flatten tensor leaves from a checkpoint payload into layout records."""
205 records = {}
206 if torch.is_tensor(value):
207 records[prefix or "<root>"] = tensor_layout(value)
208 elif isinstance(value, Mapping):
209 for name, child in value.items():
210 child_prefix = f"{prefix}.{name}" if prefix else str(name)
211 records.update(flatten_tensor_layouts(child, child_prefix))
212 elif isinstance(value, (list, tuple)):
213 for index, child in enumerate(value):
214 child_prefix = f"{prefix}.{index}" if prefix else str(index)
215 records.update(flatten_tensor_layouts(child, child_prefix))
216 return records
219__all__ = [
220 "flatten_tensor_layouts",
221 "optimizer_leaves",
222 "optimizer_state_for_parameter",
223 "snapshot_training_layout",
224 "tensor_layout",
225 "validate_optimizer_layout",
226]