Star139
197
代码
代码
Issues344
Pull Requests268
项目讨论
Wiki
分析
Star139
197
已合并
swap_adam_optimizer #982
songjiaqi创建于 6月29日
swap_adam_optimizer #982
已合并
songjiaqi创建于 6月29日
从silkage_jiajia/hyper-parallelswap_adam_optimizer合入到mindspore/hyper-parallelmaster
共 22 个文件变更+8866-4
Mhyper_parallel/core/optimizer/__init__.py+17-4文件内容审核中，请稍后刷新重试
Ahyper_parallel/core/optimizer/swap_optimizer.py+126-0
@@ -0,0 +1,126 @@
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+"""Public API for optimizer state swap."""
16+ 
17+from __future__ import annotations
18+ 
19+from dataclasses import dataclass, field, replace
20+from typing import Any, Optional, Sequence
21+ 
22+from hyper_parallel.core.optimizer.swap_optimizer_base import validate_state_keys
23+from hyper_parallel.platform import get_platform
24+from hyper_parallel.platform.platform import PlatformType
25+ 
26+ 
27+def _default_packed_swap() -> bool:
28+ """Return the packed swap default for the active backend."""
29+ return get_platform().platform_type == PlatformType.PYTORCH
30+ 
31+ 
32+def _is_mindformers_adamw(optimizer: Any) -> bool:
33+ """Return whether an optimizer is the supported MindFormers AdamW type."""
34+ optimizer_type = type(optimizer)
35+ return (
36+ optimizer_type.__name__ == "AdamW"
37+ and optimizer_type.__module__ == "mindformers.pynative.optimizer.adamw"
38+ )
39+ 
40+ 
41+def _resolve_packed_swap_default(
42+ optimizer: Any,
43+ config: "SwapOptimizerConfig",
44+ platform_type: PlatformType,
45+) -> "SwapOptimizerConfig":
46+ """Enable packed swap by default for MindFormers AdamW on MindSpore."""
47+ if (
48+ platform_type == PlatformType.MINDSPORE
49+ and not config.packed_swap_was_explicit
50+ and _is_mindformers_adamw(optimizer)
51+ ):
52+ return replace(config, packed_swap=True)
53+ return config
54+ 
55+ 
56+@dataclass(frozen=True)
57+class SwapOptimizerConfig:
58+ """Configuration for Adam/AdamW optimizer state swap.
59+ 
60+ The runtime uses a fixed one-batch-ahead prefetch pipeline.
61+ 
62+ Args:
63+ swap_times: Number of pipeline partitions.
64+ state_keys: Logical state keys to swap. ``None`` uses adapter defaults.
65+ min_numel: Tensor states smaller than this element count are not swapped.
66+ include_master_params: Whether optimizer-owned fp32 master params are swapped.
67+ packed_swap: Whether supported backends use two packed A/B staging buffers.
68+ Defaults to ``True`` on PyTorch and for MindFormers AdamW on
69+ MindSpore; other MindSpore optimizers default to ``False``. When
70+ ``False``, optimizer states are swapped tensor by tensor. An
71+ explicit value always takes precedence over these defaults.
72+ """
73+ 
74+ swap_times: int = 16
75+ state_keys: Optional[Sequence[str]] = None
76+ min_numel: int = 1024
77+ include_master_params: bool = False
78+ packed_swap: Optional[bool] = None
79+ _packed_swap_explicit: bool = field(init=False, repr=False, compare=False)
80+ 
81+ def __post_init__(self) -> None:
82+ packed_swap_explicit = self.packed_swap is not None
83+ object.__setattr__(self, "_packed_swap_explicit", packed_swap_explicit)
84+ if not packed_swap_explicit:
85+ object.__setattr__(self, "packed_swap", _default_packed_swap())
86+ if self.swap_times <= 0:
87+ raise ValueError("SwapOptimizerConfig.swap_times must be positive.")
88+ if self.min_numel < 0:
89+ raise ValueError("SwapOptimizerConfig.min_numel must be non-negative.")
90+ object.__setattr__(self, "state_keys", validate_state_keys(self.state_keys))
91+ 
92+ @property
93+ def packed_swap_was_explicit(self) -> bool:
94+ """Return whether ``packed_swap`` was explicitly supplied by the caller."""
95+ return self._packed_swap_explicit
96+ 
97+ 
98+class SwapOptimizer:
99+ """Core facade that dispatches to the active backend implementation."""
100+ 
101+ def __new__(cls, optimizer: Any, config: Optional[SwapOptimizerConfig] = None):
102+ return swap_optimizer(optimizer, config)
103+ 
104+ 
105+def swap_optimizer(optimizer: Any, config: Optional[SwapOptimizerConfig] = None) -> Any:
106+ """Wrap a supported Adam/AdamW optimizer with optimizer-state swap.
107+ 
108+ Args:
109+ optimizer: Base optimizer instance.
110+ config: Swap optimizer configuration.
111+ 
112+ Returns:
113+ Backend-specific swap optimizer wrapper.
114+ 
115+ Raises:
116+ ValueError: If the active backend or optimizer type is unsupported.
117+ """
118+ platform = get_platform()
119+ cfg = config or SwapOptimizerConfig()
120+ cfg = _resolve_packed_swap_default(optimizer, cfg, platform.platform_type)
121+ return platform.get_swap_optimizer()(optimizer, cfg)
122+ 
123+ 
124+def is_swap_optimizer(optimizer: Any) -> bool:
125+ """Return whether ``optimizer`` is a swap optimizer wrapper."""
126+ return bool(getattr(optimizer, "_is_swap_optimizer", False))
Ahyper_parallel/core/optimizer/swap_optimizer_base.py+519-0
@@ -0,0 +1,519 @@
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+"""Common abstractions for optimizer state swap."""
16+ 
17+from __future__ import annotations
18+ 
19+import contextlib
20+from dataclasses import dataclass
21+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence
22+ 
23+ 
24+STATE_KEYS = ("exp_avg", "exp_avg_sq", "max_exp_avg_sq")
25+MASTER_PARAM_KEY = "master_param"
26+SUPPORTED_STATE_KEYS = STATE_KEYS + (MASTER_PARAM_KEY,)
27+ 
28+ 
29+@dataclass
30+class SwapSlot:
31+ """One logical optimizer state tensor that may be swapped."""
32+ 
33+ name: str
34+ tensor: Any
35+ cpu_tensor: Optional[Any] = None
36+ storage_nbytes: int = 0
37+ swappable: bool = True
38+ state: str = "device"
39+ event: Optional[Any] = None
40+ shape: tuple[int, ...] = ()
41+ dtype: Optional[Any] = None
42+ device: Optional[Any] = None
43+ numel: int = 0
44+ host_offset: int = 0
45+ packed: bool = False
46+ logical_tensor: Optional[Any] = None
47+ 
48+ def bind_tensor(self, tensor: Any) -> None:
49+ """Bind the slot to a host or staging tensor view."""
50+ device = getattr(tensor, "device", None)
51+ device_type = getattr(device, "type", None)
52+ if self.logical_tensor is not None and device_type != "cpu":
53+ # DTensor exposes no public setter for replacing its local shard.
54+ setattr(self.logical_tensor, "_local_tensor", tensor)
55+ self.logical_tensor.data = tensor
56+ self.tensor = self.logical_tensor
57+ else:
58+ self.tensor = tensor
59+ 
60+ 
61+ @property
62+ def checkpoint_tensor(self) -> Any:
63+ """Return the CPU tensor when present, otherwise the live tensor."""
64+ return self.cpu_tensor if self.cpu_tensor is not None else self.tensor
65+ 
66+ 
67+@dataclass
68+class UpdateUnit:
69+ """Per-parameter optimizer update unit used by the pipeline runtime.
70+ 
71+ ``adapter_index`` is interpreted by the backend adapter: it identifies a
72+ Torch parameter group or a MindSpore parameter/state entry.
73+ """
74+ 
75+ adapter_index: int
76+ param: Any
77+ grad: Any
78+ slots: List[SwapSlot]
79+ 
80+ 
81+class OptimizerSwapAdapter:
82+ """Backend optimizer adapter interface."""
83+ 
84+ def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
85+ self.optimizer = optimizer
86+ self.config = config
87+ self.runtime = runtime
88+ 
89+ @classmethod
90+ def matches(cls, optimizer: Any) -> bool:
91+ """Return whether this adapter supports ``optimizer``."""
92+ raise NotImplementedError
93+ 
94+ def validate(self) -> None:
95+ """Validate optimizer flags and unsupported modes."""
96+ raise NotImplementedError
97+ 
98+ def prepare_step(self, *args: Any, **kwargs: Any) -> Any:
99+ """Prepare one outer optimizer step."""
100+ raise NotImplementedError
Mhyper_parallel/platform/mindspore/platform.py+8-0
@@ -758,6 +758,14 @@ class MindSporePlatform(Platform):
758 self._custom_ops_cls = MindSporeCustomOps758 self._custom_ops_cls = MindSporeCustomOps
759 return self._custom_ops_cls759 return self._custom_ops_cls
760 760 
761+ @staticmethod
762+ def get_swap_optimizer():
763+ """Return the MindSpore optimizer-state swap wrapper class."""
764+ from hyper_parallel.platform.mindspore.swap_optimizer.swap_optimizer import ( # pylint: disable=import-outside-toplevel
765+ get_swap_optimizer,
766+ )
767+ return get_swap_optimizer()
768+ 
761 def __init__(self):769 def __init__(self):
762 # Ensure MindSpore ``nn.Cell.to_empty`` is patched as soon as the770 # Ensure MindSpore ``nn.Cell.to_empty`` is patched as soon as the
763 # MindSpore platform instance is created.771 # MindSpore platform instance is created.
Ahyper_parallel/platform/mindspore/swap_optimizer/__init__.py+19-0
@@ -0,0 +1,19 @@
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+"""MindSpore optimizer swap backend."""
16+ 
17+from hyper_parallel.platform.mindspore.swap_optimizer.swap_optimizer import MindSporeSwapOptimizer
18+ 
19+__all__ = ["MindSporeSwapOptimizer"]
Ahyper_parallel/platform/mindspore/swap_optimizer/adapters.py+685-0
@@ -0,0 +1,685 @@
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+"""MindSpore Adam/AdamW swap optimizer adapters."""
16+# pylint: disable=protected-access
17+ 
18+from __future__ import annotations
19+ 
20+import importlib
21+from typing import Any, Dict, Iterable, List, Tuple
22+ 
23+import mindspore as ms
24+from mindspore import nn
25+from mindspore.common import dtype as mstype
26+from mindspore.ops import functional as F
27+ 
28+from hyper_parallel.core.dtensor.dtensor import SkipDTensorDispatch
29+from hyper_parallel.core.optimizer.swap_optimizer_base import (
30+ OptimizerSwapAdapter,
31+ SUPPORTED_STATE_KEYS,
32+ SwapSlot,
33+ UpdateUnit,
34+)
35+ 
36+ 
37+def _to_tuple(value: Any) -> Tuple[Any, ...]:
38+ if isinstance(value, tuple):
39+ return value
40+ if isinstance(value, list):
41+ return tuple(value)
42+ return tuple(value)
43+ 
44+ 
45+class MindSporeAdamBaseAdapter(OptimizerSwapAdapter):
46+ """Common MindSpore optimizer adapter logic."""
47+ 
48+ def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
49+ super().__init__(optimizer, config, runtime)
50+ self._slots: Dict[Tuple[int, str], SwapSlot] = {}
51+ 
52+ def validate(self) -> None:
53+ """Base validation."""
54+ if getattr(self.optimizer, "use_parallel", False):
55+ raise ValueError("MindSpore swap optimizer does not support parallel optimizer yet.")
56+ 
57+ def iter_update_units(self, step_context: Dict[str, Any]) -> List[UpdateUnit]:
58+ """Return units collected in prepare_step."""
59+ return step_context["units"]
60+ 
61+ def all_slots(self) -> Iterable[SwapSlot]:
62+ """Iterate known slots."""
63+ return tuple(self._slots.values())
64+ 
65+ def initial_slots(self) -> Iterable[SwapSlot]:
66+ """Build optimizer state slots that can be offloaded before the first update."""
67+ return self._checkpoint_slots()
68+ 
69+ def packed_layout_units(self) -> List[UpdateUnit]:
70+ """Return stable optimizer units used to build the packed host layout."""
71+ return []
72+ 
73+ def checkpoint_state_dict(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
74+ """Return checkpoint-safe optimizer state dict."""
75+ del args, kwargs
76+ state = self._state_dict()
77+ slot_by_name = self._checkpoint_slot_map()
78+ for name, slot in slot_by_name.items():
79+ if name not in state or not slot.swappable:
80+ continue
81+ if slot.cpu_tensor is None:
82+ if slot.state == "host":
83+ raise RuntimeError(f"Swap slot {slot.name!r} is host-resident but has no CPU mirror.")
84+ slot.cpu_tensor = self.runtime.make_cpu_tensor(slot.tensor)
85+ state[name] = ms.Parameter(self.runtime.make_cpu_tensor(slot.cpu_tensor), name=name)
86+ return state
87+ 
88+ def load_checkpoint_state_dict(
89+ self,
90+ state_dict: Dict[str, Any],
91+ *args: Any,
92+ **kwargs: Any,
93+ ) -> None:
94+ """Load checkpoint-safe parameter dict."""
95+ del args, kwargs
96+ slot_by_name = self._checkpoint_slot_map(promote_checkpoint_swappable=True)
97+ remaining = dict(state_dict)
98+ for name, slot in slot_by_name.items():
99+ if name not in remaining or not slot.swappable:
100+ continue
Mhyper_parallel/platform/platform.py+5-0
@@ -149,6 +149,11 @@ class Platform:
149 "Platform subclasses must implement custom_ops"149 "Platform subclasses must implement custom_ops"
150 )150 )
151 151 
152+ @staticmethod
153+ def get_swap_optimizer():
154+ """Return the active backend's optimizer-state swap wrapper class."""
155+ raise NotImplementedError("Platform subclasses must implement get_swap_optimizer")
156+ 
152 @staticmethod157 @staticmethod
153 def get_rank():158 def get_rank():
154 """Get the rank of the current process in the default process group.159 """Get the rank of the current process in the default process group.
Mhyper_parallel/platform/torch/platform.py+8-0
@@ -461,6 +461,14 @@ class TorchPlatform(Platform):
461 self._custom_ops_cls = TorchCustomOps461 self._custom_ops_cls = TorchCustomOps
462 return self._custom_ops_cls462 return self._custom_ops_cls
463 463 
464+ @staticmethod
465+ def get_swap_optimizer():
466+ """Return the Torch optimizer-state swap wrapper class."""
467+ from hyper_parallel.platform.torch.swap_optimizer.swap_optimizer import ( # pylint: disable=import-outside-toplevel
468+ get_swap_optimizer,
469+ )
470+ return get_swap_optimizer()
471+ 
464 @staticmethod472 @staticmethod
465 def is_linear_module(module) -> bool:473 def is_linear_module(module) -> bool:
466 """Check whether *module* is a ``torch.nn.Linear`` instance."""474 """Check whether *module* is a ``torch.nn.Linear`` instance."""
Ahyper_parallel/platform/torch/swap_optimizer/__init__.py+19-0
@@ -0,0 +1,19 @@
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+"""Torch optimizer swap backend."""
16+ 
17+from hyper_parallel.platform.torch.swap_optimizer.swap_optimizer import TorchSwapOptimizer
18+ 
19+__all__ = ["TorchSwapOptimizer"]
Ahyper_parallel/platform/torch/swap_optimizer/adapters.py+621-0
@@ -0,0 +1,621 @@
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+"""Torch Adam/AdamW swap optimizer adapters."""
16+# pylint: disable=protected-access
17+ 
18+from __future__ import annotations
19+ 
20+import copy
21+from collections import defaultdict
22+from typing import Any, Dict, Iterable, List, Optional, Tuple
23+ 
24+import torch
25+ 
26+from hyper_parallel.core.optimizer.adamw import AdamW as HyperAdamW
27+from hyper_parallel.core.optimizer.adamw import adamw as hyper_adamw
28+from hyper_parallel.core.optimizer.swap_optimizer_base import (
29+ OptimizerSwapAdapter,
30+ SwapSlot,
31+ UpdateUnit,
32+)
33+ 
34+ 
35+class TorchAdamBaseAdapter(OptimizerSwapAdapter):
36+ """Common Torch Adam/AdamW adapter logic."""
37+ 
38+ functional_name = "adam"
39+ supported_cls = ()
40+ decoupled_weight_decay = False
41+ is_hyper_adamw = False
42+ supports_fused = False
43+ 
44+ @classmethod
45+ def matches(cls, optimizer: Any) -> bool:
46+ """Return whether this adapter supports ``optimizer``."""
47+ return isinstance(optimizer, cls.supported_cls)
48+ 
49+ def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
50+ super().__init__(optimizer, config, runtime)
51+ self._slots: Dict[Tuple[int, str], SwapSlot] = {}
52+ 
53+ def validate(self) -> None:
54+ """Validate unsupported optimizer flags."""
55+ for group in self.optimizer.param_groups:
56+ if group.get("foreach", False) is True:
57+ raise ValueError("Swap optimizer does not support foreach=True.")
58+ if group.get("fused", False) is True and not self.supports_fused:
59+ raise ValueError("Swap optimizer does not support fused=True.")
60+ if group.get("differentiable", False):
61+ raise ValueError("Swap optimizer does not support differentiable=True.")
62+ if group.get("capturable", False):
63+ raise ValueError("Swap optimizer does not support capturable=True.")
64+ 
65+ def prepare_step(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
66+ """Initialize lazy state and collect this step's update units."""
67+ if args or kwargs:
68+ raise ValueError("Torch swap optimizer step does not support closure or extra arguments.")
69+ if self.runtime.packed_enabled:
70+ return self._prepare_packed_step()
71+ 
72+ units = []
73+ for group_index, group in enumerate(self.optimizer.param_groups):
74+ if self.is_hyper_adamw:
75+ group["step"] = (group.get("step") or 0) + 1
76+ for param in group["params"]:
77+ grad = getattr(param, "grad", None)
78+ if grad is None:
79+ continue
80+ if getattr(grad, "is_sparse", False):
81+ raise ValueError("Swap optimizer only supports dense Adam/AdamW gradients.")
82+ state = self.optimizer.state[param]
83+ self._init_param_state(param, grad, group)
84+ slots = self._build_slots(param, state)
85+ units.append(UpdateUnit(
86+ adapter_index=group_index,
87+ param=param,
88+ grad=grad,
89+ slots=slots,
90+ ))
91+ return {"units": units}
92+ 
93+ def _prepare_packed_step(self) -> Dict[str, Any]:
94+ """Build a stable packed layout while retaining inactive materialized states."""
95+ records = []
96+ for group_index, group in enumerate(self.optimizer.param_groups):
97+ if self.is_hyper_adamw:
98+ group["step"] = (group.get("step") or 0) + 1
99+ for param in group["params"]:
100+ grad = getattr(param, "grad", None)
Ahyper_parallel/platform/torch/swap_optimizer/swap_optimizer.py+666-0
@@ -0,0 +1,666 @@
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+"""Torch optimizer state swap wrapper."""
16+# pylint: disable=protected-access
17+ 
18+from __future__ import annotations
19+ 
20+import contextlib
21+from dataclasses import dataclass, field
22+from typing import Any, Dict, Iterable, List, Optional, Sequence
23+ 
24+import torch
25+ 
26+from hyper_parallel.core.optimizer.swap_optimizer_base import PipelineSwapRuntime, SwapSlot
27+from hyper_parallel.platform import get_platform
28+from hyper_parallel.platform.torch.swap_optimizer.adapters import (
29+ TorchHyperAdamWAdapter,
30+ TorchNativeAdamAdapter,
31+ TorchNativeAdamWAdapter,
32+)
33+ 
34+ 
35+_PACKED_ALIGNMENT_BYTES = 512
36+platform = get_platform()
37+ 
38+ 
39+@dataclass
40+class _PackedBatchRegion:
41+ """One dtype-contiguous host range transferred for a pipeline batch."""
42+ 
43+ dtype: Any
44+ host_offset: int
45+ numel: int
46+ slots: List[SwapSlot]
47+ 
48+ 
49+@dataclass
50+class _PackedBatchPlan:
51+ """Packed transfer regions for one optimizer pipeline batch."""
52+ 
53+ regions: Dict[Any, _PackedBatchRegion] = field(default_factory=dict)
54+ 
55+ 
56+@dataclass
57+class _StagingArena:
58+ """One raw device allocation and its dtype-specific views."""
59+ 
60+ raw_buffer: Any
61+ dtype_views: Dict[Any, Any] = field(default_factory=dict)
62+ layout_signature: Any = None
63+ 
64+ 
65+class TorchSwapRuntime(PipelineSwapRuntime):
66+ """Torch tensor storage/copy runtime."""
67+ 
68+ def __init__(self, config: Any) -> None:
69+ super().__init__(config)
70+ self._packed_enabled = bool(getattr(config, "packed_swap", True))
71+ self._host_buffers: Dict[Any, Any] = {}
72+ self._host_layout_signature = ()
73+ self._packed_batch_plans: List[_PackedBatchPlan] = []
74+ self._staging_arenas: List[Optional[_StagingArena]] = [None, None]
75+ self._packed_ready_events: Dict[int, Any] = {}
76+ self._packed_offload_events: Dict[int, Any] = {}
77+ self._packed_tail_event: Optional[Any] = None
78+ self._packed_device_views: Dict[Any, Any] = {}
79+ 
80+ @property
81+ def packed_enabled(self) -> bool:
82+ """Return whether this runtime may build packed state candidates."""
83+ return self._packed_enabled
84+ 
85+ def populate_slot_metadata(self, slot: SwapSlot, template: Any) -> None:
86+ """Populate stable logical tensor metadata without allocating device state."""
87+ storage_tensor = self._storage_tensor(template)
88+ slot.shape = tuple(storage_tensor.shape)
89+ slot.dtype = storage_tensor.dtype
90+ slot.device = storage_tensor.device
91+ slot.numel = int(storage_tensor.numel())
92+ slot.storage_nbytes = slot.numel * int(storage_tensor.element_size())
93+ 
94+ def is_packable_template(self, tensor: Any, min_numel: int) -> bool:
95+ """Return whether a state shaped like ``tensor`` can use packed staging."""
96+ if not self._packed_enabled:
97+ return False
98+ return self.is_swappable_tensor(tensor, min_numel)
99+ 
100+ @staticmethod
Atests/ut/platform/mindspore/swap_optimizer/__init__.py+15-0
Atests/ut/platform/torch/swap_optimizer/__init__.py+15-0