File size: 7,757 Bytes
4b2c8d9 be3048f 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 ac0b05c 4835f75 be3048f 4835f75 4b2c8d9 4835f75 4b2c8d9 4835f75 4b2c8d9 c47cbb6 4835f75 ac0b05c 4835f75 4b2c8d9 ac0b05c be3048f ac0b05c 4835f75 ac0b05c 4835f75 4b2c8d9 29c9647 ac0b05c 4b2c8d9 be3048f 4b2c8d9 be3048f 4b2c8d9 be3048f 4b2c8d9 ac0b05c 4b2c8d9 e7d1e67 4b2c8d9 ac0b05c 4b2c8d9 e7d1e67 4b2c8d9 ac0b05c 4b2c8d9 ac0b05c 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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 |
import gradio as gr
import auto_schedule
import v_schedule
import hand_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
baseline_result=None
else:
baseline_result = hand_schedule.get_hand_schedule(p, m, f, b + w, 0, c)
baseline_result = [
list(filter(lambda x: x.type in {'F', 'B'}, r)) for r in baseline_result
]
baseline_time = get_schedule_time(baseline_result)
baseline_bubble=percentage(baseline_time/(f+b+w)/m - 1)
baseline_acceleration=percentage(0)
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
zbv_result=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(filter(lambda x: x is not None, [baseline_time, zb_time, zbv_time]))
print(max_time)
if baseline_result is not None:
baseline_image = get_schedule_image(baseline_result, max_time)
if zb_result is not None:
zb_image = get_schedule_image(zb_result, max_time)
if zbv_result is not None:
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")
presets = {
'Ideal Case 1p': (4, 12, 20, 20, 20, 0, '1p (Same as 1F1B)'),
'Ideal Case 2p': (4, 12, 20, 20, 20, 0, '2p'),
'Real Case 1p': (4, 12, 1049, 1122, 903, 79, '1p (Same as 1F1B)'),
'Real Case 2p': (4, 12, 1049, 1122, 903, 79, '2p'),
}
preset_buttons = {}
with gr.Group():
gr.Markdown("Preset Setups")
with gr.Row():
for (k, v) in presets.items():
preset_buttons[k] = gr.Button(k, variant="secondary")
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=100, interactive=True, precision=0)
b=gr.Number(label="Time of B", value=110, interactive=True, precision=0)
w=gr.Number(label="Time of W", value=90, interactive=True, precision=0)
c=gr.Number(label="Time of one P2P communication", value=5, 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 int(p*float(s.split('p')[0]) + 0.5)
memsel=gr.Radio(choices=["1p (Same as 1F1B)", "1.5p", "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", variant="primary")
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", show_label=False)
with gr.Group():
gr.Markdown("ZB 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", show_label=False)
with gr.Group():
gr.Markdown("ZBV Schedule")
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", show_label=False)
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])
for (k, v) in presets.items():
def update_preset(pb, p, m, f, b, w, c, mem):
print(pb)
print(presets[pb])
print(presets[pb][-1])
return *presets[pb],*calculate(*presets[pb][:-1], update_mem(p, presets[pb][-1], -1))
preset_buttons[k].click(
update_preset,
inputs=[preset_buttons[k], p, m, f, b, w, c, mem],
outputs=[p, m, f, b, w, c, memsel, 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()
|