Spaces:
Running
Running
File size: 6,270 Bytes
2ae9b28 bb52925 d3e7e66 bb52925 dc262c1 bb52925 99a6901 bb52925 16ed969 a49be3b 06107a3 bb52925 d3e7e66 99a6901 dc262c1 a49be3b bb52925 a49be3b bb52925 a49be3b 2ae9b28 a49be3b bb52925 a49be3b bb52925 c684983 a49be3b 2ae9b28 a49be3b 06107a3 bb52925 06107a3 bb52925 06107a3 bb52925 f140d7b bb52925 d3e7e66 bb52925 99a6901 dc262c1 a49be3b bb52925 |
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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
from src.execution_model import ScheduleConfig
from src.strategies import (
generate_1f1b_interleave_overlap_schedule,
generate_1f1b_interleave_schedule,
generate_1f1b_overlap_schedule,
generate_1f1b_schedule,
generate_dualpipe_v_schedule,
generate_zero_bubble_1p_schedule,
generate_dualpipe_schedule,
)
from src.visualizer import visualize_pipeline_parallelism_dash
import hydra
from omegaconf import DictConfig, OmegaConf
@hydra.main(config_path="conf", config_name="config", version_base=None)
def main(cfg: DictConfig) -> None:
"""Run pipeline parallelism simulation with the specified configuration."""
print(f"Running with configuration: {cfg}")
if cfg.strategy == "1f1b":
run_1f1b(cfg)
elif cfg.strategy == "interleave":
run_interleave(cfg)
elif cfg.strategy == "zb1p":
run_zero_bubble_1p(cfg)
elif cfg.strategy == "1f1b_overlap":
run_1f1b_overlap(cfg)
elif cfg.strategy == "1f1b_interleave_overlap":
run_1f1b_interleave_overlap(cfg)
elif cfg.strategy == "dualpipe":
run_dualpipe(cfg)
elif cfg.strategy == "dualpipe_v":
run_dualpipe_v(cfg)
else:
raise ValueError(f"Unknown strategy: {cfg.strategy}")
def run_1f1b(cfg: DictConfig) -> None:
"""Run 1F1B pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
op_times=op_times,
placement_strategy="standard",
)
schedule = generate_1f1b_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
def run_interleave(cfg: DictConfig) -> None:
"""Run interleaved pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
placement_strategy="interleave",
op_times=op_times,
microbatch_group_size_per_vp_stage=cfg.microbatch_group_size_per_vp_stage,
)
schedule = generate_1f1b_interleave_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
def run_zero_bubble_1p(cfg: DictConfig) -> None:
"""Run zero bubble 1P pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
op_times=op_times,
split_backward=True,
)
schedule = generate_zero_bubble_1p_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
def run_1f1b_overlap(cfg: DictConfig) -> None:
"""Run 1F1B overlap pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
op_times=op_times,
split_backward=False,
)
schedule = generate_1f1b_overlap_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
def run_1f1b_interleave_overlap(cfg: DictConfig) -> None:
"""Run 1F1B interleave overlapped pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
placement_strategy="interleave",
op_times=op_times,
)
schedule = generate_1f1b_interleave_overlap_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
def run_dualpipe(cfg: DictConfig) -> None:
"""Run DualPipe pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
op_times=op_times,
split_backward=True,
placement_strategy="dualpipe",
)
schedule = generate_dualpipe_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
def run_dualpipe_v(cfg: DictConfig) -> None:
"""Run DualPipeV pipeline parallelism simulation."""
# Convert OmegaConf to dict for op_times if it exists
op_times = (
OmegaConf.to_container(cfg.op_times) if hasattr(cfg, "op_times") else None
)
schedule_config = ScheduleConfig(
num_devices=cfg.num_devices,
num_stages=cfg.num_stages,
num_batches=cfg.num_batches,
p2p_latency=cfg.p2p_latency,
op_times=op_times,
split_backward=True,
placement_strategy="dualpipe_v",
)
schedule = generate_dualpipe_v_schedule(schedule_config)
schedule.execute()
visualize_pipeline_parallelism_dash(schedule, port=cfg.visualization_port)
if __name__ == "__main__":
main()
|