Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / api.py: 98%
260 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-26 05:43 +0800
1# Copyright 2025 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"""shard"""
16import inspect
17from typing import Union, Callable, Dict, List
18from functools import wraps
20import torch.distributed as dist
21from torch import nn
23from hyper_parallel.core.dtensor.layout import Layout, DeviceMesh
24from hyper_parallel.core.dtensor.dtensor import DTensor, _is_alias_placements
25from hyper_parallel.core.dtensor.placement_types import Placement
26from hyper_parallel.core.shard.utils import (
27 get_cell_construct,
28 search_parameter_by_name,
29 set_layout_into_parameter,
30 update_parameter_by_name,
31)
32from hyper_parallel.core.shard.sharding_plan import ShardingPlan
35def _has_kwargs(func):
36 """_has_kwargs"""
37 sig = inspect.signature(func)
38 return any(
39 param.default != inspect.Parameter.empty
40 for param in sig.parameters.values()
41 )
44def _get_param_name(func):
45 """_get_param_name"""
46 sig = inspect.signature(func)
47 return list(sig.parameters.keys())
50def _convert_sharding_plan(sharding_plan: Dict, device_mesh: DeviceMesh) -> Dict:
51 """
52 Convert sharding_plan values to Layout objects.
54 This function recursively traverses the sharding_plan and converts
55 placement tuples (e.g., (Shard(0), Replicate())) to Layout objects.
57 Args:
58 sharding_plan: The original sharding plan with tuple specifications
59 device_mesh: The DeviceMesh to use for conversion
61 Returns:
62 Dict: Converted sharding plan with Layout objects
63 """
65 def _is_placement_tuple(value):
66 """Check if value is a placement specification tuple.
68 A placement tuple contains Placement instances (Shard, Replicate) or
69 alias strings ("dp", "None"). It should NOT be a tuple of placement tuples.
71 Examples of placement tuples:
72 (Shard(0), Replicate(), Shard(1)) -> True
73 ("dp", "tp", "None") -> True
74 (("dp", "tp"), "None") -> True (multi-axis sharding)
76 Examples of NON-placement tuples:
77 ((Shard(0),), (Shard(1),)) -> False (tuple of placement tuples)
78 """
79 if not isinstance(value, tuple) or len(value) == 0:
80 return False
82 for item in value:
83 # Placement instance is valid
84 if isinstance(item, Placement):
85 continue
86 # String (alias name) is valid
87 if isinstance(item, str):
88 continue
89 # Nested tuple needs special handling
90 if isinstance(item, tuple):
91 # Nested tuple of strings is valid (multi-axis sharding)
92 if len(item) > 0 and all(isinstance(x, str) for x in item):
93 continue
94 # Nested tuple containing Placement means this is a tuple of placement tuples
95 if len(item) > 0 and any(isinstance(x, Placement) for x in item):
96 return False
97 # Empty tuple or other cases - not valid
98 return False
99 # Any other type is not valid in a placement tuple
100 return False
102 return True
104 def _to_layout(value):
105 """Convert a single sharding specification to Layout."""
106 layout = Layout.from_device_mesh(device_mesh)
107 if _is_alias_placements(value):
108 result = layout(*value)
109 else:
110 result = layout(value)
111 return result
113 def _convert_value(value, wrap_single_as_list=False):
114 """Recursively convert value based on its structure."""
115 if value is None:
116 return None
118 # Case 1: It's a placement tuple - convert to Layout
119 if _is_placement_tuple(value):
120 layout = _to_layout(value)
121 # Wrap single layout in list if required (for input/output in forward)
122 return [layout] if wrap_single_as_list else layout
124 # Case 2: It's a dict - recursively process each value
125 if isinstance(value, dict):
126 converted_dict = {}
127 for k, v in value.items():
128 converted_dict[k] = _convert_value(v, wrap_single_as_list=False)
129 return converted_dict
131 # Case 3: It's a list - recursively process each element
132 if isinstance(value, list):
133 return [_convert_value(v, wrap_single_as_list=False) for v in value]
135 # Case 4: It's a tuple but not a placement tuple - treat as list
136 if isinstance(value, tuple):
137 return [_convert_value(v, wrap_single_as_list=False) for v in value]
139 # Case 5: Other types (e.g., primitives) - return as is
140 return value
142 def _convert_forward_plan(forward_plan):
143 """Convert forward plan with special handling for input/output."""
144 if forward_plan is None:
145 return None
147 converted = {}
148 for key, value in forward_plan.items():
149 if key.endswith("input") or key.endswith("output"):
150 # input/output need special handling:
151 # - dict format: convert each value, keep as dict
152 # - list/tuple format: convert each element, keep as list
153 # - single placement tuple: convert and wrap in list
154 if value is None:
155 converted[key] = None
156 elif isinstance(value, dict):
157 # Dict format for kwargs: {"x": placements, "activation": placements}
158 converted[key] = {k: _convert_value(v) for k, v in value.items()}
159 elif isinstance(value, (list, tuple)):
160 # Check if it's a single placement tuple or a list/tuple of placement tuples
161 if _is_placement_tuple(value):
162 # Single placement tuple - wrap in list
163 converted[key] = [_to_layout(value)]
164 else:
165 # List/tuple of placements for multiple positional args
166 converted[key] = [_convert_value(v) for v in value]
167 else:
168 converted[key] = _convert_value(value, wrap_single_as_list=True)
169 else:
170 # Other keys in forward plan
171 converted[key] = _convert_value(value)
172 return converted
174 # Main conversion logic
175 converted_plan = {}
177 for key, value in sharding_plan.items():
178 if key == "forward":
179 converted_plan[key] = _convert_forward_plan(value)
180 elif key.endswith("input") or key.endswith("output"):
181 # Top-level input/output (for callable sharding)
182 if value is None:
183 converted_plan[key] = None
184 elif isinstance(value, dict):
185 converted_plan[key] = {k: _convert_value(v) for k, v in value.items()}
186 elif isinstance(value, (list, tuple)):
187 if _is_placement_tuple(value):
188 converted_plan[key] = [_to_layout(value)]
189 else:
190 converted_plan[key] = [_convert_value(v) for v in value]
191 else:
192 converted_plan[key] = _convert_value(value, wrap_single_as_list=True)
193 else:
194 # parameter and other keys - use standard recursive conversion
195 converted_plan[key] = _convert_value(value)
197 return converted_plan
200def _parallel_in(func, args, kwargs, layouts):
201 """_parallel_in"""
202 if not isinstance(layouts, (list, dict, tuple)):
203 raise ValueError(f"The in_layout must be a list, tuple or dict, but got {type(layouts)}.")
205 params_name = _get_param_name(func)
206 processed_args = list(args)
207 processed_kwargs = dict(kwargs)
209 def _get_layout(index, is_list):
210 """_get_layout"""
211 if is_list:
212 return layouts[index]
213 param_name = params_name[index]
214 return layouts[param_name]
216 is_list = isinstance(layouts, (list, tuple))
217 for i, arg in enumerate(args):
218 if not isinstance(arg, DTensor):
219 continue
221 to_layout = _get_layout(i, is_list)
222 processed_args[i] = arg.redistribute(to_layout.mesh, to_layout.alias_placements)
223 for k, v in kwargs.items():
224 if not isinstance(v, DTensor) or layouts.get(k) is None:
225 processed_kwargs[k] = v
226 continue
227 to_layout = layouts[k]
228 processed_kwargs[k] = v.redistribute(to_layout.mesh, to_layout.alias_placements)
230 return tuple(processed_args), processed_kwargs
233def _parallel_out(outputs, layouts):
234 """_parallel_out"""
235 if not isinstance(layouts, (list, tuple)):
236 raise ValueError(f"The out_layout must be a list or tuple, but got {type(layouts)}.")
237 if isinstance(outputs, (tuple, list)):
238 if len(outputs) != len(layouts):
239 raise ValueError(f"The size of outputs and out_layout must be equal, but got {len(outputs)} and "
240 f"{len(layouts)}")
241 new_outputs = []
242 for i, arg in enumerate(outputs):
243 if not isinstance(arg, DTensor) or arg is None:
244 new_outputs.append(arg)
245 continue
246 to_layout = layouts[i]
247 new_outputs.append(arg.redistribute(to_layout.mesh, to_layout.alias_placements))
248 return tuple(new_outputs)
249 if len(layouts) != 1:
250 raise ValueError(f"The size of outputs and out_layout must be equal, but got 1 and "
251 f"{len(layouts)}")
253 if isinstance(outputs, DTensor):
254 return outputs.redistribute(
255 layouts[0].mesh, layouts[0].alias_placements)
256 return outputs
259def _forward_pre_hook(cell, args):
260 """_forward_pre_hook"""
261 if cell.in_layout is None:
262 return args
263 processed_args, _ = _parallel_in(get_cell_construct(cell), args, {}, cell.in_layout)
264 return processed_args
267def _forward_pre_with_kwargs_hook(cell, args, kwargs):
268 """_forward_pre_with_kwargs_hook"""
269 if cell.in_layout is None:
270 return args, kwargs
271 return _parallel_in(get_cell_construct(cell), args, kwargs, cell.in_layout)
274def _forward_hook(cell, inputs, outputs): # pylint: disable=unused-argument
275 """_forward_hook"""
276 if cell.out_layout is None:
277 return outputs
278 return _parallel_out(outputs, cell.out_layout)
281def _forward_with_kwargs_hook(cell, inputs, kwargs, outputs): # pylint: disable=unused-argument
282 """_forward_with_kwargs_hook"""
283 return _forward_hook(cell, inputs, outputs)
286def _split_sharding_plan_key(key: str) -> tuple[str, str]:
287 """Split a sharding plan key into its module prefix and layout suffix."""
288 if '.' not in key:
289 return "", key
290 return tuple(key.rsplit('.', 1))
293def _register_hook(model: nn.Module, sharding_plan: Dict):
294 """_register_hook"""
296 def _register_cell_hook(model, has_inputs_layout, has_outputs_layout):
297 """_register_cell_hook"""
298 has_kwargs = _has_kwargs(get_cell_construct(model))
299 pre_hook = _forward_pre_with_kwargs_hook if has_kwargs else _forward_pre_hook
300 hook = _forward_with_kwargs_hook if has_kwargs else _forward_hook
301 if has_inputs_layout:
302 model.register_forward_pre_hook(pre_hook, with_kwargs=has_kwargs)
304 if has_outputs_layout:
305 model.register_forward_hook(hook, with_kwargs=has_kwargs)
307 def _set_layouts(model, layouts, set_inputs_layout, set_outputs_layout):
308 """_set_layouts"""
309 if set_inputs_layout:
310 model.in_layout = layouts
312 if set_outputs_layout:
313 model.out_layout = layouts
315 cell_dict = {}
316 for name, cell in model.named_modules():
317 cell_dict[name] = cell
319 valid_suffix = ["input", "output"]
320 for key, value in sharding_plan.items():
321 if value is None:
322 continue
323 prefix, suffix = _split_sharding_plan_key(key)
324 if suffix not in valid_suffix:
325 raise ValueError(f"In python shard_module, sharding_plan's forward key must end with input or output, "
326 f"but got type {suffix}")
328 set_inputs_layout = suffix == "input"
329 set_outputs_layout = not set_inputs_layout
330 register_cell = cell_dict.get(prefix)
331 if register_cell is None:
332 raise ValueError(f"Cannot find target cell {prefix!r} in sharding_plan")
334 _set_layouts(register_cell, value, set_inputs_layout, set_outputs_layout)
335 _register_cell_hook(register_cell, set_inputs_layout, set_outputs_layout)
338def _register_local_tensor_hook(cell: nn.Module, return_local_tensor_list: List[str]):
339 """_register_local_tensor_hook"""
341 def hook_func(cell, inputs, outputs): # pylint: disable=unused-argument
342 def _recursive_to_local(out):
343 if isinstance(out, (tuple, list)):
344 new_out = []
345 for item in out:
346 new_out.append(_recursive_to_local(item))
347 return tuple(new_out) if isinstance(out, tuple) else new_out
348 if isinstance(out, DTensor):
349 return out.to_local()
350 return out
352 return _recursive_to_local(outputs)
354 cell_dict = {}
355 for name, sub_cell in cell.named_modules():
356 cell_dict[name] = sub_cell
358 for cell_name in return_local_tensor_list:
359 register_cell = cell_dict.get(cell_name)
360 if register_cell is None:
361 raise KeyError(f"Cannot find cell {cell_name!r} in sharded cell.")
362 register_cell.register_forward_hook(hook_func)
365def _shard_callable(func: Callable, sharding_plan: Dict):
366 """_shard_callable"""
367 forward_sharding_plan = sharding_plan.get("forward")
368 if forward_sharding_plan is None:
369 return func
371 @wraps(func)
372 def _shard_wrapper(*args, **kwargs):
373 """_shard_wrapper"""
374 input_layout = forward_sharding_plan.get("input")
375 output_layout = forward_sharding_plan.get("output")
376 if input_layout is not None:
377 args, kwargs = _parallel_in(func, args, kwargs, input_layout)
378 outputs = func(*args, **kwargs)
379 if output_layout is not None:
380 outputs = _parallel_out(outputs, output_layout)
381 return outputs
383 return _shard_wrapper
386def shard_module(model: Union[nn.Module, Callable], device_mesh: DeviceMesh, sharding_plan: ShardingPlan):
387 """
388 Defining the input, output and parameters layouts of this cell or Callable.
390 Note:
391 - It is valid only in pynative mode.
393 .. warning::
394 The method is currently not supported in Graph mode.
396 Args:
397 model (Module or Callable): The model to be sharded.
398 device_mesh (DeviceMesh): The device mesh for sharding.
399 sharding_plan (ShardingPlan): Define the layout for the specified parameters, inputs or outputs.
400 The sharding specification can be:
401 - tuple of strings for alias format, e.g., ("dp", "None")
402 - tuple of Placements, e.g., (Shard(0), Replicate())
404 Returns:
405 Module or Callable: The sharded model.
407 Examples:
408 >>> # Usage with device_mesh and alias format
409 >>> mesh = DeviceMesh("npu", (2, 2), nesh_dim_names=("dp", "tp"))
410 >>> sharding_plan = ShardingPlan(
411 ... plan={"mlp.weight": ("None", "tp")},
412 ... input_plan={"input": ("dp", "None")},
413 ... output_plan={"output": ("dp", "tp")}
414 ... )
415 >>> model = shard_module(model, mesh, sharding_plan)
417 >>> # Usage with device_mesh and Placement format
418 >>> mesh = DeviceMesh("npu", (2, 2), nesh_dim_names=("dp", "tp"))
419 >>> sharding_plan = ShardingPlan(
420 ... plan={"mlp.weight": (Replicate(), Shard(1))},
421 ... input_plan={"input": (Shard(0), Replicate())},
422 ... output_plan={"output": (Shard(0), Shard(1))}
423 ... )
424 >>> model = shard_module(model, mesh, sharding_plan)
425 """
426 if dist.get_world_size() == 1:
427 return None
429 if not isinstance(sharding_plan, ShardingPlan):
430 raise TypeError(f"The 'sharding_plan' must be an instance of ShardingPlan, "
431 f"but got {type(sharding_plan)}. Direct dict input is not supported.")
433 normalized_plan = {}
434 return_local_tensor_list = None
436 if sharding_plan.plan:
437 normalized_plan["parameter"] = sharding_plan.plan
439 forward_part = {}
441 if sharding_plan.input_plan:
442 if not isinstance(sharding_plan.input_plan, dict):
443 raise TypeError(f"input_plan must be a dict, but got {type(sharding_plan.input_plan)}")
444 forward_part.update(sharding_plan.input_plan)
446 if sharding_plan.output_plan:
447 if not isinstance(sharding_plan.output_plan, dict):
448 raise TypeError(f"output_plan must be a dict, but got {type(sharding_plan.output_plan)}")
449 forward_part.update(sharding_plan.output_plan)
451 if forward_part:
452 normalized_plan["forward"] = forward_part
454 if sharding_plan.return_local_tensor:
455 return_local_tensor_list = sharding_plan.return_local_tensor
457 # Convert sharding_plan to Layout objects
458 converted_plan = _convert_sharding_plan(normalized_plan, device_mesh)
460 if not isinstance(model, nn.Module):
461 return _shard_callable(model, converted_plan)
463 param_sharding_plan = converted_plan.get("parameter")
464 forward_sharding_plan = converted_plan.get("forward")
466 if param_sharding_plan is not None:
467 for param_name, layout in param_sharding_plan.items():
468 if not isinstance(layout, Layout):
469 raise ValueError(f"In python shard_module, the type of setting in parameter_plan must be Layout, "
470 f"but got type {type(layout)}")
471 result = search_parameter_by_name(model, param_name)
472 if not result:
473 raise ValueError(f"{param_name} is configured with a layout, but no instance was found.")
474 _, _, param = result
475 if layout.tensor_map is None:
476 layout.placement_to_tensor_map(param.dim())
477 param = set_layout_into_parameter(param, layout)
478 update_parameter_by_name(result, param)
480 if forward_sharding_plan is not None:
481 _register_hook(model, forward_sharding_plan)
483 if return_local_tensor_list is not None:
484 _register_local_tensor_hook(model, return_local_tensor_list)
486 return model