Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / optim / mixed_precision_optimizer.py: 93%
167 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"""Megatron-aligned fp32 main-parameter optimizer composition."""
17import logging
18from typing import Any
20import torch # pylint: disable=forbidden-backend-import
21from torch import nn # pylint: disable=forbidden-backend-import
23from hyper_parallel.core.optimizer.dtensor_compat import (
24 device_meshes_are_compatible,
25 is_dtensor,
26 to_local_if_dtensor,
27)
28from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer
31MIXED_PRECISION_OPTIMIZER_STATE_KEY = "_mixed_precision_optimizer"
32FP32_MAIN_PARAM_STATE_KEY = "fp32_from_fp16_params"
34logger = logging.getLogger(__name__)
37def _copy_tensor(destination: torch.Tensor, source: torch.Tensor) -> None:
38 """Copy one local parameter shard without invoking DTensor dispatch."""
39 destination_local = to_local_if_dtensor(destination)
40 source_local = to_local_if_dtensor(source)
41 source_local = source_local.to(
42 device=destination_local.device,
43 dtype=destination_local.dtype,
44 )
45 destination_local.copy_(source_local)
48def _gradient_for_param(
49 optimizer_param: nn.Parameter,
50 gradient: torch.Tensor,
51) -> torch.Tensor:
52 """Validate gradient layout and cast dtype on the gradient's existing device.
54 DTensor optimizer parameters keep a DTensor gradient so lazily-created
55 optimizer state inherits its complete global layout. FSDP ``main_grad`` is
56 already detached from autograd, so the DTensor object itself is preserved.
57 Plain parameters retain the detached local-tensor path used by non-FSDP and
58 third-party integrations.
59 """
60 optimizer_is_dtensor = is_dtensor(optimizer_param)
61 gradient_is_dtensor = is_dtensor(gradient)
62 if optimizer_is_dtensor != gradient_is_dtensor:
63 raise TypeError(
64 "Optimizer parameter and gradient must both be DTensors or both be local tensors"
65 )
67 optimizer_local = to_local_if_dtensor(optimizer_param)
68 gradient_local = to_local_if_dtensor(gradient)
69 if gradient_local.device != optimizer_local.device:
70 raise ValueError(
71 "Optimizer gradient must already be on its parameter's device: "
72 f"got {gradient_local.device} and {optimizer_local.device}"
73 )
74 if tuple(gradient_local.shape) != tuple(optimizer_local.shape):
75 raise ValueError(
76 "Optimizer gradient local shape must match its parameter shard: "
77 f"got {tuple(gradient_local.shape)} and {tuple(optimizer_local.shape)}"
78 )
79 if optimizer_is_dtensor:
80 if tuple(gradient.shape) != tuple(optimizer_param.shape):
81 raise ValueError(
82 "Optimizer gradient global shape must match its parameter: "
83 f"got {tuple(gradient.shape)} and {tuple(optimizer_param.shape)}"
84 )
85 if not device_meshes_are_compatible(
86 gradient.device_mesh,
87 optimizer_param.device_mesh,
88 ):
89 raise ValueError("Optimizer gradient and parameter must use the same device mesh")
90 if tuple(gradient.placements) != tuple(optimizer_param.placements):
91 raise ValueError("Optimizer gradient and parameter must use the same placements")
92 # FSDP main_grad is already an autograd-detached accumulation buffer.
93 # Keep that DTensor object intact: detach() rebuilds an uneven empty
94 # shard without its logical tensor metadata, which makes a subsequent
95 # zeros_like() infer the optimizer state's global shape as zero.
96 prepared_gradient = gradient
97 else:
98 prepared_gradient = gradient_local.detach()
99 if gradient_local.dtype != optimizer_local.dtype:
100 prepared_gradient = prepared_gradient.to(dtype=optimizer_local.dtype)
101 return prepared_gradient
104class MixedPrecisionOptimizer:
105 """Base composition for an optimizer that owns precision-conversion state."""
107 def __init__(self, optimizer: ChainedOptimizer, model: nn.Module) -> None:
108 """Bind the inner optimizer and expose its scheduler-facing state.
110 Args:
111 optimizer: Chained leaf-optimizer composition.
112 model: Module containing the computation parameters.
113 """
114 self.optimizer = optimizer
115 self.model = model
116 self.optimizers_dict: dict[str, Any] = optimizer.optimizers_dict
117 self.chained_optimizers: list[Any] = optimizer.chained_optimizers
118 self.param_groups: list[dict[str, Any]] = optimizer.param_groups
120 def prepare_grads(self) -> None:
121 """Prepare model gradients for the inner optimizer."""
122 raise NotImplementedError
124 def step_with_ready_grads(self, closure: Any = None) -> Any:
125 """Update parameters after :meth:`prepare_grads` completed.
127 Args:
128 closure: Optional callable that reevaluates the model and returns loss.
129 """
130 raise NotImplementedError
132 def reload_model_params(self) -> None:
133 """Reload optimizer-owned parameters from the registered model."""
134 raise NotImplementedError
137class Float16OptimizerWithFloat16Params(MixedPrecisionOptimizer):
138 """Update fp32 main params and copy them back to fp16/bf16 model params."""
140 def __init__(self, optimizer: ChainedOptimizer, model: nn.Module) -> None:
141 """Replace leaf optimizer params with fp32 main parameters.
143 Args:
144 optimizer: Chained AdamW/Muon optimizer over model parameters.
145 model: Module containing the computation parameters.
146 """
147 super().__init__(optimizer, model)
148 self.float16_groups: list[list[nn.Parameter]] = []
149 self.fp32_from_float16_groups: list[list[nn.Parameter]] = []
150 self.fp32_from_fp32_groups: list[list[nn.Parameter]] = []
151 self.model_param_by_optimizer_param: dict[nn.Parameter, nn.Parameter] = {}
152 self._build_model_and_main_param_groups()
153 self.optimizer_param_by_model_param = {
154 model_param: optimizer_param
155 for optimizer_param, model_param in self.model_param_by_optimizer_param.items()
156 }
157 self.optimizer.reset_optimizer_parameters(
158 self.model_param_by_optimizer_param
159 )
160 self.param_groups = self.optimizer.param_groups
162 parameter_names = {
163 parameter: name for name, parameter in model.named_parameters()
164 }
165 self._main_param_by_fqn = {
166 parameter_names[model_param]: optimizer_param
167 for optimizer_param, model_param in self.model_param_by_optimizer_param.items()
168 }
170 def _build_model_and_main_param_groups(self) -> None:
171 """Build Megatron-compatible model and main-parameter group views."""
172 parameter_names = {
173 parameter: name for name, parameter in self.model.named_parameters()
174 }
175 for leaf_optimizer in self.chained_optimizers:
176 for optimizer_group in leaf_optimizer.param_groups:
177 float16_group = []
178 fp32_from_float16_group = []
179 fp32_from_fp32_group = []
180 for param_index, model_param in enumerate(optimizer_group["params"]):
181 model_param.main_grad = None
182 if model_param.dtype == torch.float32:
183 model_param.main_param = model_param
184 fp32_from_fp32_group.append(model_param)
185 continue
186 if model_param.dtype not in (torch.float16, torch.bfloat16):
187 raise TypeError(
188 "Float16OptimizerWithFloat16Params supports trainable "
189 "float16, bfloat16, or float32 parameters; "
190 f"{parameter_names[model_param]} has dtype {model_param.dtype}"
191 )
193 main_param = nn.Parameter(
194 model_param.detach().to(dtype=torch.float32),
195 requires_grad=model_param.requires_grad,
196 )
197 main_param.model_name = model_param.model_name
198 model_param.main_param = main_param
199 optimizer_group["params"][param_index] = main_param
200 if model_param in leaf_optimizer.state:
201 leaf_optimizer.state[main_param] = leaf_optimizer.state.pop(model_param)
203 float16_group.append(model_param)
204 fp32_from_float16_group.append(main_param)
205 self.model_param_by_optimizer_param[main_param] = model_param
207 self.float16_groups.append(float16_group)
208 self.fp32_from_float16_groups.append(fp32_from_float16_group)
209 self.fp32_from_fp32_groups.append(fp32_from_fp32_group)
211 def _copy_model_grads_to_main_grads(self) -> None:
212 """Transfer model gradients to fp32 optimizer params, casting dtype only.
214 The producer must supply the optimizer parameter's device and layout.
215 Rebind existing ``main_grad`` buffers to the prepared gradient and clear
216 model ``grad``. The local source reference is overwritten before the
217 next conversion or released on return, so reference counting reclaims
218 converted source buffers before the optimizer step.
219 External owners of shared gradient storage remain responsible for it.
220 """
221 for model_group, main_group in zip(
222 self.float16_groups,
223 self.fp32_from_float16_groups,
224 ):
225 for model_param, main_param in zip(model_group, main_group):
226 model_gradient = model_param.main_grad
227 if model_gradient is None:
228 model_gradient = model_param.grad
229 main_param.grad = (
230 _gradient_for_param(main_param, model_gradient)
231 if model_gradient is not None else None
232 )
233 if model_param.main_grad is not None:
234 model_param.main_grad = main_param.grad
235 model_param.grad = None
237 for fp32_group in self.fp32_from_fp32_groups:
238 for model_param in fp32_group:
239 model_gradient = model_param.main_grad
240 if model_gradient is not None:
241 model_param.grad = _gradient_for_param(model_param, model_gradient)
242 model_param.main_grad = model_param.grad
244 @torch.no_grad()
245 def _copy_main_params_to_model_params(self) -> None:
246 """Copy updated fp32 main params to fp16/bf16 model params."""
247 for model_group, main_group in zip(
248 self.float16_groups,
249 self.fp32_from_float16_groups,
250 ):
251 for model_param, main_param in zip(model_group, main_group):
252 _copy_tensor(model_param, main_param)
254 @torch.no_grad()
255 def _copy_model_params_to_main_params(self) -> None:
256 """Refresh fp32 main params from fp16/bf16 model params."""
257 for model_group, main_group in zip(
258 self.float16_groups,
259 self.fp32_from_float16_groups,
260 ):
261 for model_param, main_param in zip(model_group, main_group):
262 _copy_tensor(main_param, model_param)
264 def prepare_grads(self) -> None:
265 """Make fp32 main gradients ready for AdamW/Muon."""
266 self._copy_model_grads_to_main_grads()
268 def step_with_ready_grads(self, closure: Any = None) -> Any:
269 """Update main params, then copy them back to model params.
271 Args:
272 closure: Optional callable that reevaluates the model and returns loss.
274 Returns:
275 The inner optimizer step result.
276 """
277 loss = self.optimizer.step(closure=closure)
278 self._copy_main_params_to_model_params()
279 return loss
281 def step(self, closure: Any = None) -> Any:
282 """Prepare gradients and execute one mixed-precision optimizer step.
284 Args:
285 closure: Optional callable that reevaluates the model and returns loss.
287 Returns:
288 The inner optimizer step result.
289 """
290 self.prepare_grads()
291 return self.step_with_ready_grads(closure=closure)
293 def reload_model_params(self) -> None:
294 """Initialize main params from the current model parameter values."""
295 self._copy_model_params_to_main_params()
297 def zero_grad(self, set_to_none: bool = True) -> None:
298 """Clear inner grads plus runtime-only model grad/main_grad state.
300 Args:
301 set_to_none: Whether inner gradients are reset to ``None``.
302 """
303 self.optimizer.zero_grad(set_to_none=set_to_none)
304 for model_param in self.optimizer_param_by_model_param:
305 model_param.grad = None
306 model_param.main_grad = None
307 if hasattr(model_param, "grad_added_to_main_grad"):
308 model_param.grad_added_to_main_grad = False
309 for fp32_group in self.fp32_from_fp32_groups:
310 for model_param in fp32_group:
311 model_param.grad = None
312 model_param.main_grad = None
313 if hasattr(model_param, "grad_added_to_main_grad"):
314 model_param.grad_added_to_main_grad = False
316 def state_dict(self) -> dict[str, Any]:
317 """Append fp32 main weights to the existing optimizer FQN state."""
318 state_dict = self.optimizer.state_dict()
319 state_dict[MIXED_PRECISION_OPTIMIZER_STATE_KEY] = {
320 FP32_MAIN_PARAM_STATE_KEY: dict(self._main_param_by_fqn),
321 }
322 return state_dict
324 def load_state_dict(self, state_dict: dict[str, Any]) -> None:
325 """Restore optimizer moments and explicit fp32 main weights.
327 Args:
328 state_dict: Optimizer state containing leaf and main-param entries.
329 """
330 optimizer_state = dict(state_dict)
331 if MIXED_PRECISION_OPTIMIZER_STATE_KEY not in optimizer_state:
332 self.optimizer.load_state_dict(optimizer_state)
333 self.param_groups = self.optimizer.param_groups
334 self.reload_model_params()
335 logger.warning(
336 "Checkpoint has no %s state; initialized fp32 main parameters "
337 "from the loaded model parameters. Values not representable in "
338 "the model precision cannot be recovered.",
339 MIXED_PRECISION_OPTIMIZER_STATE_KEY,
340 )
341 return
343 mixed_precision_state = optimizer_state.pop(
344 MIXED_PRECISION_OPTIMIZER_STATE_KEY
345 )
346 if not isinstance(mixed_precision_state, dict):
347 raise RuntimeError(
348 "Mixed-precision optimizer checkpoint state must be a mapping"
349 )
350 checkpoint_main_params = mixed_precision_state.get(
351 FP32_MAIN_PARAM_STATE_KEY
352 )
353 if not isinstance(checkpoint_main_params, dict):
354 raise RuntimeError(
355 "Mixed-precision optimizer checkpoint is missing "
356 f"{FP32_MAIN_PARAM_STATE_KEY}"
357 )
358 missing_fqns = sorted(
359 set(self._main_param_by_fqn) - set(checkpoint_main_params)
360 )
361 if missing_fqns:
362 raise RuntimeError(
363 "Mixed-precision optimizer checkpoint is missing fp32 main "
364 f"parameters: {', '.join(missing_fqns)}"
365 )
367 self.optimizer.load_state_dict(optimizer_state)
368 self.param_groups = self.optimizer.param_groups
369 with torch.no_grad():
370 for parameter_fqn, main_param in self._main_param_by_fqn.items():
371 _copy_tensor(main_param, checkpoint_main_params[parameter_fqn])
374__all__ = [
375 "FP32_MAIN_PARAM_STATE_KEY",
376 "Float16OptimizerWithFloat16Params",
377 "MIXED_PRECISION_OPTIMIZER_STATE_KEY",
378 "MixedPrecisionOptimizer",
379]