Raghu Bharadwaj

Posted on Aug 26 Originally published at techveda.live

Why the Inference Process Keeps Growing After It Starts

#linux #embedded #machinelearning #performance

Written for the person who owns the DRAM budget, the memory cgroup and the OOM postmortem on an embedded board — not the person tuning the model.

To the kernel, an edge AI model is not layers. It is two mappings with very different properties: the weights, usually mapped file-backed and therefore reclaimable, and the tensor arena, a single anonymous mapping holding every intermediate tensor, which on a board without swap cannot be reclaimed at all. The arena is roughly four to ten and a half times smaller than the sum of the tensor sizes, it becomes resident only as its pages are written, and no allocation strategy can take it below a size fixed by one layer of the graph.

Somebody hands you a model file and asks whether it will run on the board. You are the one who has to answer, because you own the DRAM budget, the memory cgroup, and the postmortem when the OOM killer takes the inference process at three in the morning. The weights are the easy half: that number is on disk and you can read it. The hard half is the tensor arena, the block the runtime carves up for intermediate results, and nobody on the model side can tell you how large it will be. This article is about sizing that tensor arena with kernel tools rather than guesses.

What the kernel sees when a model loads

Attach to the process after it has loaded a model and look at its mappings. You will find two objects that matter — the weights and the tensor arena — and the difference between them decides how the system behaves under pressure.

The first is the model file itself. Where the platform supports it, the runtime maps the file rather than reading it into a buffer. In LiteRT this is MMAPAllocation, selected by FlatBufferModel::BuildFromFile after an IsSupported() check, with FileCopyAllocation as the fallback on platforms where mapping is not available. When mapping is used, the weights are file-backed pages. They appear in /proc/<pid>/maps with the model's pathname, they are clean, and the kernel can drop them under memory pressure and fault them back in from storage later.

The second is the tensor arena. It is one large anonymous allocation with no backing file, and it normally exceeds glibc's mmap threshold, so it is served by mmap as its own mapping rather than carved out of the heap. It appears in the maps with no pathname. Two caveats before you rely on that: glibc's threshold starts at 128 KiB but is raised dynamically, up to 32 MiB on 64-bit, once a large block has been freed, and a musl or uClibc-ng rootfs behaves differently again. Check it on your own board rather than assuming.

That difference is the whole reason a kernel engineer should care about this. Under memory pressure the kernel can reclaim clean file-backed pages for free. It cannot do that with the arena's anonymous pages: once written, they have nowhere to go but swap. Most embedded products ship without swap at all. Where zram is configured there is a swap device, so the page frame really is freed and a compressed copy is kept in the zram pool — but activation tensors are dense numeric data and compress far worse than ordinary heap, and every refault costs decompression on the inference path. So of the two halves of your model's memory, the weights are the half the kernel can drop for free, and the tensor arena is the half that stays where it is until the process exits. When you size DRAM, the arena is the number with no give in it.

The weights half is its own problem, governed by the numeric format you ship and by how the memory system moves those bytes, covered separately in 4-Bit Weight Quantization: Why the Memory System Decides. The rest of this article is about the half that cannot be reclaimed.

The number the model team gives you is wrong

Ask for the memory requirement and you will usually get the sum of the tensor sizes, because that is easy to compute. For MobileNet v2 at 32-bit float the intermediate portion of that sum is about 26 MB, roughly twice the size of the model file. Measure the process on the board and you will see nothing like 26 MB of anonymous memory. You will see something near 6 MB. Here is the calculation you were handed, in full, which runs on the host rather than the board because it needs the full framework:

import numpy as np
import tensorflow as tf

interp = tf.lite.Interpreter(model_path="model.tflite")
interp.allocate_tensors()

total = 0
for t in interp.get_tensor_details():
    shape = t["shape"]
    if shape.size == 0:
        continue
    total += int(np.prod(shape)) * np.dtype(t["dtype"]).itemsize

print("sum of all tensor bytes:", total)

Note what that loop actually sums: every tensor the interpreter reports, weights and constants included, for the primary subgraph only. It is the shape of the estimate you will be handed, not a measurement of the arena. The intermediate share of it is what the naive column of the table further down measures.

The gap is the allocation plan, and because it depends on the shape of the graph rather than on the weight count, it produces a specific failure. A team sizes DRAM from one model they measured, then ships a second model with similar accuracy and a similar file size, and the second one needs considerably more anonymous memory because its graph has a different shape. Nothing in the model file warns you, and the arena grows quietly across a project: every skip connection added to improve a metric raises the peak by an amount nobody computed.

What a tensor arena actually is

A network is a directed acyclic graph: nodes are operators, edges are tensors holding intermediate results. The runtime allocates memory for those edges before inference rather than during it, for reasons a kernel engineer will recognise immediately. Calling the allocator inside the inference loop adds time that is neither small nor predictable, and an allocation that can happen mid-inference can also fail mid-inference. So the runtime asks a different question up front: not how much memory all these tensors need, but what is the smallest single block in which all of them can live, given that they do not all need to be alive at once. That block is the tensor arena, and computing it is a scheduling problem rather than an allocation problem.

The answer can be far smaller than the sum because execution is sequential. At any moment one operator is running and only its inputs and outputs must be valid; everything produced earlier and already consumed is dead memory that can be handed to a later tensor. If the graph were a plain chain, two buffers sized for the widest tensor would do. Real graphs are not chains, and residual connections that keep a tensor alive across many operators are what make the problem hard.

The record that makes reuse possible

Every runtime that does this well builds the same small record per tensor: its size, the index of the first operator that touches it and the index of the last, where the indices come from the topological sort that is also the execution order. In LiteRT you can read it directly in tensorflow/lite/simple_memory_arena.h:

struct ArenaAllocWithUsageInterval {
  size_t offset;
  size_t size;
  int32_t tensor;
  int32_t first_node;
  int32_t last_node;
};

That is the idea in five fields. first_node and last_node define a usage interval, and two tensors whose intervals overlap can never share bytes. The planner assigns every tensor an offset into one buffer such that no two overlapping tensors overlap in memory, and makes the buffer as small as it can. The same class tracks a high_water_mark_ and ships a debug dump behind a weak symbol, so a build that links lite:simple_memory_arena_debug_dump can print arena utilisation and live tensors per operator.

One layer decides your minimum arena size

Three derived quantities explain everything downstream. The operator profile of an operator is the set of tensors alive while it runs. The operator breadth is the sum of their sizes. And for a contiguous tensor arena, the theoretical minimum arena size is the maximum operator breadth across the whole graph. It is a lower bound, not a promise: the published measurements reach it on most networks and miss it on some.

If the arena does not fit your budget, one specific layer is responsible. Not the model in general, not the weight count. Every tensor in an operator's profile must be resident while that operator runs, so no allocation strategy can produce an arena smaller than the widest operator. That changes what you say in the review meeting: you can name the layer and hand it back, and the request stops being "please make the model smaller" and becomes "this layer requires N megabytes minimum, can it be restructured".

The planning itself is cheap. These are approximation algorithms for an NP-complete problem, they run in milliseconds, and a runtime can afford to compute several plans at startup and keep the smallest. Be careful where you go looking for them, though. The MemoryStrategy enum — NAIVE, EQUALITY, GREEDY_IN_ORDER, GREEDY_BY_BREADTH, GREEDY_BY_SIZE, GREEDY_BEST and MINCOSTFLOW — belongs to the GPU delegate's planner. The CPU arena you are measuring does not consult it. SimpleMemoryArena::Allocate implements its own search for the smallest gap between already-placed allocations, ordered by ArenaPlanner, so that pair of files is what to read if you want to know what produced your arena.

Where the naive number ranks two models backwards

MobileNet v2 and DeepLab v3 differ by a factor of 1.85 in naive intermediate memory. After packing they need almost the same arena, 5.742 MB against 4.653 MB. Size a board from the naive numbers and you would provision DeepLab for more than ten times the memory it needs, and you would have the relative cost of the two models backwards. The full set is below.

Network	Naive (MB)	Packed arena (MB)	Ratio
DeepLab v3	48.642	4.653	10.5x
Inception v3	54.010	7.914	6.8x
BlazeFace	2.698	0.492	5.5x
MobileNet v2	26.313	5.742	4.6x
PoseNet	28.556	6.271	4.6x
MobileNet v1	19.248	4.594	4.2x

Intermediate tensor memory only, 32-bit float, packed greedy by size. Published measurements; see Further reading. At int8 the absolute figures fall by roughly four times, but the ratios are what the sizing argument rests on and those are a property of the graph.

The ratio column is why a naive estimate is not a conservative estimate. It is wrong by a different amount for every graph, so it does not preserve the ordering between models, and it cannot be scaled into a tensor arena figure by applying a fixed factor.

Greedy by size reaches the theoretical minimum on five of these six networks; on DeepLab v3 a strip-packing method beats it by 7.2 per cent, which is the argument for computing two plans and keeping the smaller. Tighter packing should also mean a smaller working set and better cache behaviour. The source claims up to 10 per cent better inference speed from that effect but publishes no measurement for it, so treat it as something to test on your board rather than a number to budget against.

Why the process keeps growing after it starts

Here is where a straightforward measurement misleads you, and the cause is demand paging. The tensor arena is an anonymous mapping, so its pages are not resident until they are written. The runtime reserves the address range up front, but the kernel allocates physical pages only on the first write, one minor fault at a time. A read of untouched anonymous memory maps the shared zero page and costs nothing.

The consequence is that the process footprint climbs across the first several inferences and settles only once every offset in the tensor arena has been written at least once. A model whose plan places a large tensor in a branch that only executes for certain inputs may not touch those pages for a long time. So a single-shot measurement understates the requirement, sometimes badly, and the number you budget against must be the high-water mark rather than the current value.

Two files give you that. VmHWM in the process status file under /proc is the peak resident set size the process has ever reached, and writing 5 to clear_refs in the same directory resets it, which is how you scope a peak to one run. The smaps_rollup file gives a single summed entry across all mappings, including Rss, Anonymous, AnonHugePages and the rollup-only fields Pss_Anon, Pss_File and Pss_Shmem. That anonymous-versus-file split is exactly the reclaimable-versus-not split from earlier, measured rather than assumed, and it is the single most informative thing you can read about an inference process. Pss_Shmem is worth watching too: dma-buf and most accelerator allocations land there rather than in Pss_Anon. Reading the file walks every page table in the process, so it perturbs what it measures.

Sizing the cgroup instead of guessing

Once you know what the tensor arena costs in practice, put a boundary around it rather than hoping. In cgroup v2, memory.max is the hard limit: reach it and fail to reclaim, and the OOM killer is invoked inside the cgroup. memory.high is the throttle: exceed it and the processes are throttled and put under heavy reclaim pressure, and the documentation is explicit that going over it never invokes the OOM killer.

For an inference workload on an embedded board that distinction needs care rather than a default, and the reason is specific. On a swapless board the kernel does not scan the anonymous lists at all: reclaim is steered onto the file lists instead. So exceeding memory.high does not shrink the arena by a single byte. What it does is evict the other half of your model — the mapped weights — and force them back in from storage, while the task is separately throttled by a penalty sleep. You pay in refault latency and scheduling delay and get nothing back.

The practical arrangement follows from that. Use memory.max as the real boundary. Put memory.high a little below it as a monitoring tripwire rather than as a control, so the counter moves before anything is killed. And set memory.min at roughly the mapped weight size, so the reclaim that does happen cannot evict the model out from under the process.

For measurement, memory.peak records the maximum usage for the cgroup and its descendants since the cgroup was created or since the most recent reset. The reset semantics have a trap in them: writing any non-empty string resets the value only for subsequent reads through that same file descriptor. A shell sequence that writes with tee and then reads with cat opens a second descriptor and quietly reports the never-reset watermark instead. Either hold one descriptor open across the run, or use a fresh cgroup per run and skip the reset. And memory.events carries the counters that tell you afterwards what happened, including high, max, oom and oom_kill; memory.events.local is the non-hierarchical version, which is the one you want for a single-cgroup harness. A rising high count with no oom_kill is a system quietly paying for reclaim, which is the state most teams never notice.

That state has its own instrument, and it is the one to reach for. Pressure stall information reports how much work was lost to waiting on memory: system-wide in /proc/pressure/memory, and per cgroup in memory.pressure. Both carry two lines. The some line is the share of time in which at least one task was stalled on memory. The full line is the share in which every non-idle task was stalled at once, which the kernel documentation describes as thrashing, and which on an inference board means the CPU is not doing your work at all. Each line gives ten, sixty and three hundred second averages as percentages, plus a cumulative total in microseconds that catches short spikes the averages flatten out.

root@rock-5b:~# cat /sys/fs/cgroup/infer/memory.pressure
root@rock-5b:~# cat /proc/pressure/memory

For a latency-sensitive product the number to alarm on is full avg10 for the inference cgroup. Anything sustained there means the board is spending real time refaulting rather than inferring, and given what the previous paragraphs established, what it is refaulting is your model weights rather than the tensor arena. The system-wide file also accepts a trigger, so a supervisor can wait on a threshold rather than sample: writing some 150000 1000000 to it asks to be woken when partial memory stall passes 150 ms within any one-second window, with the file descriptor polled for POLLPRI. Accepted windows run from 500 ms to 10 s. Pressure stall information has to be built into the kernel, so confirm the files exist on your board before designing a supervisor around them.

Why the arena cannot be handed to an accelerator

One more property matters if the model runs on anything other than the CPU. The tensor arena is contiguous in virtual address space, not in physical memory. Its pages are ordinary anonymous pages scattered across whatever the page allocator had, and nothing about a tensor arena guarantees the physical adjacency a device may need.

That is fine for a CPU, and fine for a device behind an IOMMU or one that can scatter-gather from a descriptor list. It is not fine for a master that cannot do either, which needs physically contiguous memory from CMA or a DMA heap. Even an accelerator that could technically reach the arena usually wants its own buffers anyway, for cache-coherency and vendor layout reasons. This is why a hardware delegate allocates separately rather than being pointed at the arena, and why the arena figure and the process footprint diverge once you enable an NPU or GPU path. The arena tells you what the CPU-side plan costs. It does not tell you what the accelerator's buffers cost, and those come from a different allocator with tighter constraints. I have not verified how any specific vendor delegate accounts for its own memory, so measure that path separately rather than reasoning from the arena number.

Where this stops working

The tensor arena scheme rests on one assumption: that every intermediate tensor's size is known before the first inference. The clearest break is dynamic shapes. A graph where a tensor's extent depends on the input cannot be planned in one pass; the workaround is to plan repeatedly, allocating what is known, executing until the first dynamic size resolves and planning again. That works, but it puts allocation back inside the inference path along with the latency variance that pre-allocation existed to remove. If you are consulted on architecture for a latency-sensitive product, this is a real reason to argue for static shapes.

Worse, the arena does not shrink again. ResizableAlignedBuffer::Resize skips reallocation when the new size is smaller, so once a dynamic-shape run has grown the arena it stays grown for the life of the process. One unusually large input sets your footprint for the next three months of uptime.

Two further limits. The execution order is taken as given, so the planner optimises for a fixed topological sort and does not reorder operators to reduce the peak — offline tools exist that do reorder, which is worth knowing before you accept the number as final. And none of this applies to training, where activations must stay alive for the backward pass, so any on-device fine-tuning you are asked to support has a memory profile unrelated to the inference numbers you measured.

Measuring a tensor arena on your own board

The board below runs a Yocto image built with MACHINE = "rock-5b". Put the benchmark tool in the image rather than copying a binary onto a running board, so that what you measure is the rootfs you ship. The community meta-tensorflow-lite layer carries a recipe for it:

raghu@techveda.org:~$ bitbake-layers add-layer ../meta-tensorflow-lite
raghu@techveda.org:~$ echo 'IMAGE_INSTALL:append = " tensorflow-lite-benchmark"' >> conf/local.conf
raghu@techveda.org:~$ bitbake core-image-base

Everything from here runs on the board. Start on the kernel side, because those numbers are the ones you will defend in a design review. The process must still be alive when you read /proc, so run it long and in the background rather than sequentially:

root@rock-5b:~# cd /usr/share/tensorflow/lite/tools/benchmark
root@rock-5b:~# ./benchmark_model --graph=model.tflite --num_runs=2000 & BM=$!
root@rock-5b:~# sleep 5
root@rock-5b:~# grep VmHWM /proc/$BM/status
root@rock-5b:~# grep -E "^(Rss|Pss_Anon|Pss_File|Pss_Shmem|Anonymous|AnonHugePages|Swap):" /proc/$BM/smaps_rollup

Confirm the arena really is one anonymous mapping by listing the mappings with no pathname, largest first. Selecting on field count rather than on the absence of a slash is what keeps the heap, the stack and the vDSO out of the list, since those carry a bracketed name in the sixth field:

root@rock-5b:~# awk 'NF==5 {split($1,a,"-"); print ((strtonum("0x" a[2]) - strtonum("0x" a[1]))/1048576) " MB " $1}' /proc/$BM/maps | sort -rn | head

Watch the first-write behaviour directly. The minor fault count bounds the arena from above rather than isolating it, because it also counts the binary, the libraries and any delegate buffers. Disable the warm-up run so only one inference is measured, and diff a dry run against a real one if you want the arena's own share:

root@rock-5b:~# perf stat -e minor-faults,major-faults ./benchmark_model --graph=model.tflite --num_runs=1 --warmup_runs=0
root@rock-5b:~# perf stat -e minor-faults,major-faults ./benchmark_model --graph=model.tflite --dry_run=true

Bound it with a cgroup. The memory controller has to be enabled in the parent before the interface files exist in the child, and the workload has to actually be moved into the cgroup, or you will read zeroes from a well-formed but empty directory:

root@rock-5b:~# echo +memory > /sys/fs/cgroup/cgroup.subtree_control
root@rock-5b:~# mkdir -p /sys/fs/cgroup/infer
root@rock-5b:~# echo 64M > /sys/fs/cgroup/infer/memory.max
root@rock-5b:~# sh -c 'echo $$ > /sys/fs/cgroup/infer/cgroup.procs; exec ./benchmark_model --graph=model.tflite --num_runs=2000'
root@rock-5b:~# cat /sys/fs/cgroup/infer/memory.peak /sys/fs/cgroup/infer/memory.events.local

A fresh cgroup per run is what makes that final read trustworthy. Do not reach for the reset unless you are driving it from a program that can hold the descriptor open, for the reason given earlier.

Then go to the runtime for the plan itself. This flag prints the interpreter's internal state before the first inference, including the allocated size of each tensor, which is the plan rather than a guess about it. The second samples the tool's own footprint at a fixed interval, 50 ms by default; the sampling perturbs latency, so do not read latency and peak memory from the same run:

root@rock-5b:~# ./benchmark_model --graph=model.tflite --print_preinvoke_state=true --num_runs=1
root@rock-5b:~# ./benchmark_model --graph=model.tflite --report_peak_memory_footprint=true --memory_footprint_check_interval_ms=20

If the model has dynamic-shaped tensors, use the post-invoke state rather than the pre-invoke state, because the sizes are not known until the graph has run. Two flags then change the trade-off, but they are not independent: enabling the second also enables the first, so measure baseline, then the first alone, then both:

root@rock-5b:~# ./benchmark_model --graph=model.tflite --print_postinvoke_state=true
root@rock-5b:~# ./benchmark_model --graph=model.tflite --release_dynamic_tensors=true
root@rock-5b:~# ./benchmark_model --graph=model.tflite --optimize_memory_for_large_tensors=1

Finally, back on the host, read the planner your build actually uses rather than trusting an article about it. For a CPU inference process that means the arena and its planner, not the GPU delegate's strategy enum:

raghu@techveda.org:~$ grep -nE "smallest gap|best_offset" tensorflow/lite/simple_memory_arena.cc
raghu@techveda.org:~$ grep -rn "MemoryStrategy" tensorflow/lite/simple_memory_arena.cc

The second command returns nothing, and that is the point. The CPU arena never consults MemoryStrategy.

Key takeaways

Budget DRAM off the tensor arena figure, not the weights figure. The weights are the half the kernel can drop for free.
A naive estimate is not a conservative one. It is wrong by a different factor for every graph, so it does not even preserve the ordering between two models.
Name the layer instead of asking for a smaller model. The maximum operator breadth is the bound no allocation strategy can beat.
Measure VmHWM over a long run, never a single inference, and reset it with clear_refs to scope a peak to one run.
Set memory.max as the boundary, memory.high below it only as a tripwire, and memory.min at the weight size so reclaim cannot evict the model.
Alarm on full avg10 in the cgroup's memory.pressure. It is the one number that tells you the board is refaulting rather than inferring.
Measure the accelerator path separately. The arena is virtually contiguous only, and delegate buffers come from a different allocator.
Dynamic shapes, training and delegate buffers all fall outside the plan, and an arena grown by one large input never shrinks back.

Frequently asked questions

Why does my inference process keep growing for the first few seconds?

Because the tensor arena is an anonymous mapping whose pages become resident only when first written. The address range is reserved up front, but physical pages arrive one minor fault at a time, so the resident set climbs until every offset in the arena has been touched.

Which part of an edge model can the kernel reclaim under memory pressure?

Only the weights, and only when the runtime maps the model file rather than copying it. Those pages are clean and file-backed, so the kernel can drop them and read them back later. The arena is anonymous, so without a swap device it cannot be reclaimed at all.

Which layer decides how big my tensor arena has to be?

The operator with the largest breadth, meaning the largest total size of all tensors that must be resident while it runs. That is a lower bound no allocation strategy can beat, though a given strategy may not reach it exactly.

Should I use memory.high or memory.max for an inference cgroup?

Use memory.max as the real boundary and memory.high below it only as a tripwire. On a swapless board the kernel steers reclaim entirely onto the file lists, so exceeding memory.high evicts the mapped weights and throttles the task without shrinking the arena at all.

How do I tell whether the board is losing time to memory pressure?

Read memory.pressure in the inference cgroup, or /proc/pressure/memory system-wide. The full line is the share of time in which every non-idle task was stalled on memory, so a sustained full avg10 means the board is refaulting rather than inferring.

Can an NPU or GPU use the tensor arena directly?

Not in general. The arena is contiguous in virtual address space but its physical pages are scattered, so a DMA master without an IOMMU cannot use it. Delegates allocate their own buffers through a different path, which is why the arena figure and the process footprint diverge once acceleration is enabled.

DEV Community