Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / loss_parallel.py: 96%
47 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-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"""loss_parallel context manager.
17Context manager and utilities for distributed cross-entropy computation:
18 - loss_parallel(): Context manager to enable distributed CE semantics
19 - is_loss_parallel_active(): Debug/assert function
20 - Internal state management via ContextVar
21"""
23from __future__ import annotations
25import contextlib
26import contextvars
27import threading
28from typing import Optional, ContextManager, TYPE_CHECKING
30if TYPE_CHECKING:
31 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
33__all__ = [
34 "loss_parallel",
35 "is_loss_parallel_active",
36]
39# ========================================================================
40# ContextVar state management
41# ========================================================================
43_loss_parallel_active: contextvars.ContextVar[bool] = contextvars.ContextVar(
44 "_loss_parallel_active", default=False
45)
47_loss_parallel_count: contextvars.ContextVar[int] = contextvars.ContextVar(
48 "_loss_parallel_count", default=0
49)
51_loss_parallel_context_id: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar(
52 "_loss_parallel_context_id", default=None
53)
55_global_context_id_counter: int = 0
56_global_context_id_lock: threading.Lock = threading.Lock()
58_loss_parallel_mesh: contextvars.ContextVar[Optional["DeviceMesh"]] = contextvars.ContextVar(
59 "_loss_parallel_mesh", default=None
60)
62_loss_parallel_strict: contextvars.ContextVar[bool] = contextvars.ContextVar(
63 "_loss_parallel_strict", default=True
64)
67# ========================================================================
68# loss_parallel context manager
69# ========================================================================
71@contextlib.contextmanager
72def loss_parallel(
73 *,
74 mesh: Optional["DeviceMesh"] = None,
75 strict: bool = True,
76) -> ContextManager[None]:
77 """Context manager to enable distributed CE semantics.
79 Args:
80 mesh: Explicitly specify TP sub-mesh; None infers from participating DTensors;
81 inference failure with strict=True raises error.
82 strict: When True, layout contract violation raises ValueError or dedicated exception;
83 False only warns or takes optional fallback path (if fallback_gather implemented).
85 Nesting semantics:
86 Reentrant counting; count must return to zero on exit. Supports nested usage.
88 Example:
89 >>> with loss_parallel():
90 ... logits = model(input) # logits is DTensor, Shard(-1)
91 ... loss = F.cross_entropy(logits, target)
92 ... loss.backward()
94 Note:
95 PyTorch upstream loss_parallel() is often parameterless.
96 HyperParallel adds parameters; see docs/compatibility/pytorch_diff.md.
97 """
98 global _global_context_id_counter
100 old_count = _loss_parallel_count.get()
101 old_active = _loss_parallel_active.get()
102 old_mesh = _loss_parallel_mesh.get()
103 old_strict = _loss_parallel_strict.get()
104 old_context_id = _loss_parallel_context_id.get()
106 _loss_parallel_count.set(old_count + 1)
107 _loss_parallel_active.set(True)
108 _loss_parallel_mesh.set(mesh)
109 _loss_parallel_strict.set(strict)
111 if old_count == 0:
112 with _global_context_id_lock:
113 _global_context_id_counter += 1
114 current_context_id = _global_context_id_counter
115 _loss_parallel_context_id.set(current_context_id)
116 else:
117 current_context_id = old_context_id
119 try:
120 yield
121 finally:
122 _loss_parallel_count.set(old_count)
123 _loss_parallel_active.set(old_active)
124 _loss_parallel_mesh.set(old_mesh)
125 _loss_parallel_strict.set(old_strict)
126 _loss_parallel_context_id.set(old_context_id)
129# ========================================================================
130# is_loss_parallel_active (optional debug API)
131# ========================================================================
133def is_loss_parallel_active() -> bool:
134 """Debug or assert: whether current thread is in loss_parallel context.
136 Returns:
137 True: currently in loss_parallel context
138 False: not in context
139 """
140 return _loss_parallel_active.get()
143def get_loss_parallel_count() -> int:
144 """Get current nesting count.
146 Returns:
147 int: Nesting depth (0 means not in context).
148 """
149 return _loss_parallel_count.get()
152# ========================================================================
153# Internal functions (for use by other modules)
154# ========================================================================
156def _get_loss_parallel_token() -> Optional[int]:
157 """Get current context token for cache key differentiation.
159 Returns:
160 Optional[int]: Current context unique ID for differentiating cache keys.
161 Returns None if not in context, returns positive integer for context ID.
162 """
163 return _loss_parallel_context_id.get()
166def _get_loss_parallel_mesh() -> Optional["DeviceMesh"]:
167 """Get the mesh set in current context.
169 Returns:
170 Optional[DeviceMesh]: Explicitly specified mesh, or None (infer from DTensor).
171 """
172 return _loss_parallel_mesh.get()
175def _get_loss_parallel_strict() -> bool:
176 """Get the strict setting in current context.
178 Returns:
179 bool: Whether in strict mode.
180 """
181 return _loss_parallel_strict.get()