Spaces:
Running
Running
File size: 3,692 Bytes
2ae9b28 bb52925 16ed969 a49be3b 06107a3 bb52925 a49be3b bb52925 a49be3b bb52925 a49be3b 2ae9b28 a49be3b bb52925 a49be3b bb52925 a49be3b 2ae9b28 a49be3b 06107a3 bb52925 06107a3 bb52925 06107a3 bb52925 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 |
from src.execution_model import ScheduleConfig
from src.strategies import (
generate_1f1b_interleave_schedule,
generate_1f1b_overlap_schedule,
generate_1f1b_schedule,
generate_zero_bubble_1p_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)
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,
)
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)
if __name__ == "__main__":
main()
|