Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / sinkhorn.py: 92%

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 

16"""Momentum updates with distributed Sinkhorn balancing (DeepSeek V4.1, Algorithm 1).""" 

17 

18import math 

19from typing import Any, Callable, Dict, Optional, Tuple 

20 

21import torch # pylint: disable=forbidden-backend-import 

22import torch.distributed as dist # pylint: disable=forbidden-backend-import 

23 

24from hyper_parallel.core.optimizer.dtensor_compat import is_dtensor, to_local_if_dtensor 

25 

26 

27def _sum_across(value: torch.Tensor, groups: Tuple[Any, ...]) -> torch.Tensor: 

28 """Reduce statistics only over the mesh axes partitioning the reduced dimension.""" 

29 for group in groups: 

30 if value.numel(): 

31 dist.all_reduce(value, group=group) 

32 return value 

33 

34 

35def _distributed_norm( 

36 update: torch.Tensor, dimension: int, groups: Tuple[Any, ...], 

37) -> torch.Tensor: 

38 """Reduce squared local norms without materializing a matrix-sized square.""" 

39 norm = torch.linalg.vector_norm(update, dim=dimension, keepdim=True) 

40 if groups: 

41 _sum_across(norm.square_(), groups).sqrt_() 

42 return norm 

43 

44 

45def _validate_balancing(steps: int, eps: float, tau: float) -> None: 

46 """Validate algorithm constants before any collective is launched.""" 

47 if not isinstance(steps, int) or isinstance(steps, bool) or steps <= 0 or steps % 2 != 1: 

48 raise ValueError("steps must be a positive odd integer") 

49 if not math.isfinite(eps) or eps <= 0: 

50 raise ValueError("eps must be finite and positive") 

51 if not math.isfinite(tau) or tau < 0: 

52 raise ValueError("tau must be finite and non-negative") 

53 

54 

55class Sinkhorn(torch.optim.Optimizer): 

56 """Nesterov momentum with row/column balancing and no weight decay. 

57 

58 Gradients must be dense, fully reduced, and consistently present or absent 

59 across all shards of a matrix. Statistics use float32; momentum is stored as 

60 a parameter-shaped float32 tensor, preserving DTensor placement metadata. 

61 """ 

62 

63 def __init__( 

64 self, 

65 params: Any, 

66 lr: float = 1e-3, 

67 momentum: float = 0.95, 

68 nesterov: bool = True, 

69 steps: int = 11, 

70 eps: float = 1e-20, 

71 tau: float = 1e-3, 

72 correction: float = 0.18, 

73 ) -> None: 

74 """Initialize independent Sinkhorn parameter groups and momentum state. 

75 

76 Args: 

77 params: Matrices or parameter groups with per-group overrides. 

78 lr: Base learning rate, before correction. 

79 momentum: Exponential moving-average coefficient. 

80 nesterov: Use Algorithm 1's Nesterov update; otherwise balance the EMA. 

81 steps: Positive odd count of normalization steps. 

82 eps: Positive L2-norm stabilizer. 

83 tau: Near-zero row masking threshold. 

84 correction: Learning-rate multiplier matching Adam update magnitude. 

85 """ 

86 defaults = {"lr": lr, "momentum": momentum, "nesterov": nesterov, "steps": steps, 

87 "eps": eps, "tau": tau, "correction": correction} 

88 super().__init__(params, defaults) 

89 for group in self.param_groups: 

90 self._validate_group(group) 

91 for param in group["params"]: 

92 self._validate_parameter(param) 

93 

94 @staticmethod 

95 def _validate_parameter(param: torch.Tensor) -> None: 

96 if param.ndim != 2 or min(param.shape) <= 0: 

97 raise ValueError("Sinkhorn parameters must be non-empty 2D matrices") 

98 for placement in getattr(param, "placements", ()): 

99 if not (placement.is_shard() or placement.is_replicate()): 

100 raise ValueError("Sinkhorn parameters require Shard or Replicate placements") 

101 

102 @staticmethod 

103 def _norm_groups(param: torch.Tensor) -> Tuple[Tuple[Any, ...], Tuple[Any, ...]]: 

104 """Read current parameter layout without parameter-identity caches.""" 

105 groups = ([], []) 

106 for axis, placement in enumerate(getattr(param, "placements", ())): 

107 if placement.is_shard(): 

108 groups[placement.dim % 2].append(param.device_mesh.get_group(axis)) 

109 return tuple(groups[0]), tuple(groups[1]) 

110 

111 @staticmethod 

112 def _validate_group(group: Dict[str, Any]) -> None: 

113 """Validate parameter-group overrides as well as constructor defaults.""" 

114 _validate_balancing(group["steps"], group["eps"], group["tau"]) 

115 for key in ("lr", "correction"): 

116 if not math.isfinite(group[key]) or group[key] < 0: 

117 raise ValueError(f"{key} must be finite and non-negative") 

118 if not math.isfinite(group["momentum"]) or not 0 <= group["momentum"] < 1: 

119 raise ValueError("momentum must be finite and in [0, 1)") 

120 if group.get("weight_decay", 0) != 0: 

121 raise ValueError("Sinkhorn does not apply weight decay") 

122 

123 def load_state_dict(self, state_dict: Dict[str, Any]) -> None: 

124 """Restore float32 momentum even when parameters are low precision. 

125 

126 Args: 

127 state_dict: Serialized parameter groups and optimizer state. 

128 """ 

129 super().load_state_dict(state_dict) 

130 for saved_group, group in zip(state_dict["param_groups"], self.param_groups): 

131 for saved_id, param in zip(saved_group["params"], group["params"]): 

132 momentum = state_dict["state"].get(saved_id, {}).get("momentum_buffer") 

133 if momentum is not None: 

134 self.state[param]["momentum_buffer"] = momentum.to( 

135 device=param.device, dtype=torch.float32, 

136 ).clone() 

137 

138 @staticmethod 

139 def _new_momentum(param: torch.Tensor) -> torch.Tensor: 

140 """Keep state layout metadata even inside SkipDTensorDispatch.""" 

141 local = torch.zeros_like(to_local_if_dtensor(param), dtype=torch.float32) 

142 if is_dtensor(param): 

143 return type(param).from_local(local, param.device_mesh, param.placements, 

144 shape=tuple(param.shape), stride=(param.shape[1], 1)) 

145 return local 

146 

147 def _update_parameter(self, param: torch.Tensor, group: Dict[str, Any]) -> None: 

148 """Update one local shard without materializing the full matrix.""" 

149 if param.grad is None: 

150 return 

151 self._validate_parameter(param) 

152 gradient = param.grad 

153 if hasattr(gradient, "placements") and any(place.is_partial() for place in gradient.placements): 

154 raise ValueError("Sinkhorn gradients must be reduced before optimizer.step()") 

155 gradient = to_local_if_dtensor(gradient) 

156 if gradient.is_sparse: 

157 raise ValueError("Sinkhorn does not support sparse gradients") 

158 state = self.state[param] 

159 if "momentum_buffer" not in state: 

160 state["momentum_buffer"] = self._new_momentum(param) 

161 buffer = to_local_if_dtensor(state["momentum_buffer"]) 

162 # Explicit allocation avoids same-dtype .to(copy=True) aliasing on NPU. 

163 # Reuse this workspace for Nesterov and balancing, preserving grad and EMA state. 

164 update = torch.empty_like(gradient, dtype=torch.float32).copy_(gradient) 

165 beta = group["momentum"] 

166 buffer.mul_(beta).add_(update, alpha=1 - beta) 

167 if group["nesterov"]: 

168 update.mul_(1 - beta).add_(buffer, alpha=beta) 

169 else: 

170 update.copy_(buffer) 

171 row_groups, column_groups = self._norm_groups(param) 

172 rows, columns = param.shape 

173 eps = group["eps"] 

174 row_norm = _distributed_norm(update, 1, column_groups) 

175 mean_norm = _sum_across(row_norm.sum(), row_groups) / rows 

176 update.masked_fill_(row_norm <= group["tau"] * mean_norm, 0) 

177 # Masking changes only rows that remain zero, so the initial norms can be reused. 

178 update.div_(row_norm.add_(eps)) 

179 for _ in range((group["steps"] - 1) // 2): 

180 column_norm = _distributed_norm(update, 0, row_groups) 

181 update.div_(column_norm.add_(eps)) 

182 row_norm = _distributed_norm(update, 1, column_groups) 

183 update.div_(row_norm.add_(eps)) 

184 update.mul_(math.sqrt(columns)) 

185 to_local_if_dtensor(param).add_(update, alpha=-group["lr"] * group["correction"]) 

186 

187 @torch.no_grad() 

188 def step(self, closure: Optional[Callable[[], Any]] = None) -> Any: 

189 """Apply balancing independently in each replica using shard statistics. 

190 

191 Args: 

192 closure: Optional callable reevaluating the training loss. 

193 

194 Returns: 

195 Closure loss when supplied, otherwise None. 

196 """ 

197 loss = None 

198 if closure is not None: 

199 with torch.enable_grad(): 

200 loss = closure() 

201 for group in self.param_groups: 

202 self._validate_group(group) 

203 group["step"] = (group.get("step") or 0) + 1 

204 for param in group["params"]: 

205 self._update_parameter(param, group) 

206 return loss