Docs: brute-force vs DC-EGM comparison, example DC-EGM variant, CHANGES#382
Open
hmgaudecker wants to merge 31 commits into
Open
Docs: brute-force vs DC-EGM comparison, example DC-EGM variant, CHANGES#382hmgaudecker wants to merge 31 commits into
hmgaudecker wants to merge 31 commits into
Conversation
…notebook. The example module gains the DC-EGM variant of the model (get_dcegm_model: declared resources/savings/inverse_marginal_utility, savings-based wealth transition, no borrowing constraint, cubically clustered savings grid); the explanation notebook replaces the placeholder with an accuracy comparison against a fine-grid brute-force reference plus per-solve timings; CHANGES.md documents the DC-EGM solver, regime-level taste shocks, and the example. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
…tion). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
hmgaudecker
changed the base branch from
feat/dcegm-discrete-choices
to
feat/dcegm-simulate
Benchmark comparison (main → HEAD)Comparing
|
The solver-comparison fixtures import resources, savings, next_wealth_from_savings, and inverse_marginal_utility from lcm_examples.iskhakov_et_al_2017 instead of redefining them. The explanation notebook's intro and the explanations index describe the brute-force vs DC-EGM comparison the notebook actually contains. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The dominant accelerator resident in a DC-EGM solve is the rolling `next_regime_to_egm_carry` — a dense-endogenous-grid, per-discrete-action carry held for *every* carry-producing regime at once. Confirmed n_assets-independent (~63 GB on the ACA benchmark at n_assets=3/4/8; the value function, which scales with the Euler/n_assets axis, is negligible by comparison). Two complementary changes shrink it: - Reachable-target carry filtering (unconditional): each `egm_step` only ever indexes its reachable targets, so it is now handed just that subset (`_reachable_carry_subset`) instead of the full all-regimes mapping. `reachable_targets` is threaded `build_egm_step_functions` -> `SolverKernels` -> `SolutionPhase`. Correctness-preserving (smaller kernel signature). - `offload_carries_to_host` on `Model.solve` / `Model.simulate` (default off): evict the rolled carries to host memory between periods (inside `_roll_continuation_inputs`), so the device holds only the reachable subset each kernel re-uploads at dispatch rather than every regime's carry at once. Trades per-period host round-trips for a large drop in peak device memory; a no-op on a CPU-only host. Value functions are bit-identical with and without the flag (new parity tests in `tests/solution/test_egm_carry_offload.py`, including the ACA-shaped passive-AIME + asset-row config). Full suite green at float64. Also renames `EgmCarry` -> `EGMCarry` (and `EgmCarryProducer`, `EgmStepFunction`) for acronym casing; mechanical, spans the same files. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Under a distributed (sharded) solve the rolling carries are device-sharded. `offload_carries_to_host` moved them with `jax.device_put(..., cpu_device)`, which produces a CPU-*committed* array; re-uploading that to an AOT-compiled egm_step rejects it — "input shardings disagree with the shardings of arguments" — because a committed array must match the compiled input sharding. Use `jax.device_get` (host NumPy) instead. Host NumPy is *uncommitted*, so JAX reshards it to the kernel's compiled `in_shardings` at dispatch rather than erroring. Frees the same device memory; a no-op on a CPU-only host. Parity tests unchanged (single-device, so device_get round-trips identically). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The host-offload (`offload_carries_to_host`) copied the rolled carries to host with `jax.device_get` but never freed the device originals, so the previous period's carry set lingered (GC lag / async dispatch) alongside the next period's — two periods co-resident, and the peak never dropped (measured identical to no-offload at ~63 GB). After `device_get`, delete this period's freshly produced kernel-output carries (`period_egm_carries`) — and only those. The carried-forward entries are either already host (a prior offload) or the regimes' shared `egm_carry_template`, which must not be deleted: it is reused next period for inactive regimes and across solves, and a blanket free corrupts the cached model (observed as "deleted array" errors when a `@cache`d model is solved twice). With only one period's carries resident, the peak should roughly halve. Parity unchanged (the host copies carry the data forward); full EGM suite green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fifth DC-EGM stack PR (stacked on #380, sibling of #381): the docs and example finalization.
lcm_examples.iskhakov_et_al_2017gains the DC-EGM variant of the retirement model —get_dcegm_model(n_periods)with the declaredresources/savings/inverse_marginal_utilityfunctions, the savings-based wealth transition, no borrowing constraint (enforced intrinsically), and a cubically clustered savings grid.get_paramsworks unchanged; the two builders are mathematically equivalent specs.NotImplementedErroronsimulate()) is documented where readers will look for policy figures.Verified: full suite 1172 passed / 46 skipped, ty clean, prek clean (notebook cells stored as line arrays, outputs stripped). The notebook's new cells were smoke-run as scripts at reduced horizon; the full notebook executes on the RTD preview build of this PR.
🤖 Generated with Claude Code