File size: 3,512 Bytes
1c2f6af a54df66 1c2f6af 98a83bc 1c2f6af e0e4e61 1c2f6af b22c748 1c2f6af a54df66 1c2f6af |
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 |
import gradio as gr
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import math
# plt.rcParams['axes.prop_cycle'] = plt.cycler(color=["#FF6600","#FFBE00","#496767", "#87A19E", "#FF9200", "#0F3538","#F8E08E", "#0F2021","#FAFAF0"])
def visualize_cluster(tp, pp, nodes, nodes_per_row, node_spacing=0.5, gpu_spacing=0.1 ):
gpus_per_row = 2
gpus_per_column = 4
nnodes_x = min(nodes, nodes_per_row)
nnodes_y = math.ceil(nodes/nodes_per_row)
fig, ax = plt.subplots(figsize=(2*nnodes_x, 2*nnodes_y), dpi=200)
ax.set_xlim(-node_spacing, 2*nnodes_x + nnodes_x*gpu_spacing + (nnodes_x-1)*node_spacing+ node_spacing)
ax.set_ylim(-node_spacing, 4*nnodes_y + nnodes_y*3*gpu_spacing + (nnodes_y-1)*node_spacing+ node_spacing)# gpus_per_column *(1 + 2*gpu_spacing))
ax.set_xticks([])
ax.set_yticks([])
ax.grid(False)
ax.set_aspect('equal', 'box')
ax.invert_yaxis()
model_instance = tp*pp
dp = (nodes*gpus_per_column*gpus_per_row)//model_instance
max_gpu_usage = dp*model_instance
gpu_i = 0
# Draw nodes with spacing
for node in range(nodes):
node_x = (node%nodes_per_row) * (gpus_per_row+(gpus_per_row-1)*gpu_spacing+node_spacing)
node_y = (node//nodes_per_row) * (gpus_per_column+(gpus_per_column-1)*gpu_spacing+node_spacing)
# Draw GPUs within the node with spacing (2 rows x 4 columns)
for j in range(gpus_per_column):
for i in range(gpus_per_row):
model_instance_group = gpu_i//model_instance
tp_instance_group = (gpu_i%model_instance)//tp
alpha = (1+tp_instance_group)/pp
if pp==1:
alpha=1
x = node_x + i * (1 + gpu_spacing)
y = node_y + j * (1 + gpu_spacing)
color = f'C{model_instance_group}'
if gpu_i >=max_gpu_usage:
color="black"
alpha=1
rect = patches.Rectangle((x, y), 1, 1, linewidth=1, edgecolor='black', facecolor=color, alpha=alpha)
ax.add_patch(rect)
if pp>1:
ax.annotate(f"{tp_instance_group+1}", (x+0.5, y+0.5), color='black', weight='bold', fontsize=9, ha='center', va='center')
gpu_i += 1
plt.tight_layout()
md = f"""## Resulting configuration
- **3D config: TP={tp}, PP={pp}, DP={dp}**
- **one model instance requires {model_instance} GPUs**
- **{nodes*8-max_gpu_usage} GPUs (in black) cannot be utilized**
- **numbers and color shades indicate pipeline stage if PP>1**"""
return md, fig
with gr.Blocks() as demo:
gr.Markdown("# 3D Parallelism")
gr.Markdown("## 3D Configuration")
with gr.Row():
tp = gr.Number(value=4, label="Tensor Parallel")
pp = gr.Number(value=4, label="Pipeline Parallel")
nodes = gr.Number(value=8, label="Number of Compute Nodes")
nodes_per_row = gr.Number(value=8, label="Number Nodes per Row")
button = gr.Button("Compute!")
with gr.Row():
with gr.Column():
md = gr.Markdown("## Resulting configuration:")
plot = gr.Plot(value=plt)
button.click(fn=visualize_cluster, inputs=[tp, pp, nodes, nodes_per_row], outputs=[md, plot])
demo.load(fn=visualize_cluster, inputs=[tp, pp, nodes, nodes_per_row], outputs=[md, plot])
demo.launch() |