File size: 6,195 Bytes
4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 29c9647 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 29c9647 4b2c8d9 |
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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import gradio as gr
import auto_schedule
import v_schedule
from PIL import Image
from svg_event import render_manual_graph
import pathlib
def greet(name, is_morning, temperature):
salutation = "Good morning" if is_morning else "Good evening"
greeting = f"{salutation} {name}. It is {temperature} degrees today"
celsius = (temperature - 32) * 5 / 9
return greeting, round(celsius, 2)
def percentage(x):
return f"{x*100:.2f}%"
def get_schedule_time(result):
result = [
list(filter(lambda x: x.type in {'F', 'B', 'W'}, r)) for r in result
]
time = max(
[
max([x.completion_time for x in stage]) - min([x.start_time for x in stage]) for stage in result
]
)
return time
img_queue = []
def get_schedule_image(result, max_time):
result = [
list(filter(lambda x: x.type in {'F', 'B', 'W'}, r)) for r in result
]
svg = render_manual_graph(result, max_time, len(result[0]) <= 72)
img_queue.append(svg)
if len(img_queue) > 32:
poped = img_queue.pop(0)
pathlib.Path(poped).unlink()
return pathlib.Path(svg)
def calculate(p, m, f, b, w, c, mem):
if mem < p:
baseline_time=None
baseline_bubble=None
baseline_acceleration=None
baseline_image=None
else:
baseline_time=(f+b+w)*m + (f+b+w+c)*(p-1)
baseline_bubble=percentage(baseline_time/(f+b+w)/m - 1)
baseline_acceleration=percentage(0)
baseline_image=None
zb_result = auto_schedule.auto_schedule(p, m, auto_schedule.GraphConfig(
cost_f=f,
cost_b=b,
cost_w=w,
cost_comm=c,
max_mem=mem * 2,
print_scaling=1000
))
zb_time=get_schedule_time(zb_result)
zb_bubble=percentage(zb_time/(f+b+w)/m - 1)
zb_acceleration=percentage(baseline_time/zb_time - 1) if baseline_time is not None else None
if mem < p:
zbv_time=None
zbv_bubble=None
zbv_acceleration=None
zbv_image=None
else:
zbv_graph = v_schedule.PipelineGraph(
n_stage=p,
n_micro=m,
f_cost=f/2,
b_cost=b/2,
w_cost=w/2,
c_cost=c,
f_mem=2,
b_mem=-1,
w_mem=-1,
max_mem=mem * 4,
)
zbv_result = zbv_graph.get_v_schedule()
zbv_time = get_schedule_time(zbv_result)
zbv_bubble=percentage(zbv_time/(f+b+w)/m - 1)
zbv_acceleration=percentage(baseline_time/zbv_time - 1) if baseline_time is not None else None
max_time = max([baseline_time, zb_time, zbv_time])
zb_image = get_schedule_image(zb_result, max_time)
zbv_image = get_schedule_image(zbv_result, max_time)
return [baseline_time, baseline_bubble, baseline_acceleration, baseline_image, zb_time, zb_bubble, zb_acceleration, zb_image, zbv_time, zbv_bubble, zbv_acceleration, zbv_image]
with gr.Blocks() as demo:
gr.Markdown(open("description1.md").read())
gr.Markdown("# Pipeline Scheduler Playground")
with gr.Row():
with gr.Column(scale=1):
with gr.Group():
gr.Markdown("Basic Parameters")
with gr.Row():
p=gr.Number(label="Number of stages (p)", value=4, interactive=True, precision=0)
m=gr.Number(label="Number of microbatches (m)", value=12, interactive=True, precision=0)
with gr.Column(scale=2):
with gr.Group():
gr.Markdown("Costs. All costs are used as integers. For ZBV schedules, this is the time of two virtual stages on a stage combined.")
with gr.Row():
f=gr.Number(label="Time of F", value=8, interactive=True, precision=0)
b=gr.Number(label="Time of B", value=8, interactive=True, precision=0)
w=gr.Number(label="Time of W", value=8, interactive=True, precision=0)
c=gr.Number(label="Time of one P2P communication", value=1, interactive=True, precision=0)
with gr.Group():
gr.Markdown("Activation memory limit.")
def update_mem(p, s, mem):
print("update")
if s=="custom":
return mem
return p*int(s[:1])
memsel=gr.Radio(choices=["1p (Same as 1F1B)", "2p", "3p", "custom"], value="1p (Same as 1F1B)")
mem=gr.Number(label="Custom memory limit in terms of pending F on a stage. For ZBV schedules, this is relative to two virtual stages on a stage combined.", value=p.value, interactive=True, precision=0)
memsel.change(update_mem, inputs=[p, memsel, mem], outputs=mem)
p.change(update_mem, inputs=[p, memsel, mem], outputs=mem)
button=gr.Button("Calculate")
with gr.Group():
gr.Markdown("1F1B")
with gr.Row():
with gr.Column(scale=1):
baseline_time=gr.Textbox("", label="Longest Stage Time")
baseline_bubble=gr.Textbox("", label="Bubble Rate. Calculated as (1 - longest stage time/(F+B+W)/m).")
baseline_acceleration=gr.Textbox("", label="Acceleration compared to 1F1B")
with gr.Column(scale=4):
baseline_image=gr.Image(None, interactive=False, label="Schedule Image")
with gr.Group():
gr.Markdown("Zero Bubble Schedule")
with gr.Row():
with gr.Column(scale=1):
zb_time=gr.Textbox("", label="Longest Stage Time")
zb_bubble=gr.Textbox("", label="Bubble Rate. Calculated as (1 - longest stage time/(F+B+W)/m).")
zb_acceleration=gr.Textbox("", label="Acceleration compared to 1F1B")
with gr.Column(scale=4):
zb_image=gr.Image(None, interactive=False, label="Schedule Image")
with gr.Group():
gr.Markdown("Zero Bubble V Schedule (ZBV)")
with gr.Row():
with gr.Column(scale=1):
zbv_time=gr.Textbox("", label="Longest Stage Time")
zbv_bubble=gr.Textbox("", label="Bubble Rate. Calculated as (1 - longest stage time/(F+B+W)/m).")
zbv_acceleration=gr.Textbox("", label="Acceleration compared to 1F1B")
with gr.Column(scale=4):
zbv_image=gr.Image(None, interactive=False, label="Schedule Image")
button.click(calculate, inputs=[p, m, f, b, w, c, mem], outputs=[baseline_time, baseline_bubble, baseline_acceleration, baseline_image, zb_time, zb_bubble, zb_acceleration, zb_image, zbv_time, zbv_bubble, zbv_acceleration, zbv_image])
gr.Markdown(open("description2.md").read())
demo.launch()
|