Spaces:
Running
Running
| # PP schedule config | |
| from src.execution_model import ScheduleConfig | |
| from src.strategies import generate_1f1b_interleave_overlap_schedule, generate_dualpipe_v_schedule | |
| p = 4 # PP size | |
| v = 2 # number of virtual stages | |
| m = 16 # total microbatches | |
| # stage time config | |
| F = 2.0 # forward time in one PP rank for all stages | |
| W = 2.0 # backward_W time in one PP rank for all stages | |
| D = 2.0 # backward_D time in one PP rank for all stages | |
| B = W + D # backward time in one PP rank for all stages | |
| FwB = 5.5 # overlapped forward backward time in one PP rank for all stages | |
| op_times = { | |
| "forward": F, | |
| "backward": B, | |
| "backward_D": D, | |
| "backward_W": W, | |
| "overlapped_forward_backward": FwB | |
| } | |
| def dualpipe_v_execution_time_by_formula(): | |
| # Formula from the image | |
| item_1 = ((p - 1) / 2) * F | |
| item_2 = (p + 0.5) * F + (p / 2 + 1) * B | |
| item_3 = (m - (p / 2 + 1)) * FwB | |
| print(f"item_1: {item_1}, item_2: {item_2}, item_3: {item_3}") | |
| total_time = item_1 + item_2 + item_3 | |
| return total_time | |
| def dualpipe_v_execution_time_by_formula_detailed(): | |
| # Correct formula | |
| local_F = F / 2 | |
| local_B = B / 2 | |
| local_D = D / 2 | |
| local_W = W / 2 | |
| local_FwB = FwB / 2 | |
| forward_bubble = (p - 1) * local_F # forward bubble | |
| forward_time = 2 * p * local_F | |
| overlapped_time = (2 * (m-p)-1) * local_FwB + (p-1) * local_FwB | |
| backward_time = (2*p-1) * local_D + local_W | |
| other_time = 2 * local_B + local_F | |
| active_time = (2 * (m-p)-1) * local_FwB + (2*p+1) * (local_F + local_B) | |
| total_time = forward_bubble + forward_time + overlapped_time + backward_time + other_time | |
| bubble_time = total_time - active_time | |
| assert bubble_time == (p-1)*(local_FwB + local_B - 3*local_W) | |
| return total_time | |
| def dualpipe_v_execution_time_by_emulate(): | |
| op_times_per_stage = { | |
| "forward": F / 2, | |
| "backward": B / 2, | |
| "backward_D": D / 2, | |
| "backward_W": W / 2, | |
| "overlapped_forward_backward": FwB / 2 | |
| } | |
| print(f"op_times_per_stage: {op_times_per_stage}") | |
| dualpipe_schedule_config = ScheduleConfig( | |
| num_devices=p, | |
| num_stages=p*2, | |
| num_batches=m, | |
| p2p_latency=0.0, | |
| op_times=op_times_per_stage, | |
| split_backward=True, | |
| placement_strategy="dualpipe_v", | |
| ) | |
| dual_pipe_schedule = generate_dualpipe_v_schedule(dualpipe_schedule_config) | |
| dual_pipe_schedule.execute() | |
| return dual_pipe_schedule.get_total_execution_time() | |
| def overlap_1f1b_execution_time_by_emulate(): | |
| op_times_per_stage = { | |
| "forward": F / v, | |
| "backward": B / v, | |
| "backward_D": D / v, | |
| "backward_W": W / v, | |
| "overlapped_forward_backward": FwB / v | |
| } | |
| overlap_1f1b_schedule_config = ScheduleConfig( | |
| num_devices=p, | |
| num_stages=p*v, | |
| num_batches=m, | |
| p2p_latency=0.0, | |
| op_times=op_times_per_stage, | |
| split_backward=False, | |
| placement_strategy="interleave", | |
| ) | |
| overlap_1f1b_schedule = generate_1f1b_interleave_overlap_schedule(overlap_1f1b_schedule_config) | |
| overlap_1f1b_schedule.execute() | |
| return overlap_1f1b_schedule.get_total_execution_time() | |
| def overlap_1f1b_execution_time_by_formula(): | |
| forward_bubble = (p-1) * F / v | |
| backward_bubble = (p-1) * B / v | |
| non_overlapped_batches = p*(v - 1) + 1 | |
| forward_backward_time = non_overlapped_batches * (F + B) / v | |
| overlapped_time = (m*v - non_overlapped_batches) * FwB / v | |
| total_time = forward_bubble + backward_bubble + forward_backward_time + overlapped_time | |
| return total_time | |
| print(f"DualPipe-V by emulate: {dualpipe_v_execution_time_by_emulate()}") | |
| print(f"DualPipe-V by formula detailed: {dualpipe_v_execution_time_by_formula_detailed()}") | |
| print(f"Overlap-1f1b by emulate: {overlap_1f1b_execution_time_by_emulate()}") | |
| print(f"Overlap-1f1b by formula: {overlap_1f1b_execution_time_by_formula()}") | |