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`.
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
"""ParallelPlan — parallelism as a model contract.
|
|
|
|
Parallelism is not a launch flag; it affects cache keys, scheduling, transport,
|
|
capture, and parity, so it lives on the card. Declarative, validated, compiled to a
|
|
mesh via a ParallelDims-style builder (``parallel/mesh.py``). This module is a pure
|
|
leaf (no card/runtime imports) so ``card/`` can hold plans without a cycle.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
|
|
# Canonical axis names. cfgp is <=2.
|
|
AXIS_NAMES = ("dp", "tp", "sp", "cp", "cfgp", "pp_patch", "vae", "ep", "fsdp", "role", "replica")
|
|
|
|
|
|
@dataclass
|
|
class ParallelPlan:
|
|
axes: dict[str, int] = field(default_factory=dict) # e.g. {"tp": 2, "sp": 4, "cfgp": 2}
|
|
mesh_order: list[str] = field(default_factory=list)
|
|
placement: str = "colocated"
|
|
communication: dict[str, str] = field(default_factory=dict)
|
|
# applicability conditions travel with axes
|
|
applicability: dict[str, dict] = field(default_factory=dict)
|
|
per_axis_communication: dict[str, str] = field(default_factory=dict)
|
|
|
|
def degree(self, axis: str) -> int:
|
|
return int(self.axes.get(axis, 1))
|
|
|
|
def world_size(self) -> int:
|
|
w = 1
|
|
for v in self.axes.values():
|
|
w *= int(v)
|
|
return max(w, 1)
|
|
|
|
@property
|
|
def hash(self) -> str:
|
|
"""Stable hash — part of the CacheKey (parallel_plan_hash)."""
|
|
payload = json.dumps({"axes": self.axes, "order": self.mesh_order}, sort_keys=True)
|
|
return hashlib.sha256(payload.encode()).hexdigest()[:16]
|
|
|
|
@classmethod
|
|
def single(cls) -> ParallelPlan:
|
|
"""The default deployment: one device, all degree-one trivial groups."""
|
|
return cls(axes={"dp": 1}, mesh_order=["dp"])
|