Will Lin
1fe50c0092
The v2 top level had grown to ~27 dirs + 13 loose files — half of them
single-concept abstraction shells (memory/ 86 LOC, transport/ 222, parity/ 178,
extend/ 226) and half vendored fastvideo code sitting as peers to the actual v2
design. Regroup to mirror README section 3 "Planes & dependency order":
core/ enums+types, card, loop, program, parity, request, parallel
(the model-native contracts; no kernels)
runtime/ + folded-in substrate: cache, memory, transport, extend
(the import graph shows only runtime consumes them; compile/cudagraph
already lived here)
serving/ + deploy/ (products / fleet)
_vendor/ all copied fastvideo: models, layers, attention, configs, distributed,
platforms, api, hooks, logging_utils, third_party + fastvideo_args/
utils/logger/envs/forward_context — internal layout unchanged, still
mirrors upstream for diffing
28 dirs + 13 root files -> 8 dirs + 6 root files. Pure mechanical move: 947
absolute-import paths rewritten (v2.X -> v2.{core,runtime,serving,_vendor}.X),
boundary-anchored so platform/ (native dispatch) and platforms/ (vendored CUDA
detect) no longer collide and neither does hooks/ vs extend/. Deleted the empty
v2/loader/. README section 3 + 16 updated; stale test count corrected
(34 files/216 tests -> 22 files/143 tests).
Validated: `import v2` stays torch-free; `v2/run_tests.py` and `pytest v2/tests/`
-> 143 passed, 2 skipped on the numpy toy backend. (On a GPU box force the toy
backend with CUDA_VISIBLE_DEVICES="" or detect() picks cuda.) No external importer
changed — examples use the re-exported `from v2 import VideoGenerator`.
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
"""FakeDeviceMesh — a torchtitan-style ParallelDims builder, CPU-testable.
|
|
|
|
On a GPU box this compiles to ``torch.distributed.device_mesh.DeviceMesh``. Here it is a
|
|
pure-Python mesh of rank tuples so topology logic is unit-tested without GPUs.
|
|
Degree-one axes exist as trivial groups so component code needs no special cases.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from itertools import product
|
|
|
|
from v2.core.parallel.plan import ParallelPlan
|
|
from v2.core.parallel.validation import validate_plan
|
|
|
|
|
|
class FakeDeviceMesh:
|
|
|
|
def __init__(self, plan: ParallelPlan):
|
|
self.plan = plan
|
|
self.order = plan.mesh_order or list(plan.axes.keys())
|
|
self.shape = tuple(plan.degree(a) for a in self.order)
|
|
self.world_size = plan.world_size()
|
|
|
|
def group_ranks(self, axis: str) -> list[list[int]]:
|
|
"""All collective groups along one axis (each is a list of global ranks)."""
|
|
if axis not in self.order:
|
|
return [[0]] # degree-one trivial group
|
|
axis_idx = self.order.index(axis)
|
|
groups: list[list[int]] = []
|
|
ranges = [range(d) for d in self.shape]
|
|
seen: set[tuple] = set()
|
|
for coord in product(*ranges):
|
|
key = coord[:axis_idx] + coord[axis_idx + 1:]
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
group = []
|
|
for i in range(self.shape[axis_idx]):
|
|
c = list(coord)
|
|
c[axis_idx] = i
|
|
group.append(self._coord_to_rank(tuple(c)))
|
|
groups.append(group)
|
|
return groups
|
|
|
|
def _coord_to_rank(self, coord: tuple[int, ...]) -> int:
|
|
rank = 0
|
|
for c, d in zip(coord, self.shape, strict=False):
|
|
rank = rank * d + c
|
|
return rank
|
|
|
|
def __repr__(self) -> str:
|
|
return f"FakeDeviceMesh(order={self.order}, shape={self.shape}, world={self.world_size})"
|
|
|
|
|
|
def build_mesh(plan: ParallelPlan, card=None, **kw) -> FakeDeviceMesh:
|
|
validate_plan(plan, card=card, **kw)
|
|
return FakeDeviceMesh(plan)
|