Create app.py

app.py

@@ -0,0 +1,88 @@
+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(nodes, tp, pp, node_spacing=0.5, gpu_spacing=0.1, nodes_per_row=8):
+    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\
+
+TP={tp}, PP={pp}, DP={dp}, where one model instance requires {model_instance} GPUs and {nodes*8-max_gpu_usage} GPUs (in black) cannot be utilized in this configuration. 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=2, 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()