Skip to content

Commit 9fd23fe

Browse files
apstenku123apstenku123
committed
feat(v4/parallelism): PSpec Stage A — DeviceTopology + ShardingSpec data-layer
Adds the planning/sizing foundation for distributed training (ParallelismSpec.md §3.1-3.2, §6 A). Pure data layer — Apple-Silicon MLX target has no runtime FSDP/Megatron; this layer models the math so the visual builder can predict OOM on CUDA/TPU configs before they get launched on nanochat / cppmega. cppmega_v4/parallelism/topology.py: - DeviceKind enum (H100_80GB, H200_141GB, A100_40/80GB, B100_80GB, GB10, TPU_V5P, TPU_V6E, M3_ULTRA) - DeviceSpec frozen dataclass with __post_init__ validation (HBM > 0, interconnect ∈ 7-entry whitelist, bandwidth > 0) - DeviceTopology frozen dataclass: devices + mesh_axes where the product of axis degrees must equal device count - Built-in factories: h100_8x / h200_8x / a100_8x (40|80gb) / b100_8x / gb10_quarter / tpu_v6e_8 (2D dp×tp) / tpu_v5p_4 / m3_ultra_solo. Each accepts optional mesh-axis kwargs to support splits (h100_8x(dp=2, tp=2, ep=2) for 3D parallelism). - TOPOLOGY_BUILTINS registry maps name → dotted factory path for the GUI dropdown (Stage E). cppmega_v4/parallelism/sharding_spec.py: - ParallelismKind enum (DP, FSDP1, FSDP2, ZERO1, ZERO2, TP, SP, EP, PP, PP_VPP) - AxisAssignment: axis_name + kind + degree (with validation) - ShardingSpec: topology + axis_assignments + per-brick override + knobs (master_weights_fp32, grad_reduce_dtype, compile_mode, fp8_enabled, activation_checkpointing). Validation: every axis_name must exist in topology mesh; axis degree must match mesh degree; no duplicate axes; enum-bounded grad_dtype / compile_mode / checkpoint mode. Importantly accepts ``compile_mode="whole_model"`` (KNOWN BROKEN with FSDP2 and Megatron) so the Stage C gotcha checker can flag it loudly — refusing construction would hide the footgun. - Convenience: .axis_kinds() (frozenset), .degree_of(kind) aggregate. - Built-in factories: single_device / fsdp2_only / megatron_ep_only / fsdp2_plus_tp — mirror the production patterns from cppmega (EP=4/8) and nanochat (FSDP2 + Megatron TP regional-compiled). Tests (tests/v4/test_pspec_stage_a.py, 57 cases): - DeviceSpec rejection: non-enum kind / zero HBM / unknown interconnect / non-positive bandwidth - DeviceTopology rejection: empty devices / non-DeviceSpec entries / empty mesh / bad axis degree / product mismatch - DeviceTopology helpers: num_devices, total_hbm_bytes, axis_degree - Built-in topology factories: each well-formed; parametrized "every TOPOLOGY_BUILTINS entry instantiates" - AxisAssignment rejection (kind/blank-name/degree) - ShardingSpec rejection: non-topology / empty assignments / duplicate axis / axis-not-in-mesh / degree-mismatch / bad grad_dtype / bad compile_mode / bad checkpoint mode - ShardingSpec accepts ``"none"`` axis for pure replication - ShardingSpec accepts ``compile_mode="whole_model"`` (for gotcha detection downstream) - ShardingSpec helpers: axis_kinds, degree_of multi-axis aggregate, num_ranks - Built-in sharding factories: single_device, fsdp2_only, megatron_ep_only (with missing-axis rejection), fsdp2_plus_tp - Frozen invariant across DeviceSpec / DeviceTopology / ShardingSpec Full v4 regression: 972 passed / 5 skipped / 0 failed.
1 parent 52dd692 commit 9fd23fe

4 files changed

Lines changed: 1072 additions & 0 deletions

File tree

‎cppmega_v4/parallelism/__init__.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
"""ParallelismSpec — planning/sizing layer for FSDP/TP/EP/PP.
2+
3+
See ``ParallelismSpec.md`` (repo root) for the design + lessons imported
4+
from ../cppmega (Megatron production) and ../nanochat (multi-stack:
5+
FSDP2 + Megatron TP + SPMD).
6+
7+
Stage A surface (this commit):
8+
- topology: DeviceKind, DeviceSpec, DeviceTopology + builtin factories
9+
- sharding_spec: ParallelismKind, AxisAssignment, ShardingSpec +
10+
builtin factories (single_device, fsdp2_only, megatron_ep_only,
11+
fsdp2_plus_tp)
12+
"""
13+
14+
from __future__ import annotations
15+
16+
from cppmega_v4.parallelism.sharding_spec import (
17+
AxisAssignment,
18+
ParallelismKind,
19+
ShardingSpec,
20+
fsdp2_only,
21+
fsdp2_plus_tp,
22+
megatron_ep_only,
23+
single_device,
24+
)
25+
from cppmega_v4.parallelism.topology import (
26+
TOPOLOGY_BUILTINS,
27+
DeviceKind,
28+
DeviceSpec,
29+
DeviceTopology,
30+
a100_8x,
31+
b100_8x,
32+
gb10_quarter,
33+
h100_8x,
34+
h200_8x,
35+
m3_ultra_solo,
36+
tpu_v5p_4,
37+
tpu_v6e_8,
38+
)
39+
40+
__all__ = [
41+
"AxisAssignment",
42+
"DeviceKind",
43+
"DeviceSpec",
44+
"DeviceTopology",
45+
"ParallelismKind",
46+
"ShardingSpec",
47+
"TOPOLOGY_BUILTINS",
48+
"a100_8x",
49+
"b100_8x",
50+
"fsdp2_only",
51+
"fsdp2_plus_tp",
52+
"gb10_quarter",
53+
"h100_8x",
54+
"h200_8x",
55+
"m3_ultra_solo",
56+
"megatron_ep_only",
57+
"single_device",
58+
"tpu_v5p_4",
59+
"tpu_v6e_8",
60+
]

‎cppmega_v4/parallelism/sharding_spec.py‎

Lines changed: 291 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,291 @@
1+
"""ShardingSpec — declarative parallelism strategy over a DeviceTopology.
2+
3+
A :class:`ShardingSpec` carries:
4+
- the topology (which devices, what mesh)
5+
- a tuple of :class:`AxisAssignment` (which mesh axis runs which
6+
ParallelismKind at what degree)
7+
- global knobs that affect memory accounting and known footguns:
8+
master_weights_fp32, grad_reduce_dtype, compile_mode, fp8_enabled,
9+
activation_checkpointing
10+
11+
Pure data layer. The memory estimator (Stage B) and gotcha checker
12+
(Stage C) read these without modifying them.
13+
"""
14+
15+
from __future__ import annotations
16+
17+
from collections.abc import Mapping
18+
from dataclasses import dataclass, field
19+
from enum import Enum
20+
from typing import Final
21+
22+
from cppmega_v4.parallelism.topology import DeviceTopology
23+
24+
25+
class ParallelismKind(str, Enum):
26+
"""Recognised parallelism strategies (3D + ZeRO variants)."""
27+
28+
DP = "dp" # replicate weights, shard data
29+
FSDP1 = "fsdp1" # legacy fully_sharded_data_parallel
30+
FSDP2 = "fsdp2" # modern fully_shard (ZeRO-3)
31+
ZERO1 = "zero1" # optimizer state only
32+
ZERO2 = "zero2" # optim state + gradients
33+
TP = "tp" # tensor parallel intra-layer
34+
SP = "sp" # sequence parallel (with TP)
35+
EP = "ep" # expert parallel (MoE)
36+
PP = "pp" # pipeline parallel
37+
PP_VPP = "pp_vpp" # virtual pipeline parallel
38+
39+
40+
_VALID_GRAD_DTYPES: Final[frozenset[str]] = frozenset({"bf16", "fp32"})
41+
_VALID_COMPILE_MODES: Final[frozenset[str]] = frozenset({
42+
"off", "regional", "whole_model",
43+
})
44+
_VALID_ACTIVATION_CHECKPOINT: Final[frozenset[str]] = frozenset({
45+
"off", "full", "selective",
46+
})
47+
48+
49+
@dataclass(frozen=True)
50+
class AxisAssignment:
51+
"""One row in the sharding table: this mesh axis runs this kind.
52+
53+
Fields:
54+
axis_name: name of the mesh axis (must exist in topology.mesh_axes
55+
or be ``"none"`` for plain replication).
56+
kind: which parallelism family runs on this axis.
57+
degree: degree of parallelism. Must equal the mesh-axis degree
58+
when ``axis_name`` is a real mesh axis.
59+
"""
60+
61+
axis_name: str
62+
kind: ParallelismKind
63+
degree: int
64+
65+
def __post_init__(self) -> None:
66+
if not isinstance(self.kind, ParallelismKind):
67+
raise TypeError(
68+
f"AxisAssignment.kind must be ParallelismKind, got "
69+
f"{type(self.kind).__name__}"
70+
)
71+
if not isinstance(self.axis_name, str) or not self.axis_name.strip():
72+
raise ValueError(
73+
"AxisAssignment.axis_name must be non-empty str"
74+
)
75+
if not isinstance(self.degree, int) or self.degree < 1:
76+
raise ValueError(
77+
f"AxisAssignment.degree must be int ≥ 1, got {self.degree!r}"
78+
)
79+
80+
81+
@dataclass(frozen=True)
82+
class ShardingSpec:
83+
"""Declarative parallelism strategy.
84+
85+
Fields:
86+
topology: the underlying device mesh.
87+
axis_assignments: which mesh axis runs which parallelism kind.
88+
Every entry's ``axis_name`` must exist in topology.mesh_axes
89+
and its ``degree`` must match the mesh axis degree.
90+
brick_axes: optional per-brick override — when present, only these
91+
bricks participate in the named axes. Empty dict = global.
92+
master_weights_fp32: keep an fp32 master-weight copy alongside
93+
bf16/fp8 compute weights. Doubles param memory; required for
94+
certain optimisers / loss scaling regimes (see nanochat
95+
``megatron_optimizer.py`` for the duplication pain story).
96+
grad_reduce_dtype: ``"bf16"`` (default; faster) | ``"fp32"`` (more
97+
stable for very deep stacks; doubles grad-buffer cost).
98+
compile_mode: ``"off"`` | ``"regional"`` (per-block; required to
99+
avoid FSDP2/Megatron compile footguns — see GotchaChecker
100+
Stage C) | ``"whole_model"`` (KNOWN BROKEN with FSDP2 and Megatron;
101+
the spec accepts it so we can flag it loudly).
102+
fp8_enabled: forward in FP8 via Transformer Engine / torchao; this
103+
does NOT make grads or optimiser FP8 — known duplication pain.
104+
activation_checkpointing: ``"off"`` (peak activations) | ``"full"``
105+
(only block boundaries kept) | ``"selective"`` (per-layer
106+
cherry-pick).
107+
"""
108+
109+
topology: DeviceTopology
110+
axis_assignments: tuple[AxisAssignment, ...]
111+
brick_axes: Mapping[str, tuple[str, ...]] = field(default_factory=dict)
112+
master_weights_fp32: bool = False
113+
grad_reduce_dtype: str = "bf16"
114+
compile_mode: str = "regional"
115+
fp8_enabled: bool = False
116+
activation_checkpointing: str = "full"
117+
118+
def __post_init__(self) -> None:
119+
if not isinstance(self.topology, DeviceTopology):
120+
raise TypeError(
121+
f"ShardingSpec.topology must be DeviceTopology, got "
122+
f"{type(self.topology).__name__}"
123+
)
124+
if not self.axis_assignments:
125+
raise ValueError(
126+
"ShardingSpec.axis_assignments must declare at least one entry"
127+
)
128+
seen_axes: set[str] = set()
129+
for a in self.axis_assignments:
130+
if not isinstance(a, AxisAssignment):
131+
raise TypeError(
132+
f"ShardingSpec.axis_assignments entries must be "
133+
f"AxisAssignment, got {type(a).__name__}"
134+
)
135+
if a.axis_name in seen_axes:
136+
raise ValueError(
137+
f"ShardingSpec axis {a.axis_name!r} appears more than once"
138+
)
139+
seen_axes.add(a.axis_name)
140+
if a.axis_name == "none":
141+
continue
142+
if a.axis_name not in self.topology.mesh_axes:
143+
raise ValueError(
144+
f"ShardingSpec axis {a.axis_name!r} not in topology mesh "
145+
f"{sorted(self.topology.mesh_axes)}"
146+
)
147+
if a.degree != self.topology.mesh_axes[a.axis_name]:
148+
raise ValueError(
149+
f"ShardingSpec axis {a.axis_name!r} degree ({a.degree}) "
150+
f"must match topology mesh degree "
151+
f"({self.topology.mesh_axes[a.axis_name]})"
152+
)
153+
if self.grad_reduce_dtype not in _VALID_GRAD_DTYPES:
154+
raise ValueError(
155+
f"ShardingSpec.grad_reduce_dtype={self.grad_reduce_dtype!r} "
156+
f"not in {sorted(_VALID_GRAD_DTYPES)}"
157+
)
158+
if self.compile_mode not in _VALID_COMPILE_MODES:
159+
raise ValueError(
160+
f"ShardingSpec.compile_mode={self.compile_mode!r} not in "
161+
f"{sorted(_VALID_COMPILE_MODES)}"
162+
)
163+
if self.activation_checkpointing not in _VALID_ACTIVATION_CHECKPOINT:
164+
raise ValueError(
165+
f"ShardingSpec.activation_checkpointing="
166+
f"{self.activation_checkpointing!r} not in "
167+
f"{sorted(_VALID_ACTIVATION_CHECKPOINT)}"
168+
)
169+
170+
def axis_kinds(self) -> frozenset[ParallelismKind]:
171+
return frozenset(a.kind for a in self.axis_assignments)
172+
173+
def degree_of(self, kind: ParallelismKind) -> int:
174+
"""Return aggregate degree across all axes carrying ``kind``."""
175+
d = 1
176+
for a in self.axis_assignments:
177+
if a.kind is kind:
178+
d *= a.degree
179+
return d
180+
181+
@property
182+
def num_ranks(self) -> int:
183+
"""Total number of ranks across the mesh."""
184+
return self.topology.num_devices
185+
186+
187+
# ---------------------------------------------------------------------------
188+
# Built-in factories (the patterns from cppmega + nanochat production)
189+
# ---------------------------------------------------------------------------
190+
191+
192+
def single_device(topology: DeviceTopology) -> ShardingSpec:
193+
"""Replicate everything on every device (DP-only)."""
194+
return ShardingSpec(
195+
topology=topology,
196+
axis_assignments=(
197+
AxisAssignment(
198+
axis_name=next(iter(topology.mesh_axes)),
199+
kind=ParallelismKind.DP,
200+
degree=next(iter(topology.mesh_axes.values())),
201+
),
202+
),
203+
compile_mode="regional",
204+
)
205+
206+
207+
def fsdp2_only(
208+
topology: DeviceTopology,
209+
*,
210+
fp8_enabled: bool = False,
211+
activation_checkpointing: str = "full",
212+
) -> ShardingSpec:
213+
"""FSDP2 across the dp axis only (Zaero-3, no TP/EP)."""
214+
dp_axis = "dp" if "dp" in topology.mesh_axes else next(iter(topology.mesh_axes))
215+
return ShardingSpec(
216+
topology=topology,
217+
axis_assignments=(
218+
AxisAssignment(
219+
axis_name=dp_axis, kind=ParallelismKind.FSDP2,
220+
degree=topology.mesh_axes[dp_axis],
221+
),
222+
),
223+
fp8_enabled=fp8_enabled,
224+
activation_checkpointing=activation_checkpointing,
225+
compile_mode="regional",
226+
)
227+
228+
229+
def megatron_ep_only(
230+
topology: DeviceTopology,
231+
*,
232+
ep_axis: str = "ep",
233+
fp8_enabled: bool = False,
234+
) -> ShardingSpec:
235+
"""Megatron-style EP=N production pattern (cppmega EP=4/8)."""
236+
if ep_axis not in topology.mesh_axes:
237+
raise ValueError(
238+
f"megatron_ep_only: topology missing required mesh axis "
239+
f"{ep_axis!r}; have {sorted(topology.mesh_axes)}"
240+
)
241+
return ShardingSpec(
242+
topology=topology,
243+
axis_assignments=(
244+
AxisAssignment(
245+
axis_name=ep_axis, kind=ParallelismKind.EP,
246+
degree=topology.mesh_axes[ep_axis],
247+
),
248+
),
249+
fp8_enabled=fp8_enabled,
250+
compile_mode="regional",
251+
)
252+
253+
254+
def fsdp2_plus_tp(
255+
topology: DeviceTopology,
256+
*,
257+
dp_axis: str = "dp",
258+
tp_axis: str = "tp",
259+
) -> ShardingSpec:
260+
"""FSDP2 across dp + Megatron TP across tp (3D parallel base)."""
261+
for ax in (dp_axis, tp_axis):
262+
if ax not in topology.mesh_axes:
263+
raise ValueError(
264+
f"fsdp2_plus_tp: topology missing required mesh axis "
265+
f"{ax!r}; have {sorted(topology.mesh_axes)}"
266+
)
267+
return ShardingSpec(
268+
topology=topology,
269+
axis_assignments=(
270+
AxisAssignment(
271+
axis_name=dp_axis, kind=ParallelismKind.FSDP2,
272+
degree=topology.mesh_axes[dp_axis],
273+
),
274+
AxisAssignment(
275+
axis_name=tp_axis, kind=ParallelismKind.TP,
276+
degree=topology.mesh_axes[tp_axis],
277+
),
278+
),
279+
compile_mode="regional", # mandatory — see GotchaChecker
280+
)
281+
282+
283+
__all__ = [
284+
"AxisAssignment",
285+
"ParallelismKind",
286+
"ShardingSpec",
287+
"fsdp2_only",
288+
"fsdp2_plus_tp",
289+
"megatron_ep_only",
290+
"single_device",
291+
]

0 commit comments

Comments
 (0)