Huawei recomputation scheduler

What it is

Huawei's custom challenge at StormHacks in October 2025 was called "Automatic Decision for Re-computation in AI model optimization." Running a model's computation graph produces intermediate results that take up memory. Under a strict memory limit, you have to decide for each one: keep it, or throw it away and recompute it later when it's needed again.

Keeping everything is fast but blows the memory limit. Recomputing everything fits in memory but is slow. The job was to schedule the graph and choose what to recompute, so peak memory stays under the limit and total runtime stays low.

It's the same tradeoff as activation checkpointing in ML training and inference, and the same problem ML compilers face when fitting a model onto an accelerator. We placed 1st of 220+ and presented the solution to Huawei Research.

Why the order matters

A result has to stay in memory until every node that uses it has run. Any valid order runs the same work, but different orders keep different results alive at the same time. In a small graph where A feeds B, C and E, two valid orders give very different peaks:

Finding the best order is NP-hard in general, so exact methods don't scale to graphs with 10,000 nodes.

How we solved it

We were a team of four: Winston Thov, Dennis Kritchko, James Nguyen and me. We built a scheduler in C++ with several algorithms behind one interface, so we could compare them on the same graphs.

When memory gets tight, the scheduler evicts intermediate results aggressively and recomputes the nodes that produced them later, which cuts the memory peak. To know whether the fast solvers were any good, we wrote an MILP formulation in Gurobi that finds the optimal answer on small graphs, and checked the others against it.

Why beam search won

Beam search keeps the best few partial schedules at each step instead of just one. That made it much faster than branch-and-bound and gave it much lower peak memory than greedy.

Each candidate schedule is ranked by peak memory, then current memory, then a critical-path tie-breaker that prefers nodes with the longest path still to the output. Without it, the search tends to put off long chains, which keeps their inputs alive and blows up memory late in the schedule. With it, long chains clear early and free their results.

Every state update is incremental, and the state is compact: bitsets and flat arrays instead of node objects, with no allocation per step. That's why it's C++. The inner loop runs millions of times, and that compactness is what fits a 10,000-node graph into under 2 seconds. The beam width is the one knob: wider gives better schedules and takes proportionally longer.

Schedule beam_search(const Dag& g, size_t K, Bytes limit) {
    std::vector<State> beam{ State::initial(g) };
    for (size_t step = 0; step < g.size(); ++step) {
        std::vector<State> next;
        for (const State& s : beam) {
            for (NodeId v : s.ready()) {
                State c = s.extend(v, g);        // incremental alloc/free
                if (c.peak() <= limit) next.push_back(std::move(c));
            }
        }
        dedupe(next);                            // same scheduled set + live set
        std::partial_sort(next.begin(),
                          next.begin() + std::min(K, next.size()),
                          next.end(), by_peak_mem_then_critical_path);
        next.resize(std::min(K, next.size()));
        beam = std::move(next);
    }
    return best_of(beam);
}

Across 200 graphs, beam search had 40% lower peak memory than greedy. On the 50 graphs small enough to solve exactly, it came within 8% of Gurobi's optimum.

Results

nodes

10K+

runtime

<2s

peak memory

-40% vs greedy (200 graphs)

gap to optimal

within 8% (50 graphs)

team

4

What I'd change

Our own write-up admits the recomputation logic and fallbacks needed work, and on some cases peak memory hit the ceiling because we pushed hard for the best runtime. I'd make smarter choices about what to recompute, and make the fallbacks consistent, before squeezing out more speed.