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 f140d7b 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()
|