from __future__ import annotations import sys import unittest import os from pathlib import Path ROOT = Path(__file__).resolve().parents[1] WORKSPACE = ROOT.parent os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") os.environ.setdefault("XDG_CACHE_HOME", "/tmp") sys.path.insert(0, str(ROOT)) sys.path.insert(0, str(WORKSPACE / "pp-scheduler")) from pp_scheduler import ( # noqa: E402 BigMac, BigMacVPP, DistTrainEncoder1F1B, DistTrain1F1B, GPipe, OpType, UnifiedBigMacVPP, VPP1F1BSchedule, ) from pp_simulator import OpDurationSpec, PipelineSimulator # noqa: E402 from web.app import run_comparison # noqa: E402 class PipelineSimulatorTest(unittest.TestCase): def test_gpipe_fixed_duration_result_is_valid(self): scheduler = GPipe(pp_size=2, num_microbatches=4) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate( { OpType.FORWARD: OpDurationSpec(mean=1.0), OpType.BACKWARD: OpDurationSpec(mean=2.0), }, seed=1, ) simulator.validate_result(result) self.assertGreater(result.summary["makespan"], 0.0) self.assertEqual(result.metadata["op_count"], 16) def test_vpp_dependencies_are_valid(self): scheduler = VPP1F1BSchedule(pp_size=2, vpp_size=2, num_microbatches=4) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate( { "F": OpDurationSpec(mean=1.0, variance=0.01), "B": OpDurationSpec(mean=2.0, variance=0.01), }, seed=2, ) simulator.validate_result(result) self.assertEqual(result.metadata["vpp_size"], 2) def test_chrome_trace_export(self): scheduler = GPipe(pp_size=2, num_microbatches=2) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate( { OpType.FORWARD: OpDurationSpec(mean=1.0), OpType.BACKWARD: OpDurationSpec(mean=2.0), } ) trace = result.to_chrome_trace_dict() complete_events = [event for event in trace["traceEvents"] if event["ph"] == "X"] thread_name_events = [ event for event in trace["traceEvents"] if event["ph"] == "M" and event["name"] == "thread_name" ] self.assertEqual(len(complete_events), result.metadata["op_count"]) self.assertEqual(len(thread_name_events), 2) self.assertIn("summary", trace) self.assertEqual(complete_events[0]["pid"], 1) self.assertIn(complete_events[0]["tid"], {1, 2}) self.assertIn("batch_id", complete_events[0]["args"]) self.assertIn("dep_reasons", complete_events[0]["args"]) def test_chrome_trace_keeps_pp4_vpp2_f3_c0_and_chunk_categories(self): scheduler = UnifiedBigMacVPP(pp_size=4, vpp_size=2, num_microbatches=8, ve_forward_limit=3) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate(_default_specs()) trace = result.to_chrome_trace_dict() complete_events = [event for event in trace["traceEvents"] if event["ph"] == "X"] f3_c0_events = [event for event in complete_events if event["name"] == "F3-C0"] self.assertEqual(len(f3_c0_events), 4) self.assertEqual({event["args"]["pp_rank"] for event in f3_c0_events}, {0, 1, 2, 3}) f0_c0 = next(event for event in complete_events if event["name"] == "F0-C0") f0_c1 = next(event for event in complete_events if event["name"] == "F0-C1") self.assertEqual(f0_c0["cat"], "F/chunk_a") self.assertEqual(f0_c1["cat"], "F/chunk_b") self.assertNotEqual(f0_c0["cat"], f0_c1["cat"]) def test_perfetto_compat_trace_uses_unique_event_names(self): scheduler = UnifiedBigMacVPP(pp_size=4, vpp_size=2, num_microbatches=8, ve_forward_limit=3) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate(_default_specs()) trace = result.to_chrome_trace_dict(perfetto_compat=True) complete_events = [event for event in trace["traceEvents"] if event["ph"] == "X"] f3_c0_events = [ event for event in complete_events if event["args"]["label"] == "F3-C0" ] self.assertEqual(len(f3_c0_events), 4) self.assertEqual( {event["name"] for event in f3_c0_events}, { "rank0/F/chunk_a/mb3", "rank1/F/chunk_a/mb3", "rank2/F/chunk_a/mb3", "rank3/F/chunk_a/mb3", }, ) def test_monte_carlo_report_statistics(self): scheduler = GPipe(pp_size=2, num_microbatches=2) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) report = simulator.monte_carlo( { OpType.FORWARD: OpDurationSpec(mean=1.0, variance=0.01), OpType.BACKWARD: OpDurationSpec(mean=2.0, variance=0.01), }, num_trials=5, seed=11, validate=True, ) self.assertEqual(report.metadata["num_trials"], 5) self.assertEqual(len(report.trial_summaries), 5) self.assertEqual(report.statistics["makespan"]["count"], 5) self.assertIn("0", report.statistics["rank_utilization"]) self.assertIn("F", report.statistics["op_type_time"]) def test_monte_carlo_fixed_duration_has_zero_makespan_std(self): scheduler = GPipe(pp_size=2, num_microbatches=2) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) report = simulator.monte_carlo( { OpType.FORWARD: OpDurationSpec(mean=1.0), OpType.BACKWARD: OpDurationSpec(mean=2.0), }, num_trials=3, seed=3, ) self.assertEqual(report.statistics["makespan"]["std"], 0.0) self.assertEqual(report.statistics["makespan"]["min"], report.statistics["makespan"]["max"]) def test_monte_carlo_rejects_non_positive_trial_count(self): scheduler = GPipe(pp_size=2, num_microbatches=2) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) with self.assertRaises(ValueError): simulator.monte_carlo( { OpType.FORWARD: OpDurationSpec(mean=1.0), OpType.BACKWARD: OpDurationSpec(mean=2.0), }, num_trials=0, ) def test_bigmac_llm_forward_waits_for_ve_unit(self): scheduler = BigMac(pp_size=2, num_microbatches=4, ve_forward_limit=3) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) target = _find_op(simulator, rank=0, op_type="F", microbatch_id=0, chunk_id=None) vf0 = _find_op(simulator, rank=0, op_type="VF", microbatch_id=0, chunk_id=None) vf1 = _find_op(simulator, rank=1, op_type="VF", microbatch_id=1, chunk_id=None) reasons = simulator.graph.dependency_reasons[target.id] self.assertIn("ve_forward_ready", reasons[vf0.id]) self.assertNotIn(vf1.id, reasons) def test_bigmac_llm_forward_waits_only_for_matching_ve_forward(self): scheduler = BigMac( pp_size=4, num_microbatches=16, ve_forward_limit=3, check_ve_forward_limit=False, ) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) f12 = _find_op(simulator, rank=0, op_type="F", microbatch_id=12, chunk_id=None) vf12 = _find_op(simulator, rank=0, op_type="VF", microbatch_id=12, chunk_id=None) peer_vfs = [ _find_op(simulator, rank=rank, op_type="VF", microbatch_id=12 + rank, chunk_id=None) for rank in range(1, 4) ] reasons = simulator.graph.dependency_reasons[f12.id] self.assertIn("ve_forward_ready", reasons[vf12.id]) for peer_vf in peer_vfs: self.assertNotIn(peer_vf.id, reasons) def test_bigmac_low_ve_forward_limit_exposes_bubbles(self): scheduler = BigMac( pp_size=4, num_microbatches=16, ve_forward_limit=2, check_ve_forward_limit=False, ) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate(_default_specs(), seed=23) simulator.validate_result(result) idle_counts = { rank: sum(1 for op in ops if op.op_type == OpType.IDLE) for rank, ops in scheduler.schedules.items() } self.assertGreater(sum(idle_counts.values()), 0) def test_bigmac_low_ve_forward_limit_interleaves_ve_backward_before_next_ve_forward(self): scheduler = BigMac( pp_size=4, num_microbatches=16, ve_forward_limit=3, check_ve_forward_limit=False, ) scheduler.generate_schedule() rank0_ops = [ str(op) for op in scheduler.schedules[0] if op.op_type != OpType.IDLE ] self.assertLess(rank0_ops.index("B8"), rank0_ops.index("VB0")) self.assertEqual(rank0_ops[rank0_ops.index("VB0") + 1], "VF12") self.assertLess(rank0_ops.index("VF12"), rank0_ops.index("B9")) def test_bigmac_vpp_chunk_wrap_dependencies(self): scheduler = BigMacVPP(pp_size=2, vpp_size=2, num_microbatches=4, ve_forward_limit=3) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) f_rank0_chunk1 = _find_op(simulator, rank=0, op_type="F", microbatch_id=0, chunk_id=1) f_rank1_chunk0 = _find_op(simulator, rank=1, op_type="F", microbatch_id=0, chunk_id=0) b_rank1_chunk0 = _find_op(simulator, rank=1, op_type="B", microbatch_id=0, chunk_id=0) b_rank0_chunk1 = _find_op(simulator, rank=0, op_type="B", microbatch_id=0, chunk_id=1) reasons = simulator.graph.dependency_reasons self.assertIn("vpp_forward_from_previous_chunk", reasons[f_rank0_chunk1.id][f_rank1_chunk0.id]) self.assertIn("vpp_backward_from_next_chunk", reasons[b_rank1_chunk0.id][b_rank0_chunk1.id]) def test_unified_generator_dependencies_use_matching_microbatch(self): scheduler = UnifiedBigMacVPP(pp_size=2, vpp_size=2, num_microbatches=4, ve_forward_limit=3) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) gf0 = _find_op(simulator, rank=0, op_type="GF", microbatch_id=0, chunk_id=1) gb0 = _find_op(simulator, rank=0, op_type="GB", microbatch_id=0, chunk_id=0) b_last = _find_op(simulator, rank=1, op_type="B", microbatch_id=0, chunk_id=1) final_f0 = _find_op(simulator, rank=1, op_type="F", microbatch_id=0, chunk_id=1) final_f1 = _find_op(simulator, rank=1, op_type="F", microbatch_id=1, chunk_id=1) gb1 = _find_op(simulator, rank=1, op_type="GB", microbatch_id=1, chunk_id=0) reasons = simulator.graph.dependency_reasons self.assertIn("llm_output_ready", reasons[gf0.id][final_f0.id]) self.assertNotIn(final_f1.id, reasons[gf0.id]) self.assertIn("generator_backward_consumes_forward", reasons[gb0.id][gf0.id]) self.assertNotIn(gb1.id, reasons[gb0.id]) self.assertIn("generator_backward_ready", reasons[b_last.id][gb0.id]) self.assertNotIn(gb1.id, reasons[b_last.id]) def test_unified_generator_backward_does_not_wait_for_peer_generator_forward(self): scheduler = UnifiedBigMacVPP(pp_size=4, vpp_size=2, num_microbatches=8, ve_forward_limit=3) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) gb4 = _find_op(simulator, rank=0, op_type="GB", microbatch_id=4, chunk_id=0) gf4 = _find_op(simulator, rank=0, op_type="GF", microbatch_id=4, chunk_id=1) peer_gfs = [ _find_op(simulator, rank=rank, op_type="GF", microbatch_id=4 + rank, chunk_id=1) for rank in range(1, 4) ] reasons = simulator.graph.dependency_reasons[gb4.id] self.assertIn("generator_backward_consumes_forward", reasons[gf4.id]) for peer_gf in peer_gfs: self.assertNotIn(peer_gf.id, reasons) def test_bigmac_ve_backward_waits_only_for_matching_llm_backward(self): scheduler = BigMac( pp_size=4, num_microbatches=32, ve_forward_limit=3, check_ve_forward_limit=False, ) scheduler.generate_schedule() simulator = PipelineSimulator.from_scheduler(scheduler) vb30 = _find_op(simulator, rank=2, op_type="VB", microbatch_id=30, chunk_id=None) b30 = _find_op(simulator, rank=0, op_type="B", microbatch_id=30, chunk_id=None) peer_bs = [ _find_op(simulator, rank=0, op_type="B", microbatch_id=mb, chunk_id=None) for mb in (28, 29, 31) ] reasons = simulator.graph.dependency_reasons[vb30.id] self.assertIn("llm_input_gradient_ready", reasons[b30.id]) for peer_b in peer_bs: self.assertNotIn(peer_b.id, reasons) def test_disttrain_1f1b_layout_and_dependencies(self): scheduler = DistTrain1F1B(pp_size=4, num_microbatches=8) scheduler.generate_schedule() self.assertEqual(scheduler.pp_size, 6) self.assertEqual(scheduler.llm_pp_size, 4) self.assertEqual(scheduler.encoder_stage, 0) self.assertEqual(scheduler.generator_stage, 5) simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate(_default_specs(), seed=5) simulator.validate_result(result) vf0 = _find_op(simulator, rank=0, op_type="VF", microbatch_id=0, chunk_id=None) f0_llm0 = _find_op(simulator, rank=1, op_type="F", microbatch_id=0, chunk_id=None) f0_llm_last = _find_op(simulator, rank=4, op_type="F", microbatch_id=0, chunk_id=None) gf0 = _find_op(simulator, rank=5, op_type="GF", microbatch_id=0, chunk_id=None) gb0 = _find_op(simulator, rank=5, op_type="GB", microbatch_id=0, chunk_id=None) b0_llm_last = _find_op(simulator, rank=4, op_type="B", microbatch_id=0, chunk_id=None) b0_llm0 = _find_op(simulator, rank=1, op_type="B", microbatch_id=0, chunk_id=None) vb0 = _find_op(simulator, rank=0, op_type="VB", microbatch_id=0, chunk_id=None) reasons = simulator.graph.dependency_reasons self.assertIn("encoder_forward_ready", reasons[f0_llm0.id][vf0.id]) self.assertIn("llm_output_ready", reasons[gf0.id][f0_llm_last.id]) self.assertIn("generator_backward_consumes_forward", reasons[gb0.id][gf0.id]) self.assertIn("generator_backward_ready", reasons[b0_llm_last.id][gb0.id]) self.assertIn("llm_input_gradient_ready", reasons[vb0.id][b0_llm0.id]) gf1 = _find_op(simulator, rank=5, op_type="GF", microbatch_id=1, chunk_id=None) self.assertNotIn(gf1.id, reasons[gb0.id]) def test_disttrain_1f1b_activation_limits(self): scheduler = DistTrain1F1B(pp_size=4, num_microbatches=8) scheduler.generate_schedule() forward_ops = {OpType.VEFORWARD, OpType.FORWARD, OpType.GENFORWARD} backward_ops = {OpType.VEBACKWARD, OpType.BACKWARD, OpType.GENBACKWARD} peaks = {} for rank, ops in scheduler.schedules.items(): live = 0 peak = 0 for op in ops: if op.op_type in forward_ops: live += 1 peak = max(peak, live) elif op.op_type in backward_ops: live -= 1 peaks[rank] = peak self.assertLessEqual(peak, scheduler.memory_limits[rank]) self.assertEqual(live, 0) self.assertEqual(peaks[5], 1) self.assertEqual(peaks[4], 2) def test_disttrain_uses_uniform_module_durations(self): scheduler = DistTrain1F1B(pp_size=4, num_microbatches=8) scheduler.generate_schedule() forward_ops = {OpType.VEFORWARD, OpType.FORWARD, OpType.GENFORWARD} backward_ops = {OpType.VEBACKWARD, OpType.BACKWARD, OpType.GENBACKWARD} for ops in scheduler.schedules.values(): for op in ops: if op.op_type in forward_ops: self.assertEqual(op.duration, scheduler.forward_duration) elif op.op_type in backward_ops: self.assertEqual(op.duration, scheduler.backward_duration) rank4_ops = [ (op.op_type, op.microbatch_id) for op in scheduler.schedules[4] if op.op_type is not OpType.IDLE ] self.assertLess( rank4_ops.index((OpType.BACKWARD, 0)), rank4_ops.index((OpType.FORWARD, 2)), ) def test_disttrain_encoder_1f1b_layout_and_dependencies(self): scheduler = DistTrainEncoder1F1B(pp_size=4, num_microbatches=8) scheduler.generate_schedule() self.assertEqual(scheduler.pp_size, 5) self.assertEqual(scheduler.llm_pp_size, 4) self.assertEqual(scheduler.encoder_stage, 0) self.assertIsNone(scheduler.generator_stage) self.assertFalse(scheduler.pipeline_layout["has_generator"]) all_op_types = {op.op_type for ops in scheduler.schedules.values() for op in ops} self.assertNotIn(OpType.GENFORWARD, all_op_types) self.assertNotIn(OpType.GENBACKWARD, all_op_types) simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate(_default_specs(), seed=11) simulator.validate_result(result) vf0 = _find_op(simulator, rank=0, op_type="VF", microbatch_id=0, chunk_id=None) f0_llm0 = _find_op(simulator, rank=1, op_type="F", microbatch_id=0, chunk_id=None) f0_llm_last = _find_op(simulator, rank=4, op_type="F", microbatch_id=0, chunk_id=None) b0_llm_last = _find_op(simulator, rank=4, op_type="B", microbatch_id=0, chunk_id=None) b0_llm0 = _find_op(simulator, rank=1, op_type="B", microbatch_id=0, chunk_id=None) vb0 = _find_op(simulator, rank=0, op_type="VB", microbatch_id=0, chunk_id=None) reasons = simulator.graph.dependency_reasons self.assertIn("encoder_forward_ready", reasons[f0_llm0.id][vf0.id]) self.assertIn("backward_consumes_forward_activation", reasons[b0_llm_last.id][f0_llm_last.id]) self.assertNotIn("generator_backward_ready", { reason for dep_reasons in reasons[b0_llm_last.id].values() for reason in dep_reasons }) self.assertIn("llm_input_gradient_ready", reasons[vb0.id][b0_llm0.id]) def test_compare_uses_shared_duration_per_rank_op_microbatch_chunk(self): payload = { "scheduler_a": "GPipe", "scheduler_b": "Schedule1F1B", "workload": { "pp_size": 2, "vpp_size": 1, "num_microbatches": 4, "ve_forward_limit": 3, "include_encoder": False, "include_generator": False, }, "duration_specs": { "F": {"mean": 1.0, "variance": 0.1}, "B": {"mean": 2.0, "variance": 0.1}, }, "seed": 17, } response = run_comparison(payload) result_a = response["result_a"] result_b = response["result_b"] for rank in (0, 1): a_op = _find_result_op(result_a, rank=rank, op_type="F", microbatch_id=0, chunk_id=None) b_op = _find_result_op(result_b, rank=rank, op_type="F", microbatch_id=0, chunk_id=None) self.assertEqual(a_op["duration"], b_op["duration"]) rank0 = _find_result_op(result_a, rank=0, op_type="F", microbatch_id=0, chunk_id=None) rank1 = _find_result_op(result_a, rank=1, op_type="F", microbatch_id=0, chunk_id=None) self.assertNotEqual(rank0["duration"], rank1["duration"]) def test_compare_rejects_incompatible_scheduler_shape(self): payload = { "scheduler_a": "GPipe", "scheduler_b": "BigMac", "workload": { "pp_size": 2, "vpp_size": 1, "num_microbatches": 4, "ve_forward_limit": 3, "include_encoder": False, "include_generator": False, }, "duration_specs": _default_web_specs(), "seed": 17, } with self.assertRaises(ValueError): run_comparison(payload) def test_compare_can_allow_vpp_and_non_vpp_with_llm_duration_scale(self): payload = { "scheduler_a": "GPipe", "scheduler_b": "VPP1F1BSchedule", "workload": { "pp_size": 2, "vpp_size": 2, "num_microbatches": 4, "ve_forward_limit": 3, "include_encoder": False, "include_generator": False, }, "duration_specs": _default_web_specs(), "seed": 17, } response = run_comparison(payload) self.assertTrue(response["comparison"]["workload"]["allow_cross_vpp"]) self.assertEqual(response["result_b"]["metadata"]["vpp_size"], 2) non_vpp_f = _find_result_op( response["result_a"], rank=0, op_type="F", microbatch_id=0, chunk_id=None, ) vpp_f_c0 = _find_result_op( response["result_b"], rank=0, op_type="F", microbatch_id=0, chunk_id=0, ) vpp_f_c1 = _find_result_op( response["result_b"], rank=0, op_type="F", microbatch_id=0, chunk_id=1, ) self.assertEqual(vpp_f_c0["duration"], non_vpp_f["duration"] / 2) self.assertEqual(vpp_f_c1["duration"], non_vpp_f["duration"] / 2) def test_compare_adjusts_disttrain_stage_budget_and_llm_duration_scale(self): payload = { "scheduler_a": "BigMac", "scheduler_b": "DistTrainEncoder1F1B", "workload": { "pp_size": 4, "vpp_size": 1, "num_microbatches": 8, "ve_forward_limit": 3, "include_encoder": True, "include_generator": False, }, "duration_specs": _default_web_specs(), "seed": 17, } response = run_comparison(payload) disttrain = response["result_b"] self.assertEqual(disttrain["metadata"]["pp_size"], 4) self.assertEqual(disttrain["metadata"]["pipeline_layout"]["llm_pp_size"], 3) bigmac_f = _find_result_op( response["result_a"], rank=1, op_type="F", microbatch_id=0, chunk_id=None, ) disttrain_f = _find_result_op( disttrain, rank=1, op_type="F", microbatch_id=0, chunk_id=None, ) self.assertAlmostEqual(disttrain_f["duration"], bigmac_f["duration"] * 4 / 3) def _find_op(simulator, *, rank, op_type, microbatch_id, chunk_id): for op in simulator.plan.ops_by_rank[rank]: if ( op.op_type == op_type and op.microbatch_id == microbatch_id and op.chunk_id == chunk_id ): return op raise AssertionError( f"missing op rank={rank} op_type={op_type} " f"microbatch_id={microbatch_id} chunk_id={chunk_id}" ) def _find_result_op(result, *, rank, op_type, microbatch_id, chunk_id): for op in result["ops"]: if ( op["rank"] == rank and op["op_type"] == op_type and op["microbatch_id"] == microbatch_id and op["chunk_id"] == chunk_id ): return op raise AssertionError( f"missing result op rank={rank} op_type={op_type} " f"microbatch_id={microbatch_id} chunk_id={chunk_id}" ) def _default_specs(): return { "F": OpDurationSpec(mean=1.0), "B": OpDurationSpec(mean=2.0), "VF": OpDurationSpec(mean=0.8), "VB": OpDurationSpec(mean=0.9), "GF": OpDurationSpec(mean=0.6), "GB": OpDurationSpec(mean=0.7), } def _default_web_specs(): return { "F": {"mean": 1.0, "variance": 0.0}, "B": {"mean": 2.0, "variance": 0.0}, "VF": {"mean": 0.8, "variance": 0.0}, "VB": {"mean": 0.9, "variance": 0.0}, "GF": {"mean": 0.6, "variance": 0.0}, "GB": {"mean": 0.7, "variance": 0.0}, } if __name__ == "__main__": unittest.main()