Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/auto_parallel/_op

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Op records: what each op of a layer keeps for its backward, as data (shared decision S1).

17``op_records.yaml`` states, for every op of the op vector, the recompute

18switch that drops it, if any, the ops whose outputs recomputing it takes,

19and whether it is communication; the *parts* of an op's saves that a switch

20of their own drops, each named for the op it is part of; and, for each of the

21memory model's activation formulas, a *slot*, the formula and the elements

22each of its ops and parts keeps. The memory model prices a layer's activations from the slots, so

23the records are the one source of what a layer keeps, op by op: a caller

24prices any setting of the switches from them without setting the switches

25on a config.

27The expressions are Python arithmetic, parsed once and refused unless every

28node is one of a closed set: numbers, names, the arithmetic and comparison

29operators, ``and``, ``or``, ``not``, conditional expressions, and calls of

30``max`` and ``min``. A name is a field of the layer config, or one of the

31names the caller binds.

32"""

33import ast

34import functools

35import os

36from dataclasses import dataclass, field, fields

37from types import MappingProxyType

38from typing import Any, Callable, Dict, FrozenSet, Mapping, Optional, Tuple

40import yaml

42from hyper_parallel.auto_parallel._model_spec import ModelSpecError, OpCounts

44RECORDS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "op_records.yaml")

46# The parts of a layer a slot belongs to, as the memory logs name them.

47PARTS = ("attention", "ffn", "norm")

49# The recompute switches: a selective layer drops what the ops they name keep.

50# attUp drops the query, key and value heads an MLA layer's up-projections

51# build from the latents it keeps, a part of attMM's saves.

52SWITCHES = ("attBMM", "headCast", "dropout", "softmax", "normOp", "gather", "ffAct", "attUp")

54_NODES = (

55 ast.Expression, ast.BinOp, ast.UnaryOp, ast.BoolOp, ast.Compare, ast.IfExp, ast.Call, ast.Name,

56 ast.Constant, ast.Load, ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, ast.Mod, ast.Pow,

57 ast.USub, ast.UAdd, ast.Not, ast.And, ast.Or, ast.Eq, ast.NotEq, ast.Lt, ast.LtE, ast.Gt, ast.GtE,

58)

59_CALLS = MappingProxyType({"max": max, "min": min})

60# What an expression sees beside its names: max and min, and no builtins.

61_GLOBALS = {"__builtins__": {}, **_CALLS}

64class _Names(dict):

65 """The names an expression reads, each looked up the first time it is read.

67 A conditional expression reads only its branch's names, as the code it

68 stands for read only its branch's fields.

69 """

71 def __init__(self, value: Callable[[str], Any]) -> None:

72 """Look each name up with *value*."""

73 super().__init__()

74 self._value = value

76 def __missing__(self, name: str) -> Any:

77 if name in _CALLS:

78 # Not a name to look up: the call the expression makes.

79 raise KeyError(name)

80 found = self._value(name)

81 self[name] = found

82 return found

85@dataclass(frozen=True)

86class Expression:

87 """One record expression, checked and compiled.

89 Attributes:

90 source: The expression as written.

91 names: The names it reads, the calls it makes apart.

92 """

94 source: str

95 names: FrozenSet[str]

96 code: Any

98 @classmethod

99 def parse(cls, source: Any, where: str) -> "Expression":

100 """Parse *source*, refusing a node outside the closed set or a call of anything but max and min."""

101 if not isinstance(source, (str, int, float)) or isinstance(source, bool):

102 raise ModelSpecError(f"{where} must be an expression, got {source!r}")

103 text = str(source).strip()

104 try:

105 tree = ast.parse(text, mode="eval")

106 except SyntaxError as exc:

107 raise ModelSpecError(f"{where}: {text!r} is not an expression: {exc.msg}") from exc

108 names = set()

109 for node in ast.walk(tree):

110 if not isinstance(node, _NODES):

111 raise ModelSpecError(f"{where}: {text!r} uses {type(node).__name__}, which a record may not")

112 if isinstance(node, ast.Call):

113 if not isinstance(node.func, ast.Name) or node.func.id not in _CALLS or node.keywords:

114 raise ModelSpecError(f"{where}: {text!r} calls something other than max or min")

115 elif isinstance(node, ast.Name):

116 names.add(node.id)

117 elif isinstance(node, ast.Constant) and not isinstance(node.value, (int, float)):

118 raise ModelSpecError(f"{where}: {text!r} holds {node.value!r}, which is not a number")

119 return cls(text, frozenset(names - set(_CALLS)), compile(tree, where, "eval"))

120

121 def __call__(self, names: Dict[str, Any]) -> Any:

122 """Evaluate the expression with *names* bound."""

123 return eval(self.code, _GLOBALS, names) # pylint: disable=eval-used

124

125

126@dataclass(frozen=True)

127class OpRecord:

128 """What the records state of one op, beside what it keeps in each slot.

129

130 Attributes:

131 name: The op, one of the op vector's, or a part of one's saves.

132 switch: The recompute switch that drops what it keeps, or ``None``

133 for an op every selective layer keeps.

134 needs: The ops whose outputs recomputing it takes.

135 comm: Whether it is communication rather than compute.

136 part_of: For a part, the op of the op vector whose saves it is part

137 of; ``None`` for an op.

138 """

139

140 name: str

141 switch: Optional[str]

142 needs: Tuple[str, ...]

143 comm: bool

144 part_of: Optional[str] = None

145

146

147@dataclass(frozen=True)

148class Slot:

149 """One of the memory model's activation formulas, over the elements its ops keep.

150

151 Attributes:

152 name: The slot.

153 part: The part of the layer its bytes count toward.

154 formula: The bytes a layer keeps in it, over ``total``.

155 ops: Each op's elements in it, in the order the formula sums them.

156 reads: Every name its expressions read but ``total``.

157 when: Of which layers it is part, or ``None`` for every layer.

158 """

159

160 name: str

161 part: str

162 formula: Expression

163 ops: Tuple[Tuple[str, Expression], ...]

164 reads: FrozenSet[str]

165 when: Optional[Expression] = None

166

167 def holds(self, value: Callable[[str], Any]) -> bool:

168 """Whether the slot is part of a layer whose names *value* gives."""

169 return self.when is None or bool(self.when(_Names(value)))

170

171

172@dataclass(frozen=True)

173class OpRecords:

174 """Every op's record and every slot.

175

176 Attributes:

177 ops: The record of each op of the op vector, by name.

178 slots: The slots, by name.

179 defaults: The value of a field a layer config may not hold.

180 parts: The record of each part of an op's saves, by name.

181 """

182

183 ops: Mapping[str, OpRecord]

184 slots: Mapping[str, Slot]

185 defaults: Mapping[str, float]

186 parts: Mapping[str, OpRecord] = field(default_factory=lambda: MappingProxyType({}))

187

188 def record(self, name: str) -> OpRecord:

189 """The record of an op or of a part."""

190 return self.ops[name] if name in self.ops else self.parts[name]

191

192 def _counts_toward(self, name: str, op: str) -> bool:

193 """Whether what *name* keeps counts toward *op*: it is the op, or a part of its saves."""

194 return name == op or (name in self.parts and self.parts[name].part_of == op)

195

196 def keeps(self, op: str) -> Dict[str, Expression]:

197 """The slots *op* keeps anything in, its parts' among them, each with the elements it keeps there.

198

199 For an op of the op vector, the elements of the first of its

200 entries in each slot, itself or a part; for a part, its own.

201 """

202 if op not in self.ops and op not in self.parts:

203 raise ModelSpecError(f"no op {op!r}; the ops are {sorted(self.ops)}, the parts {sorted(self.parts)}")

204 found: Dict[str, Expression] = {}

205 for slot in self.slots.values():

206 for name, elements in slot.ops:

207 if self._counts_toward(name, op):

208 found.setdefault(slot.name, elements)

209 return found

210

211 def evaluate(self, slot: str, value: Callable[[str], Any], keep: Callable[[str], Any],

212 only: Optional[str] = None) -> float:

213 """The bytes a layer keeps in *slot*.

214

215 Args:

216 slot: The slot.

217 value: The value of each name the slot reads.

218 keep: For each op and part, 1 where the layer keeps what it

219 keeps, 0 where it drops it.

220 only: An op to price alone: the slot's formula over its

221 elements and its parts', whatever *keep* says; or a part,

222 over its own.

223

224 Returns:

225 The bytes, one rank's share of a micro-batch.

226 """

227 record = self.slots[slot]

228 names = _Names(value)

229 total = 0

230 for op, elements in record.ops:

231 kept = int(self._counts_toward(op, only)) if only is not None else keep(op)

232 if kept:

233 total += kept * elements(names)

234 names["total"] = total

235 return record.formula(names)

236

237

238def _op_names() -> Tuple[str, ...]:

239 """The ops of the op vector, in its order."""

240 return tuple(field.name for field in fields(OpCounts))

241

242

243def _op_record(name: str, data: Any, where: str = "", part: bool = False) -> OpRecord:

244 """Build one op's record from its ``{switch, needs, comm}`` mapping, a part's with its ``part_of``."""

245 where = where or f"op_records.ops.{name}"

246 keys = {"switch", "needs", "comm"} | ({"part_of"} if part else set())

247 if not isinstance(data, Mapping):

248 raise ModelSpecError(f"{where} must map {', '.join(sorted(keys))}, got {data!r}")

249 unknown = sorted(set(data) - keys)

250 if unknown:

251 raise ModelSpecError(f"{where} has unknown keys {unknown}; a record has {', '.join(sorted(keys))}")

252 switch = data.get("switch")

253 if switch is not None and switch not in SWITCHES:

254 raise ModelSpecError(f"{where}.switch must be one of {list(SWITCHES)}, got {switch!r}")

255 needs = data.get("needs", [])

256 ops = _op_names()

257 if not isinstance(needs, list) or any(need not in ops for need in needs):

258 raise ModelSpecError(f"{where}.needs must list ops of {ops}, got {needs!r}")

259 comm = data.get("comm", False)

260 if not isinstance(comm, bool):

261 raise ModelSpecError(f"{where}.comm must be true or false, got {comm!r}")

262 part_of = data.get("part_of") if part else None

263 if part and (part_of not in ops or switch is None):

264 raise ModelSpecError(f"{where} must name the op of {ops} it is part of, and the switch that drops it")

265 return OpRecord(str(name), switch, tuple(needs), comm, part_of)

266

267

268def _parts(data: Any) -> Mapping[str, OpRecord]:

269 """The records of the parts of ops' saves that a switch of their own drops."""

270 parts = data or {}

271 if not isinstance(parts, Mapping):

272 raise ModelSpecError(f"op records' parts must map names to records, got {parts!r}")

273 clash = sorted(set(parts) & set(_op_names()))

274 if clash:

275 raise ModelSpecError(f"op records' parts {clash} are named as ops")

276 return MappingProxyType({str(name): _op_record(name, record, f"op_records.parts.{name}", part=True)

277 for name, record in parts.items()})

278

279

280def _slot(name: str, data: Any, ops: Mapping[str, OpRecord]) -> Slot: # ops: the ops and the parts

281 """Build one slot from its ``{part, formula, ops}`` mapping."""

282 where = f"op_records.slots.{name}"

283 if not isinstance(data, Mapping):

284 raise ModelSpecError(f"{where} must map part, formula and ops, got {data!r}")

285 unknown = sorted(set(data) - {"part", "formula", "ops", "when"})

286 if unknown:

287 raise ModelSpecError(f"{where} has unknown keys {unknown}; a slot has part, formula, ops and when")

288 part = data.get("part")

289 if part not in PARTS:

290 raise ModelSpecError(f"{where}.part must be one of {list(PARTS)}, got {part!r}")

291 formula = Expression.parse(data.get("formula"), f"{where}.formula")

292 if "total" not in formula.names:

293 raise ModelSpecError(f"{where}.formula must read total, what the slot's ops keep")

294 members = data.get("ops")

295 if not isinstance(members, Mapping) or not members:

296 raise ModelSpecError(f"{where}.ops must map ops to the elements they keep")

297 unknown_ops = sorted(set(members) - set(ops))

298 if unknown_ops:

299 raise ModelSpecError(f"{where}.ops names {unknown_ops}, which have no record")

300 members = tuple((str(op), Expression.parse(elements, f"{where}.ops.{op}")) for op, elements in members.items())

301 reads = set(formula.names)

302 for _, elements in members:

303 reads |= elements.names

304 when = None if data.get("when") is None else Expression.parse(data["when"], f"{where}.when")

305 return Slot(str(name), part, formula, members, frozenset(reads - {"total"}), when)

306

307

308def _defaults(data: Any) -> Mapping[str, float]:

309 """The values the records give a field a layer config may not hold."""

310 defaults = data or {}

311 if not isinstance(defaults, Mapping) or any(

312 isinstance(value, bool) or not isinstance(value, (int, float)) for value in defaults.values()):

313 raise ModelSpecError(f"op records' defaults must map fields to numbers, got {defaults!r}")

314 return MappingProxyType(dict(defaults))

315

316

317@functools.lru_cache(maxsize=None)

318def load_op_records(path: str = RECORDS_PATH) -> OpRecords:

319 """Read and check the op records.

320

321 Raises:

322 ModelSpecError: If the file misses an op of the op vector, holds an

323 unknown key, op or switch, or an expression outside the closed

324 set of nodes.

325 """

326 with open(path, encoding="utf-8") as handle:

327 data = yaml.safe_load(handle) or {}

328 unknown = sorted(set(data) - {"defaults", "ops", "parts", "slots"})

329 if unknown:

330 raise ModelSpecError(f"op records have unknown keys {unknown}")

331 stated = data.get("ops") or {}

332 names = _op_names()

333 missing = [name for name in names if name not in stated]

334 extra = sorted(set(stated) - set(names))

335 if missing or extra:

336 raise ModelSpecError(f"op records must state every op of {list(names)}: missing {missing}, unknown {extra}")

337 ops = {name: _op_record(name, stated[name]) for name in names}

338 parts = _parts(data.get("parts"))

339 slots = {str(name): _slot(name, slot, {**ops, **parts}) for name, slot in (data.get("slots") or {}).items()}

340 return OpRecords(MappingProxyType(ops), MappingProxyType(slots), _defaults(data.get("defaults")), parts)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _op_records.py: 93%

177 statements