lerobot/robot-learning-tutorial

1import plotly.graph_objects as go2import plotly.io as pio3import numpy as np4 5"""6Stacked bar chart: GPU memory breakdown vs sequence length, with menus for Model Size and Recomputation.7Responsive, no zoom/pan, clean hover; styled to match the minimal theme.8"""9 10# Axes11seq_labels = ["1024", "2048", "4096", "8192"]12seq_scale = np.array([1, 2, 4, 8], dtype=float)13 14# Components and colors (aligned with the provided example)15components = [16    ("parameters", "rgb(78, 165, 183)"),17    ("gradients",  "rgb(227, 138, 66)"),18    ("optimizer",  "rgb(232, 137, 171)"),19    ("activations", "rgb(206, 192, 250)"),20]21 22# Model sizes and base memory (GB) for params/grad/opt (constant vs seq), by size23model_sizes = ["1B", "3B", "8B", "70B", "405B"]24params_mem = {25    "1B": 4.0,26    "3B": 13.3,27    "8B": 26.0,28    "70B": 244.0,29    "405B": 1520.0,30}31# Optimizer ~= 2x params; gradients ~= params (illustrative)32 33# Activations base coefficient per size (growth ~ coeff * (seq/1024)^2)34act_coeff = {35    "1B": 3.6,36    "3B": 9.3,37    "8B": 46.2,38    "70B": 145.7,39    "405B": 1519.9,40}41 42def activations_curve(size_key: str, recompute: str) -> np.ndarray:43    base = act_coeff[size_key] * (seq_scale ** 2)44    if recompute == "selective":45        return base * 0.2546    if recompute == "full":47        return base * (1.0/16.0)48    return base49 50def stack_for(size_key: str, recompute: str):51    p = np.full_like(seq_scale, params_mem[size_key], dtype=float)52    g = np.full_like(seq_scale, params_mem[size_key], dtype=float)53    o = np.full_like(seq_scale, 2.0 * params_mem[size_key], dtype=float)54    a = activations_curve(size_key, recompute)55    return {56        "parameters": p,57        "gradients": g,58        "optimizer": o,59        "activations": a,60    }61 62# Precompute all combinations63recomp_modes = ["none", "selective", "full"]64Y = {mode: {size: stack_for(size, mode) for size in model_sizes} for mode in recomp_modes}65 66# Build traces: 4 traces per size (20 total). Start with size index 0 visible67fig = go.Figure()68for size in model_sizes:69    for comp_name, color in components:70        fig.add_bar(71            x=seq_labels,72            y=Y["none"][size][comp_name],73            name=comp_name,74            marker=dict(color=color),75            hovertemplate="Seq len=%{x}<br>Mem=%{y:.1f}GB<br>%{data.name}<extra></extra>",76            showlegend=True,77            visible=(size == model_sizes[0]),78        )79 80# Compute y-axis ranges per size and recomputation81def max_total(size: str, mode: str) -> float:82    stacks = Y[mode][size]83    totals = stacks["parameters"] + stacks["gradients"] + stacks["optimizer"] + stacks["activations"]84    return float(np.max(totals))85 86layout_y_ranges = {mode: {size: 1.05 * max_total(size, mode) for size in model_sizes} for mode in recomp_modes}87 88# Layout89fig.update_layout(90    barmode="stack",91    autosize=True,92    paper_bgcolor="rgba(0,0,0,0)",93    plot_bgcolor="rgba(0,0,0,0)",94    margin=dict(l=40, r=28, t=20, b=40),95    hovermode="x unified",96    legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0),97    xaxis=dict(title=dict(text="Sequence Length"), fixedrange=True),98    yaxis=dict(title=dict(text="Memory (GB)"), fixedrange=True),99)100 101# Updatemenus: Model Size (toggle visibility)102buttons_sizes = []103for i, size in enumerate(model_sizes):104    visible = [False] * (len(model_sizes) * len(components))105    start = i * len(components)106    for j in range(len(components)):107        visible[start + j] = True108    buttons_sizes.append(dict(109        label=size,110        method="update",111        args=[112            {"visible": visible},113            {"yaxis": {"range": [0, layout_y_ranges["none"][size]]}},114        ],115    ))116 117# Updatemenus: Recomputation (restyle y across all traces)118def y_for_mode(mode: str):119    ys = []120    for size in model_sizes:121        stacks = Y[mode][size]122        for comp_name, _ in components:123            ys.append(stacks[comp_name])124    return ys125 126buttons_recomp = []127for mode, label in [("none", "None"), ("selective", "selective"), ("full", "full")]:128    ys = y_for_mode(mode)129    # Flatten into the format expected by Plotly for multiple traces130    buttons_recomp.append(dict(131        label=label,132        method="update",133        args=[134            {"y": ys},135            {"yaxis": {"range": [0, max(layout_y_ranges[mode].values())]}},136        ],137    ))138 139fig.update_layout(140    updatemenus=[141        dict(142            type="dropdown",143            x=1.03, xanchor="left",144            y=0.60, yanchor="top",145            showactive=True,146            active=0,147            buttons=buttons_sizes,148        ),149        dict(150            type="dropdown",151            x=1.03, xanchor="left",152            y=0.40, yanchor="top",153            showactive=True,154            active=0,155            buttons=buttons_recomp,156        ),157    ],158    annotations=[159        dict(text="Model Size:", x=1.03, xanchor="left", xref="paper", y=0.60, yanchor="bottom", yref="paper", showarrow=False),160        dict(text="Recomputation:", x=1.03, xanchor="left", xref="paper", y=0.40, yanchor="bottom", yref="paper", showarrow=False),161    ],162)163 164# Write fragment165fig.write_html("../../app/src/content/fragments/bar.html",166               include_plotlyjs=False,167               full_html=False,168               config={169                   'displayModeBar': False,170                   'responsive': True,171                   'scrollZoom': False,172               })173 174