Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / utils.py: 46%
92 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"""Utility functions for Torch shard operations."""
17from typing import Any, Optional
19from torch import nn
20from torch import Tensor
21from torch.nn import Parameter
22from torch._ops import OpOverload, OpOverloadPacket
24from hyper_parallel.core.dtensor.dtensor import DTensor
25from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout
26from hyper_parallel.core.utils.communication import get_group_local_rank
27from hyper_parallel.core.tensor_parallel.loss_parallel import _get_loss_parallel_strict
28from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import (
29 _check_context_and_layout,
30 _get_local_tensor,
31 _get_mesh_and_dim,
32 _is_dtensor,
33 _is_shard_on_last_dim,
34 _validate_cross_entropy_params,
35 _validate_mesh_and_shard,
36)
39def get_op_name(func):
40 """Return the registry name for a Torch callable or operator overload."""
41 if hasattr(func, "__name__"):
42 return func.__name__
43 if isinstance(func, OpOverload):
44 return func.name.split("::")[-1].split(".")[0]
45 if isinstance(func, OpOverloadPacket):
46 return func.name.split("::")[-1]
47 func_str = str(func)
48 if "built-in function" in func_str:
49 return func_str.split()[-1].strip(">")
50 if "function" in func_str:
51 return func_str.split()[1]
52 return "unknown_op"
55def get_cell_construct(cell):
56 """Return the Torch module forward callable."""
57 return cell.forward
60def search_parameter_by_name(cell, param_name: str):
61 """Find a parameter and its owning module by dotted name."""
62 param_name = param_name.replace("self.", "")
63 if param_name in cell._parameters: # pylint: disable=protected-access
64 return cell, param_name, cell._parameters[param_name] # pylint: disable=protected-access
66 if "." in param_name:
67 cell_path, param_key = param_name.rsplit(".", 1)
68 try:
69 target_cell = cell.get_submodule(cell_path)
70 except AttributeError:
71 target_cell = None
72 if target_cell is not None and param_key in target_cell._parameters: # pylint: disable=protected-access
73 return target_cell, param_key, target_cell._parameters[param_key] # pylint: disable=protected-access
75 for _, child_cell in cell.named_children():
76 if isinstance(child_cell, nn.Module):
77 result = search_parameter_by_name(child_cell, param_name)
78 if result is not None:
79 return result
80 return None
83def set_layout_into_parameter(param, layout):
84 """Convert a local Torch parameter to a DTensor-backed parameter."""
85 if isinstance(param, DTensor):
86 raise ValueError(f"Parameter {param} has been configured layout, cannot be set repeatedly.")
87 requires_grad = param.requires_grad
88 param_dtensor = DTensor.from_local(
89 _get_slice_tensor_by_layout(param, layout),
90 layout.mesh,
91 layout.alias_placements,
92 )
93 return Parameter(param_dtensor, requires_grad=requires_grad)
96def update_parameter_by_name(result: tuple, new_param) -> bool:
97 """Replace a parameter on its owning Torch module."""
98 parent_cell, param_key, _ = result
99 if param_key in parent_cell._parameters: # pylint: disable=protected-access
100 parent_cell._parameters[param_key] = new_param # pylint: disable=protected-access
101 else:
102 parent_cell.register_parameter(param_key, new_param)
103 return True
106def _is_floating_torch(tensor: Tensor) -> bool:
107 """Check if PyTorch tensor is floating point."""
108 return tensor.is_floating_point()
111def distributed_cross_entropy(
112 input_tensor: Tensor,
113 target: Tensor,
114 weight: Optional[Tensor] = None,
115 size_average: Optional[bool] = None,
116 ignore_index: int = -100,
117 reduce: Optional[bool] = None,
118 reduction: str = "mean",
119 label_smoothing: float = 0.0,
120) -> Tensor:
121 """Distributed cross_entropy entry used by shard dispatch."""
122 # Defer the components import to preserve the lightweight models import boundary.
123 from hyper_parallel.components.losses._vocab_parallel_cross_entropy import ( # pylint: disable=C0415
124 DistributedCrossEntropyFunction,
125 )
127 input_dtensor = None
128 mesh = None
129 vocab_size = None
131 if _is_dtensor(input_tensor):
132 if not _is_shard_on_last_dim(input_tensor):
133 raise ValueError(
134 "input must be Shard(-1) on class dimension. "
135 f"Got placements: {input_tensor.placements}"
136 )
137 input_dtensor = input_tensor
138 mesh, _ = _get_mesh_and_dim(input_tensor)
139 vocab_size = input_tensor.shape[-1]
141 input_for_check = input_dtensor if input_dtensor is not None else input_tensor
142 _check_context_and_layout(input_for_check)
143 _validate_cross_entropy_params(
144 input_tensor,
145 target,
146 weight,
147 size_average,
148 ignore_index,
149 reduce,
150 reduction,
151 label_smoothing,
152 _is_floating_torch,
153 )
155 if input_dtensor is None:
156 raise ValueError(
157 "input must be a DTensor when using loss_parallel. "
158 f"Got type: {type(input_tensor)}"
159 )
161 input_local = _get_local_tensor(input_dtensor)
162 local_vocab_size = input_local.shape[-1]
163 if input_dtensor.ndim > 2:
164 input_local = input_local.reshape(-1, local_vocab_size)
165 target = target.reshape(-1)
167 strict = _get_loss_parallel_strict()
168 _validate_mesh_and_shard(input_dtensor, strict)
169 return DistributedCrossEntropyFunction.apply(
170 input_local,
171 target,
172 weight,
173 ignore_index,
174 reduction,
175 vocab_size,
176 mesh,
177 0,
178 )
181def distributed_cross_entropy_from_op_call(
182 op_call: Any, # pylint: disable=unused-argument
183 args: tuple,
184 kwargs: dict,
185):
186 """Parse a cross_entropy op call and invoke the distributed Torch implementation."""
187 input_tensor = args[0] if len(args) > 0 else kwargs.get("input")
188 target = args[1] if len(args) > 1 else kwargs.get("target")
189 weight = args[2] if len(args) > 2 else kwargs.get("weight")
190 size_average = args[3] if len(args) > 3 else kwargs.get("size_average")
191 ignore_index = args[4] if len(args) > 4 else kwargs.get("ignore_index", -100)
192 reduce = args[5] if len(args) > 5 else kwargs.get("reduce")
193 reduction = args[6] if len(args) > 6 else kwargs.get("reduction", "mean")
194 label_smoothing = args[7] if len(args) > 7 else kwargs.get("label_smoothing", 0.0)
196 return distributed_cross_entropy(
197 input_tensor=input_tensor,
198 target=target,
199 weight=weight,
200 size_average=size_average,
201 ignore_index=ignore_index,
202 reduce=reduce,
203 reduction=reduction,
204 label_smoothing=label_smoothing,
205 )
208__all__ = [
209 # Op-name resolution.
210 "get_op_name",
211 # Cell/parameter plumbing used by the shard API.
212 "get_cell_construct",
213 "search_parameter_by_name",
214 "set_layout_into_parameter",
215 "update_parameter_by_name",
216 # Distributed cross-entropy entry points.
217 "distributed_cross_entropy",
218 "distributed_cross_entropy_from_op_call",
219 # Re-exported from :mod:`hyper_parallel.core.utils.communication`.
220 "get_group_local_rank",
221]