Spaces:
Running
Running
| 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() | |