Spaces:
Sleeping
Sleeping
File size: 3,173 Bytes
89cb5bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from .duration import op_type_value
COMPUTE_OP_TYPES = frozenset({"F", "B", "VF", "VB", "GF", "GB"})
@dataclass(frozen=True)
class PlanOp:
id: str
rank: int
index: int
original_index: int
op_type: str
microbatch_id: int | None
chunk_id: int | None
base_duration: float
label: str
@property
def is_forward(self) -> bool:
return self.op_type == "F"
@property
def is_backward(self) -> bool:
return self.op_type == "B"
@dataclass(frozen=True)
class Plan:
ops_by_rank: dict[int, list[PlanOp]]
pp_size: int
num_microbatches: int
scheduler_name: str
vpp_size: int = 1
pipeline_layout: dict[str, int] | None = None
@property
def ops(self) -> list[PlanOp]:
return [
op
for rank in sorted(self.ops_by_rank)
for op in self.ops_by_rank[rank]
]
@classmethod
def from_scheduler(cls, scheduler: Any, *, include_op_types=None) -> "Plan":
if not getattr(scheduler, "schedules", None):
scheduler.generate_schedule()
allowed = frozenset(include_op_types or COMPUTE_OP_TYPES)
ops_by_rank: dict[int, list[PlanOp]] = {}
pp_size = int(getattr(scheduler, "pp_size"))
for rank in range(pp_size):
rank_ops: list[PlanOp] = []
for original_index, raw_op in enumerate(scheduler.schedules.get(rank, [])):
op_value = op_type_value(getattr(raw_op, "op_type", ""))
if op_value not in allowed:
continue
rank_index = len(rank_ops)
microbatch_id = getattr(raw_op, "microbatch_id", None)
chunk_id = getattr(raw_op, "chunk_id", None)
op_id = _make_op_id(rank, rank_index, op_value, microbatch_id, chunk_id)
rank_ops.append(
PlanOp(
id=op_id,
rank=rank,
index=rank_index,
original_index=original_index,
op_type=op_value,
microbatch_id=microbatch_id,
chunk_id=chunk_id,
base_duration=float(getattr(raw_op, "duration", 0.0)),
label=str(raw_op),
)
)
ops_by_rank[rank] = rank_ops
return cls(
ops_by_rank=ops_by_rank,
pp_size=pp_size,
num_microbatches=int(getattr(scheduler, "num_microbatches")),
scheduler_name=scheduler.__class__.__name__,
vpp_size=int(getattr(scheduler, "vpp_size", 1)),
pipeline_layout=getattr(scheduler, "pipeline_layout", None),
)
def _make_op_id(
rank: int,
index: int,
op_type: str,
microbatch_id: int | None,
chunk_id: int | None,
) -> str:
mb = "none" if microbatch_id is None else str(microbatch_id)
chunk = "none" if chunk_id is None else str(chunk_id)
return f"r{rank}:i{index}:{op_type}:mb{mb}:c{chunk}"
|