Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Performance and memory tuning

Tune a measured model, not an abstract solver. Correctness comes first: choose a solver whose assumptions represent the economics, run the model at log_level="debug", and compare a reduced problem with grid search before optimizing it.

The performance workflow has four measurements:

  1. cold environment and cold compilation cache;

  2. cold model compilation in an installed environment;

  3. warm execution in the same process;

  4. peak host and device memory.

Record model size, precision, device, solver configuration, and whether the compilation cache was warm. Without those fields, two timings are not comparable.

First locate the limiting axis

List the sizes of:

The scaling discussion explains how those axes enter each solver family. The largest declared grid is not necessarily the largest intermediate: a product or envelope matrix can dominate.

Reduce discretization only when accuracy permits

Fewer nodes reduce work and memory, but the right grid is an economic approximation choice. Inspect value and policy changes as grids are refined. Put resolution near curvature, boundaries, or regions visited frequently in simulation.

PiecewiseLinSpacedGrid and PiecewiseLogSpacedGrid control density around known locations. They do not declare a budget kink or cliff to NBEGM; use the structured budget declarations for that.

Fix a planner axis width

A solver declares its execution axes by name. ExecutionConfig(axis_widths=...) fixes the width used by each applicable compiled program or host-dispatch loop:

from lcm import ExecutionConfig, Model

model = Model(
    regimes=regimes,
    ages=ages,
    regime_id_class=RegimeId,
    edges=edges,
    initial_nodes=initial_nodes,
    execution_config=ExecutionConfig(axis_widths={"action_product": 8}),
)

A width smaller than the axis bounds how many entries are processed together; a width at or above the extent selects the whole axis in one chunk. Compiled axes vectorize that chunk; a host-dispatch axis bounds pending node calls. Lower values can reduce live intermediates at the cost of more sequential execution. They do not cap surrounding arrays, retained candidate banks, compilation memory, or total device memory.

Which axis names exist depends on which solver a regime uses; each solver’s section in Solvers and capabilities names the axes it declares.

Fix it in one regime only

A bare integer applies to every regime declaring the axis, which is the wrong instrument when two regimes have opposite shapes: a width that keeps a 16-cell, 2048-action regime inside memory takes a 4096-cell, two-action regime from eight chunks to 512. Give the axis a mapping from regime name to width instead, and only the regimes it names are fixed:

ExecutionConfig(axis_widths={"cell": {"heavy": 8}})

Regimes the mapping does not name keep the width the planner chooses for them, per regime and per period. The two forms may be mixed across axes. Because only the solve phase plans per regime, an axis that only simulation programs declare — subject, for one — takes the bare-integer form; a per-regime width for it is refused at model build, as is one naming a regime the model does not declare.

axisdeclared by
action_productthe flattened Cartesian action product of a streamed GridSearch core
stochastic_nodechild stochastic nodes folded by DCEGM, ride-along NBEGM, and their nested solvers
cellindependent state cells tiled by GridSearch, DCEGM, ride-along NBEGM, and their nested solvers
intervalliquid intervals folded by ride-along NBEGM, including eligible NNBEGM inner programs
branchdiscrete-action subproblems mapped by ride-along NBEGM, including eligible NNBEGM inner programs
savings_pointexogenous savings nodes tiled by DCEGM, including NEGM inner programs
euler_pointEuler nodes tiled by DCEGM, including NEGM inner programs
envelope_cellindependent envelope cells tiled by DCEGM, including NEGM inner programs
outer_candidatethe exogenous outer post-decision nodes of a nested outer search

These are solver capability names; a particular program declares only its applicable axes. Plain EGM declares no execution-width axis. Nested solvers inherit the relevant inner axes; their outer search determines how outer_candidate is used. See the generated capability table in Solvers and capabilities.

Choose the largest width that meets the measured memory target, then verify values and runtime against the whole-axis setting on the model and backend you will use. Leaving an axis out of the mapping lets the planner choose, which is what a device-memory budget asks it to do.

Choose execution widths independently of grids

Grids specify economic support and interpolation. Set execution widths on the model with ExecutionConfig(axis_widths=...). The effect depends on retained arrays and downstream folds; for example, outer_candidate can lower a nested solver’s temporary evaluation memory without capping its retained candidate bank.

For GridSearch, cell tiles the flattened product of the states evaluated inside each sharded slice. Sharded states stay outside this product. action_product separately streams eligible hard-max action reductions; EV1 and collective models keep their canonical dense action reductions while tiling state cells. DCEGM uses the applicable axes listed above.

NBEGM also owns no compiled-width fields. Use its applicable axes in the table above: interval streams the continuation read together with the stable-identity candidate fold, branch maps the discrete subproblems before their maximum, and cell tiles independent ride cells inside each co-mapped carry slice. A route declares only the nontrivial meshes it consumes. The interval stream has no separate segment-width loop. Smaller widths can reduce live intermediates at the cost of sequential execution; they do not cap surrounding arrays, retained candidate banks, compilation memory, or total device memory.

Exact solver fields are in Solvers and capabilities, Upper envelopes, and Outer search.

Factor a fixed component out of a Markov state

A product-coded Markov state often combines a component that never changes, such as a fixed effect or a permanent type, with components that do. With a plain StochasticTransition, the law keeps the fixed part by giving zero probability to every code outside the current group. Declaring this structure lets continuation lookup select the current group and makes the fixed component available for sharding.

Declare the group of each code with fixed_component:

from lcm import StochasticTransition

# Codes 0-3 = 2 * kind + health; `kind` never changes.
state_transitions = {
    "kind_health": StochasticTransition(
        func=next_kind_health, fixed_component=(0, 0, 1, 1)
    ),
}

fixed_component[code] is the group of that code. The law must give probability zero to every target outside the current code’s group, and every group must have the same number of codes. At model construction the state is rewritten into:

Continuation lookup selects the current group by index. For a linear expectation, the original probability vector and lottery-slot order are retained, including zero-probability slots; the reduction is not shortened. This preserves the original reduction layout while keeping value reads local to the current group. Published values can still differ in floating-point rounding between the annotated and original models.

Use it when a Markov state’s transition matrix is block-diagonal across a component, that is, when most entries of each row are structurally zero. A state without such a component gains nothing, and a model that does not use the annotation is unaffected.

Because <s>_fixed is a model-level state, it can be named in ExecutionConfig(sharded_states=("kind_health_fixed",)), which spreads the groups across devices.

Model.simulate takes initial conditions in the declared state’s codes: pass a kind_health column in a DataFrame or a "kind_health" entry in a mapping, and it is split into the two parts. The public validation and feasibility methods accept the same original observations. Labels are interpreted only on rows whose initial regime carries the state; irrelevant cells may remain blank.

The annotation works in model-level, regime-level, per-target and Phased laws. Every carrier, including a terminal regime, uses the same code grouping. Each outgoing law must declare that grouping or use fixed_transition; inconsistent groupings and reset/entry laws without an established group are rejected. Generated _rest and _fixed names must be unused by user declarations.

Solve one invariant code at a time

A state that never changes, such as a preference type with fixed_transition, splits the solve into independent problems, one per code. Name it in ExecutionConfig(invariant_block_widths={"pref_type": 1}) to solve every non-terminal regime carrying it one code at a time:

from lcm import ExecutionConfig, Model

model = Model(
    ...,
    execution_config=ExecutionConfig(invariant_block_widths={"pref_type": 1}),
)

Each code is its own dispatch of one shared executable; the code is a runtime operand, so the number of compiled programs does not grow with the number of codes. A continuation that carries the state is read through that code’s block only, so a continuation replicated onto other devices, for example under a sharded continuous state, moves one code’s share at a time. The blocked solve evaluates the same Bellman problem as the unblocked one and publishes values in the same layout. Independently compiled blocked and unblocked values agree within eight ULP of the working float format; policies, states, actions, regimes, and public result structure remain exact. Routes dispatching the same compiled program retain their bitwise contract. Terminal regimes are solved whole.

The request is checked at model construction. It is refused, naming every reason, when:

A blocked regime-period cannot be captured or replayed with the period-replay tools.

Grouped simulation

When the forward phase provably never changes the state either, simulate groups the subjects by their starting code. Each code’s subjects run in chunks of their own, and every stored value carrying the state is read through that code’s block, so a device receives one code’s share of each such value rather than all of them. Each subject keeps the random draws of its original row and its original row in the result, so the panel is the ungrouped panel. A subject that starts without the state, in a regime that does not carry it, joins the first code’s group; an empty group dispatches nothing.

Simulation falls back to the ungrouped route when the simulate phase can change the state, for example through a Phased law whose simulate branch moves it, or when a regime has taste shocks, has gated edges (a Gate in its Transition in Model(edges=...)), or replays a stored policy. The plan summary’s subject_grouping names the grouping state, or is None on the ungrouped route.

Solve, simulate and release one code at a time

By default the blocked solve is period-major: each period solves every code before the next period starts, and the device keeps every period’s complete value until solve returns. When those retained values, rather than one period’s working set, are what fills the device, solve each code through its whole lifetime instead:

from lcm import ExecutionConfig, InvariantBlockSchedule, Model

model = Model(
    ...,
    execution_config=ExecutionConfig(
        invariant_block_widths={"pref_type": 1},
        invariant_block_schedule=InvariantBlockSchedule.BLOCK_MAJOR,
        device_memory_bytes=None,
    ),
)

Each code is solved through every period on the same compiled programs, then copied to the host, and its device buffers are deleted before the next code starts, so the device holds one code’s values at a time. simulate without a solution simulates each code’s subjects right after that code is solved, while its values are still on the device.

The result keeps its usual contract. solution.values lists every period and regime, and inspecting it reads nothing. Reading a value assembles it from the retained codes on the layout the period-major solve publishes, and SolutionResult.save writes every value from the host. Simulating the result reads one code’s values at a time. Values and simulated panels are those of the period-major schedule, byte for byte; where blocking itself places a value on a neighbouring float of the unblocked one, so does this schedule.

What moves where:

The schedule is refused at model construction unless a state is named in invariant_block_widths and every regime carries it on a discrete grid. Simulating a block-major solution is refused when the simulate phase does not group subjects by the state, when the simulation is budgeted (pass device_memory_bytes=None), or when log_path asks for snapshots that would hold every value at once. A solution loaded from an archive is read whole, as on any other model.

Share a large action product over devices

A regime with a small state grid and a large action product leaves most devices idle when only its states are spread. Name it in ExecutionConfig(action_partitions=...) to let several devices share its actions:

from lcm import ExecutionConfig, Model

model = Model(
    ...,
    execution_config=ExecutionConfig(action_partitions={"working": 4}),
)

The canonical action product is cut into contiguous runs of whole action blocks, one run per device, with block counts that differ by at most one. Each device evaluates its own run with the unchanged Q and feasibility functions and keeps one exact hard-max accumulator per state cell: the best value, the winning action’s canonical identity and whether any action was feasible. The devices exchange only these accumulators and every device merges them in device order. The published value is the same exact hard maximum the ordinary route computes:

At the same action block width the published values equal the ordinary route bit for bit. The route caps the action block at n_actions // count identities, so that every device owns at least one block; the ordinary route may pick a wider block, and a different width compiles a differently vectorized Q, so compare across routes at an action width fixed with axis_widths when you need bit equality.

The regime’s devices form a mesh of its sharded-state axes times the count, so the request composes with a sharded continuous state (states times actions) and with type-local blocks:

ExecutionConfig(
    devices=(0, 1, 2, 3),
    sharded_states=("assets",),
    action_partitions={"working": 2},  # assets over 2 devices, actions over 2
)

The count is never reduced: the state axis takes the largest shard count the remaining devices allow. Values stay sharded along the states and replicated along the action devices. The exchanged accumulators are temporaries of the compiled program, so the memory admission that charges each program’s compiler reservation covers them. A count of one, or an empty mapping, runs the ordinary route unchanged.

The request is checked at model construction. It is refused, naming every reason, when:

Distribute state work

Declare a discrete state at model level, then name it in ExecutionConfig(sharded_states=("preference",)) to shard its axis. Select devices with ExecutionConfig(devices=(0, 1, 2, 3), ...); omitting devices uses all devices visible to JAX. Solver-specific restrictions also apply; see Solvers and capabilities.

Sharding is resolved per regime, from the state each regime actually carries. A regime whose DAG reads the state carries its grid axis and runs on the submesh that axis defines. A regime that never reads it --- so DAG pruning drops it there --- publishes a value without that axis and runs on a single device; the planner moves the values that cross between the two placements. That holds for a non-terminal regime as much as for a terminal one, so a shock read through working life and dropped in retirement shards the working-life regimes alone. Only a state every regime prunes is refused, because then no axis is left to spread. The stricter continuous route additionally requires the state in every regime.

Ordinary singleton hard-max GridSearch also supports one model-level LinSpacedGrid as the sole sharded state. Every regime must retain that same static grid. Unsharded states may include:

Runtime state grids, other process families, mixed solvers, collective or gated routes, same-period references and taste shocks remain outside this route. Runtime action grids keep their ordinary GridSearch semantics. Process nodes retain the ordinary index-based continuation read; sharding does not enable process-aware off-grid interpolation. Additional continuous coordinates keep their named positions, so the sharded axis need not be trailing.

For example, a 24-point assets axis can use eight selected devices, with three assets coordinates on each device. Logical value-axis order stays unchanged. Interpolation still needs the complete target value: the planner explicitly replicates each needed continuation on its consumer mesh before the numerical core runs. Retained shards, replicas and transfer workspace all count toward admission. Budgeted continuous sharding currently requires every regime and transfer to use the complete selected device set. Transfer workspace is reserved conservatively across the period, even when individual copies finish earlier. These copies can limit capacity and add communication time; eight shards do not imply an eightfold memory reduction or speedup. Reducing a cell width cannot remove the full-continuation storage requirement.

If a shard remains too large, reduce an applicable execution width. For example, ExecutionConfig(sharded_states=("preference",), axis_widths={"cell": 32}) keeps the preference axis sharded and tiles the remaining GridSearch state product.

Before solving, verify the resources actually visible to JAX:

import jax

assert jax.device_count() == expected_devices

A larger GPU can run larger chunks and may benefit from more concurrent independent work. It does not automatically shorten a workload made of small sequential kernels. Measure occupancy and memory rather than extrapolating from device memory alone.

Under a device-memory budget, solve execution waits for earlier compiled work that uses any of the next core’s execution or transfer devices. Work on disjoint device sets can remain asynchronous. A shared source array can therefore serialize cores assigned to different regime submeshes. The wait covers returned auxiliary arrays and declared copies as well as values; it does not bound memory used by compiler autotuning.

Set a device-memory budget

ExecutionConfig(device_memory_bytes=...) sets a per-device ceiling for the compiler’s reservation plus the accounted live residency, and turns on budget admission: a plan that does not fit is refused with ExecutionPlanningError rather than left to the allocator. It takes one of three values:

A few routes cannot be budgeted: enable_jit=False, forward simulation through host gated or replay adapters, and simulating from a supplied foreign result that carries replay payloads. They refuse any budget, the default included, with an error that names the budget’s source and the remedies. Pass device_memory_bytes=None to run them unbudgeted.

You may pass a device’s whole allocator pool limit. The model does not plan against all of it: device_memory_headroom_fraction (default 0.15) is the share of each selected device’s pool kept outside the ceiling, an operational margin for storage the represented accounting does not see --- collective-communication buffers, library workspaces such as cuBLAS, the driver context, allocator fragmentation. The effective ceiling is the smaller of your request and every selected device’s pool limit less that fraction. The fraction is a policy awaiting workload validation, not a measured requirement: if you have measured your own envelope, set it to 0.0 and plan against the whole pool. It must be an exact float in [0, 1), and it moves only the ceiling --- every reservation and residency figure compared against that ceiling is unchanged.

execution_config = ExecutionConfig(
    devices=(0, 1),
    device_memory_bytes=40 * 2**30,
    device_memory_headroom_fraction=0.0,  # Only with a measured envelope.
)

Resolving the budget is logged, so you can see which of the two bounds actually bound:

Every admission refusal names the ceiling it was measured against and where it came from: the request, the request capped by the device headroom (with the effective bytes, the requested bytes and the fraction that separates them), or the device default. It then lists what you can change to fit:

Exhaustive width selection across periods

Under a budget, exhaustive search checks each core’s ranked width frontier from the widest candidate and selects the first admitted rank. Admission uses the current period’s residency and each required executable variant’s compiler reservation. Memory need not be monotone in frontier rank: a refused neighbour does not prove that all earlier candidates also refuse.

Batch forward simulation

Set ExecutionConfig(axis_widths={"subject": k}) on the model to process subjects in chunks. The inner compiled subject width remains fixed at k; the outer chunk is clamped to the population and rounded up to a multiple of the subject-device count. For example, width three uses outer chunks of four on four devices while retaining an inner width of three. Completed chunks are offloaded to host. Random keys retain their original population and global subject indices, so changing the width does not change simulated draws.

Under a device-memory budget, every call is planned by the same top-first outer-cohort search, whether or not axis_widths["subject"] is pinned. Replay routes that cannot be budgeted refuse the budget, whether it is the device default or explicit.

A pinned axis_widths["subject"] fixes the anchor width the frontier doubles from. Without a pin, the existing single-axis workspace search proposes the widest representable subject width the budget admits as that anchor; the frontier and its map selection then proceed identically either way. The anchor only seeds the search — it is not a memory verdict, and the admitted extent may use a different inner map than the anchor if a fallback runs.

The planner constructs device-aligned outer sizes by repeatedly doubling the anchor width, clamping to the population, and suppressing duplicates. The largest size covers the population. It derives the preferred full inner-width map from the anchor’s axis declarations, without compiling or admitting the anchor size first. It then profiles the largest size. If that complete profile fits, it returns immediately: one whole-chunk profile, no smaller-shape sweep, and no predicted memory verdict.

After a largest-size refusal, it selects the original full/bootstrap map at the anchor size. The full map is tried first; a distinct conservative bootstrap map is tried only if the full anchor fails. That complete admitted map then stays frozen while the remaining outer sizes are tried largest first. The first fitting descending candidate wins; otherwise the admitted anchor is retained. An already rejected extent/map pair is not repeated. If fallback changes the map, the largest size must receive its own new profile. If neither anchor map fits, the search refuses; unsearched configurations may still fit. There is no outer-by-inner width sweep or assumption of monotone memory use.

For K distinct sizes, a fitting preferred full-population candidate uses one complete profile. If the full map remains selected, the worst case is K profiles. A distinct bootstrap fallback uses at most K + 2 profiles (two for a singleton frontier). This trade-off favors populations that fit whole; models fitting only small cohorts can pay more planning cost than an ascending search. Compiler exceptions are not caught as memory refusals. Compilation and its host-memory cost are not bounded by the simulated payload’s device budget.

For example, 226,848 subjects, inner width 2,048 and eight subject devices give the geometric sizes 2,048, 4,096, 8,192, 16,384, 32,768, 65,536, 131,072 and 226,848. The last size is profiled first. Success selects one cohort immediately, with the inner map still fixed. Every attempted size uses its own compiled requirements, retained owners, padding, and output/assembly reservations. Different outer shapes are not interchangeable executables. Selection is largest-first feasibility, not a measured speed optimum.

Receipts use frontier_version=3. Their candidates remain in ascending geometric order; attempts record actual profile order, including refusals and map fallback. profile_count counts complete profiles, not backend compilations. Current inputs and resident owners are checked on each call, and live admission is still checked before chunk execution.

The bound includes the retained solution, full-population inputs and RNG workspace, pending period owners, published results, padding, and final assembly. It is conservative about future aliases; it does not estimate an allocator optimum. On CPU, retaining all chunks and assembling the final result can impose a floor that narrower chunks cannot remove. CPU assembly after GPU offload uses host RAM outside the GPU ceiling.

Finite NNBEGM replay includes candidate preparation, diagnostics and canonical ranking in this bound. The prepared bank has one row per subject in the outer chunk and keeps the full published candidate extent. A smaller inner width does not remove that bank’s storage requirement. Addressed policy copies and their transfer overlap are counted alongside retained originals.

Parameterized grids and user entry laws can still allocate eagerly while shared inputs are completed. Their resulting arrays are counted, but this entry work is not pre-admitted by the chunk profiles. Budgeted host replay routes remain refused until their complete stages and owners have profiles.

The same fixed seed gives the same per-subject EV1 taste-shock choices across subject chunk widths and repeated calls. Keep the seed, parameters, initial conditions, and model fixed when checking that invariance.

Each simulate(...) call takes its population from the initial conditions. A model can simulate different populations; runtime executors reuse compiled programs when their shapes and other compilation inputs match. Parameter and support validation still run for every call.

Reuse compilation

pylcm enables a persistent JAX compilation cache by default. Check:

import jax

print(jax.config.jax_compilation_cache_dir)

A None value means no persistent cache. JAX_COMPILATION_CACHE_DIR chooses the full path; LCM_COMPILATION_CACHE_NAME chooses a project leaf under the default root. Compilation keys still change when program shapes or the software environment change.

Do not create JAX device arrays at module import solely for constants. Keep tables as NumPy arrays and convert inside traced functions; this avoids initializing a device in processes that only import the model.

Runtime environment controls are listed in Runtime, results, and persistence.

Choose knobs from measurements

The knobs above change how the same program is executed. The measured examples below show why each one needs a measurement on your own model before it is set.

Wider is not always faster

Under a budget, the width search admits the widest candidate that fits. That is a memory criterion, not a speed criterion, and a narrower width can run faster because smaller tiles materialise smaller intermediates:

So measure a width on a single-period replay, confirm it with one production pair, and then pin it with ExecutionConfig(axis_widths=...), per regime where the regimes differ (see Fix a planner axis width).

Which states can be sharded

Sharded states come first in the value array’s axis order. Compare sharded and unsharded values by state name or through the simulated panel, never by axis position.

At fp32, a sharded solve is a different compiled program and can round differently. Choices whose values are nearly tied can then flip, and a flip propagates through the rest of that subject’s history. Measured on the Borella model, fp32, 5×A40 sharded against 1×A40: 17,084 of 20,000 simulated subjects’ actions diverged at some point, with a median relative value gap of 1.7e-6 at the first divergence. Judge such a difference against an fp64 control before treating it as a defect or as noise.

Subject sharding for simulate-heavy models

When simulation, not the solve, dominates a call, spread the simulated subjects over all devices with ExecutionConfig(simulation_sharding="subjects"); see subject-parallel simulation. Measured on the ACA retirement model, 8×A40, fp32, one pair: the warm call went from 1,686 s to 1,064 s and warm simulation from 801 s to 191 s, with the solve time and device peak unchanged and bitwise-equal panels.

How to measure

Time three calls in one process and report them separately:

To iterate on one slow regime-period without a full backward induction, set LCM_CAPTURE_PERIOD="<regime>@<period>" and LCM_CAPTURE_DIR for one solve. That writes the period’s kernel inputs, and _lcm.solution.period_replay then runs just that period under a candidate setting. Check that the replayed values match the in-context values bitwise before you trust its timings. benchmarks/perf_loop.py in the pylcm repository runs a cold, warm and changed-parameter protocol on a named example model.

Benchmark the decision you face

For a solver comparison, hold the economic model and accuracy target fixed. Report:

Use the external LCM solver benchmarks for evolving shared evidence. The pylcm package’s own regression benchmarks belong to the Development chapter.

Checklist

pylcm
Authoring for EGM-family solvers
pylcm
Run invariant components as independent jobs