Series contents · Numerics and Quantization · 阅读中文版
“The model has one tiny scale vector and one tiny bias vector. How can it run out of memory?” The allocation failure appears to be arguing with common sense—until someone expands the compiled constants’ shapes. Each tiny vector has become a full feature map.
In mathematics, broadcasting resembles an ellipsis. In an implementation, it may mean an address rule or thousands of copies. One word should not hide that distinction.
1. Broadcast Semantics Are Cheap; Materialization May Not Be
Let the teaching input be [N,C,H,W]=[1,12,20,28]. Each per-channel affine parameter contains C elements.
Keeping scale and bias per-channel requires 2*C=24 parameter elements. Expanding both to the input shape requires 2*N*C*H*W=13440. The ratio is N×H×W. This is an element-count derivation, not actual device memory usage; physical occupancy also depends on storage width, alignment, and tiling.
Historical implementations used several strategies: reusing channel vectors, expanding along a spatial dimension, and expanding across the full feature map. Each could make the kernel contract of its time easier to satisfy while moving complexity elsewhere.
Full expansion gives the two inputs identical shapes, making them straightforward inputs to general elementwise multiplication and addition. Its costs can include larger parameter files, more transfers, and larger buffers. Compact parameters move less data but require correct broadcast addressing and a scheduler that understands their reuse.
2. You Cannot Locate C by Finding a Conveniently Equal Dimension
Channel-first [C,H,W] and channel-last [H,W,C] have different linear sequences. Expanding parameters in the former repeats one channel’s parameter over the entire spatial region before moving to the next. In the latter, every pixel repeats the channel parameters.
With three parameters [a,b,c]:
1 | channel-first: a a a ..., b b b ..., c c c ... |
If a spatial dimension happens to equal C, finding a dimension whose size matches the parameter length becomes ambiguous. Later historical expansion logic accepted more shapes and scalar parameters, but that does not eliminate the general risk of inferring semantics from sizes.
A safer principle is to preserve layout or the channel axis explicitly. If a heuristic is unavoidable, define its precedence and reject cases that cannot be identified reliably. A parameter-length check establishes compatible counts, not correct axis identity.
Teaching tests should deliberately include both C=H or C=W and shapes whose dimensions are all different. The former exposes ambiguity; the latter exposes simple axis mistakes. Choosing only tidy square shapes lets several errors cooperate.
3. Logical Element Count Is Not Physical Memory Usage
Suppose a teaching device aligns each channel vector to A bytes, with b bytes per element. One possible channel-last storage rule is:
1 | pixel_stride = align_up(C*b, A) |
For small C, padding may exceed useful data. Suppose also that C=1 permits treating W elements as one contiguous row:
1 | row_stride = align_up(W*b, A) |
The reason for these different paths depends on storage and execution rules; neither should be copied indiscriminately to all hardware. The historical dedicated affine path did distinguish single-channel and multichannel cases, using different loop strides to reduce waste in certain layouts.
Sub-byte types need special care. Rounding the byte size of one element upward before multiplying by the element count may overestimate storage. Packed storage instead requires calculating the total bit count before converting to bytes. Read, alignment, and vector-access rules may still round again at row or block boundaries. Memory estimation and actual address generation must share one contract.
4. Does the Temporary Buffer Really Need the Whole Feature Map?
A dedicated affine kernel may multiply into a temporary region, then read it to add bias. If it handles one small block at a time and completes its multiply-add immediately, that temporary region could theoretically be only one block large.
If scheduling separates multiplication and addition, or permits asynchronous execution and pipeline overlap, lifetimes must be recalculated. Adjacent lines of code do not prove that an intermediate can immediately be overwritten.
Early historical implementations declared a workspace roughly the size of the output; later paths used the output itself to hold an intermediate. This suggests a practical review question: were resource declarations updated together with actual use? If code no longer reads a workspace but the interface still reserves it, a compiler can exhaust memory on an unused reservation.
The question extends beyond affine operators. Whenever a kernel moves from multiple stages to in-place processing, revisit extra workspace, aliasing conditions, dependencies, and alignment.
5. Why Simplifying Multiplication by One Helps but Does Not Solve Everything
One historical fix for memory exhaustion lowered unit-coefficient cases to bias addition, or to requantization when the bias was also zero. This avoided some constants and intermediate nodes.
That is reasonable but limited. General coefficients still need broadcasting and workspace, and approximately unit coefficients raise the error conditions discussed in BatchNorm Quantization. A special case compiling successfully does not establish that the broadcast-memory model is solved.
Broader candidates include deferring broadcast materialization through zero strides or broadcast descriptors, loading compact parameters by tile, reusing parameters across adjacent elementwise computations, and permitting safe output-buffer reuse. All need actual support and tests; these suggestions are not presented as historically completed capabilities.
6. Observe Performance and Resources Together
A rough traffic model is:
1 | Fully expanded parameters: additional parameter traffic is proportional to N*C*H*W. |
“Ideal” matters. Reloading compact parameters for every small block increases actual traffic beyond C. Conversely, a full constant held in suitable storage and reused can amortize its cost. Loop instructions, descriptor loading, and small DMA transfer efficiency also matter.
The following tests are proposed:
| Dimension | Coverage | Observations |
|---|---|---|
| Layout | Channel-first and channel-last | Parameter correspondence and address strides |
| Channels | 1, small counts, near alignment boundaries | Padding and special paths |
| Parameters | Scalar, per-channel, mixed | Broadcast rules |
| Shape | C equal to a spatial dimension; all dimensions different | Ambiguity in channel inference |
| Resources | Before/after materialization, different tiles | Peak local memory and lifetimes |
| Data | Different parameters per channel, coordinate-encoded inputs | Constant data cannot hide misalignment |
Historical diffs establish changes in broadcast strategy and resource simplification. They do not provide a reusable table of measured peak memory. A public article can honestly explain how to construct that table and which entries require execution.
When a tiny vector fills local memory, the allocator may not be unintelligent. It may be faithfully carrying out an ellipsis expanded too early. Keeping broadcasting as a rule until copying is actually necessary can be more valuable than squeezing a few more multiplication instructions.
7. Add a “Live at the Same Time” Column to the Memory Ledger
Adding every tensor’s size usually does not give peak memory. Looking only at the largest tensor does not either. The relevant quantity is the sum of inputs, parameters, outputs, and temporaries still live at a scheduling instant, with alignment and nonreuse constraints included. Even a moderately sized expanded constant can worsen the peak if loaded early and released late.
For teaching, draw a five-column timeline: load input, load parameters, multiply, add, write output back. Mark which buffers must remain live in each column. Once multiplication finishes, can the input space hold the addition output? Not if a bypass branch still consumes the input. One extra graph consumer can turn a legal in-place optimization into an illegal one.
Broadcast parameters also differ from activations. They may be reused across tiles, making full materialization for every tile unattractive. Keeping them in local memory for a long time, however, reduces space available to activations. One alternative is compact storage at a higher memory level and loading by channel block. More small transfers and synchronization are the cost; assess them using real transfer granularity, not just total bytes.
Estimation interfaces and address generation should share alignment rules. Allocating compact byte counts while code generation accesses aligned row strides causes overruns. Conversely, counting padding twice while actual accesses are compact can reject a model that would fit. One problem appears as incorrect execution and the other as compilation failure, yet both may originate in the same missing layout contract.
Resource tests should preserve three figures: logical size, physical size, and peak live size. Logical size explains model scale; physical size explains padding; peak live size explains scheduling. “Memory went down” does not tell readers whether constants shrank, layout changed, or lifetimes shortened, making the lesson hard to transfer.
For the historical unit-coefficient simplification, the most specific validation is not merely successful compilation. Compare which constants and temporary buffers disappear, while checking that the output numerical domain remains correct. That confirms the improvement comes from the intended cause rather than an unrelated tiling change accidentally avoiding the allocation failure.
Back to series contents · Previous · Next
If you like this blog or find it useful for you, you are welcome to comment on it. You are also welcome to share this blog, so that more people can participate in it. All the images used in the blog are my original works or AI works, if you want to take it,don't hesitate. Thank you !