Memory Budget Allocator

A stacked-bar visualisation of the four cohabiting memory buckets that fight for a single GPU's HBM during training: weights, gradients, optimizer state (Adam's m + v moments), and activations. Users drag four sliders — model size, batch size, sequence length, precision — and read off whether the configuration fits in a fixed capacity (default 80 GiB, an A100).

Installation

npx shadcn@latest add https://craftbits.dev/r/memory-budget-allocator.json

Usage

import { MemoryBudgetAllocator } from "@craft-bits/core";
 
<MemoryBudgetAllocator
  defaultParams={{
    modelSizeBn: 7,
    batchSize: 1,
    seqLen: 2048,
    precision: "bf16",
  }}
  capacityGb={80}
/>

Drive the parameters from outside:

<MemoryBudgetAllocator
  params={params}
  onParamsChange={setParams}
  capacityGb={80}
/>

Hide the per-bucket breakdown when you want a single total bar:

<MemoryBudgetAllocator showBreakdown={false} />

Understanding the component

1. One formula per bucket.
   • weights = modelSizeBn * 1e9 * bytesPerParam
   • grads = weights
   • optimizer = 2 * weights (Adam's m + v moments, simplified)
   • activations = batchSize * seqLen * hidden * layers * bytesPerParam with hidden = 4096, layers = 32 (a 7B-shaped reference).
2. Stacked bar with capacity tick. The bar normalises against max(capacityGb, totalGb) so the capacity hairline stays anchored at the right edge until the total exceeds the budget. Anything past the tick paints in cb-error.
3. Verdict line. Below the bar, the caption switches between {headroom} GiB headroom and exceeds by {overflow} GiB so the pass/fail state is stated textually too — colour is never the only signal.
4. Breakdown legend. Each of the four slices has a tinted swatch, a textual label, and a GiB readout. Tints map to the same semantic tokens as elsewhere in the library (cb-accent / cb-info / cb-warning / cb-success).
5. Controlled or uncontrolled. params follows the Radix pattern — pass params plus onParamsChange for controlled, or defaultParams for uncontrolled. The precision pills are a role="radiogroup" of role="radio" buttons.
6. SPRINGS.smooth everywhere. Bar-width changes animate with the canonical smooth spring; prefers-reduced-motion: reduce collapses every spring to an instant swap.

Props

Accessibility

• The figure is role="figure" with a hidden summary listing model size, batch, sequence, precision, total, capacity, and the verdict — screen readers hear the story whenever props change.
• The precision picker is a role="radiogroup" of role="radio" pills with aria-checked. Tab focuses the group; Space and Enter commit a selection.
• Sliders are native <input type="range"> via the library's LabeledSlider with aria-valuemin / aria-valuemax / aria-valuenow / aria-valuetext — full keyboard + screen-reader semantics for free.
• Colour is never the only signal — every bucket has a textual label and a GiB readout, and the fits / exceeds verdict is stated as words too.
• Motion respects prefers-reduced-motion: reduce.

Credits

• Extracted from: craftingattention (app/src/lessons/primitives/nn/MemoryBudgetAllocator.tsx). Stripped the inference-focused stacked bar (weights / KV cache / activations), the phase-based narration state machine (observe / optimizing / insight), the GPU-capacity threshold markers, the breathing-pulse hint, the throughput readout, and the discrete-step sliders. The library extract is the pure plotting primitive: training-time memory (weights + grads + optimizer + activations) against a single configurable capacity line, driven entirely by props with motion-library springs and a controlled / uncontrolled params bag.