ENH Move ACA execution controls to Model policy, and related solver/precision work - #17
Draft
hmgaudecker wants to merge 15 commits into
Draft
hmgaudecker wants to merge 15 commits into
hmgaudecker wants to merge 15 commits into
Conversation
`build_actions` and the nongroup builder narrowed the M1 regime for NBEGM, dropping `buy_private` and `labor_supply` as actions and fixing their former outputs to constants. That made the solved model a different model from the one brute force solves: the household lost its coverage and hours choices. The regime now declares whichever choices its structure affords under every solver. pylcm's ride-along discrete envelope is written over a single action's grid and refuses a regime declaring several, so model build under NBEGM raises until that arity widens. The refusal is the honest outcome — a solver that cannot carry a choice refuses the regime rather than being handed a narrower one. Tests split accordingly: the regime-level and solver-config assertions stay green, model-build assertions become strict xfails naming the arity, and one green test pins the refusal itself so the xfails flip when the arity widens. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018yHhuFzqsDw1MhdB1i2Ljm
The M1 regime declares both buy_private and labor_supply, and pylcm's envelope now branches over the Cartesian product of a regime's discrete action grids, so the eight strict xfails pinned to the old single-action refusal all XPASS. Removes them and converts the refusal test into the capability it replaced: the model builds with both actions live. 25 passed where 15 had failed. Every one of the 15 was a stale expectation -- including the five policy-variant builds, which are reported as FAILED rather than XPASS under strict=True and so read like real build failures until the XPASS(strict) marker is inspected. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018yHhuFzqsDw1MhdB1i2Ljm
`solver="nbegm"` attached the solver to one regime and left the other 17 on brute force, so the model's stated solver was not the one that produced most of its result. Every living regime now gets the NBEGM config, and the retiree and tied builders carry the savings-form budget (`resources`, `savings`) that the solver's contract reads -- without it those regimes cannot be built at all. The build-time affinity and interval-constancy probes cannot run on the added regimes: they differentiate the budget on scalar inputs and `assets_and_income.capital_income` declares `rate_of_return: ScalarFloat`, which the probe's one-element array violates. The probes fall back to `assume_declared` and warn, so NB-EGM's exactness precondition is asserted rather than checked there, and the solve needs validating against an independent reference. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018yHhuFzqsDw1MhdB1i2Ljm
The probes run on the first solve against the model's complete parameter vector, so the tax tables and threshold schedules the budget reads are the model's own rather than synthesized stand-ins. Affinity and constancy are checked rather than asserted.
pylcm takes the liquid roles from the regime: a `ConsumptionSavingsRegime` names the liquid state, the action paid from resources, the resources node and the post-decision state, and the solver carries numerical configuration only. `ACA_LIQUID_MARGIN` states that once for the three living-regime builders, which now route through `build_alive_regime`; a brute-force regime has no `resources` or `savings` node to name and stays a plain `Regime`. The role assertions move with the declaration: they read `regime.liquid`, which is public API, rather than the bound solver's attributes.
Both callers of `build_model_constraints` still said DC-EGM gets no broadcast constraint. It is declared under every solver: the EGM solve enforces the limit through the savings grid's lower bound, but forward simulation re-decides consumption by an argmax over the consumption grid and needs the explicit feasibility mask. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018yHhuFzqsDw1MhdB1i2Ljm
The xfail reason described the assets-law chain, but the build stops earlier: pylcm's DC-EGM refuses the broadcast borrowing constraint for reading continuous variables. Both blockers are now listed in the order they fire. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018yHhuFzqsDw1MhdB1i2Ljm
pylcm's `Model.solve` returns a `SolutionResult` and `Model.simulate` takes `solution=` in place of `period_to_regime_to_V_arr=`. The value mapping these tests read is `SolutionResult.values`; passing `None` for the old keyword meant "solve inside simulate", which is the new default. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y7QCx9tkGD3TqwNw681QwD
Forward ExecutionConfig through baseline, ACA and benchmark factories. Derive omitted accelerator budgets from selected JAX allocator limits and reject retired GridConfig execution fields. Preserve economic grids and numerical solver settings. Validation against pylcm f6d41d54: exact ASV CPU preflight, focused compatibility tests, Ruff, ty and repository hooks pass. DC-EGM construction still exposes an upstream continuation-template placement failure on three CPU devices.
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.
Opens the accumulated
codex/call-time-simulation-populationwork for review: 15 commits, 31 files, +735/−705 againstmain.What is here
Broadly three strands:
ENH Move ACA execution controls to Model policyaddssrc/aca_model/execution.py(new on this branch — absent frommain) withexecution_config_for_devices, andAPI Remove constructor population hints from model factoriesfollows through on the factory signatures.SolutionResultsolve API, benchmark migrated to breakpoint grids, explicit fp32 ACA runs, and the borrowing limit declared structurally.Known issue in this branch — being fixed upstream in pylcm
execution.py:52setsdevice_memory_bytes=min(limits), where each limit is the device's rawmemory_stats()["bytes_limit"]— the entire XLA BFC pool, nothing subtracted.pylcm documents that field as a "Per-device ceiling for represented compiler reservation plus residency", adding that "runtime storage omitted by the compiler remains outside this accounting scope". Passing the full pool therefore leaves nothing for the categories outside that scope: NCCL buffers for the device clique, cuBLAS workspace, driver context, fragmentation.
pylcm's admission accepts the largest width whose compiled reservation fits the budget, so it settles as close to the ceiling as width quantisation allows. Observed consequence, SLURM job 27624359 (3xA40, fp32, production
sharded_states=("pref_type",)):XLA_PYTHON_CLIENT_MEM_FRACTION=0.90on a 46,068 MiB A40 => 40.49 GiB budgetThis is on ACA's default GPU policy, not an experimental one —
pref_typeover three devices is whataca-slurmsets for both A100 and A40. The siblingassets-sharded leg completed, but incidentally: its per-cell cost is higher, so its admissible widths were 8x smaller (102,600 vs 820,800 cells at the same regime and age) and never approached the ceiling. A100s having 80 GB where A40s have 46 is the likely reason this had not surfaced before.The repair is being made in pylcm, not here. pylcm knows the device count and emits the collectives whose buffers the caller is being asked to anticipate, and today validates only that the budget is a positive integer without comparing it to the device it is about to plan for. This branch is therefore left unpatched on purpose: no
int(0.85 * min(limits))mitigation, so nothing here pre-empts or double-counts the upstream fix. This section will be updated to point at it once it lands.Unresolved: the 39.82 GiB allocation has not been attributed to a specific buffer —
preallocated-temparena (in-contract, a fragmentation story) versus adevice_putdestination (out-of-contract on the discrete-sharding route).aca-slurm/src/aca_slurm/experiments/_buffer_assignment.pyexists to settle it.Status
Draft. Opened to surface the branch and the issue above for review; no CI run has been inspected as part of opening it.
https://claude.ai/code/session_013LSspB3nyDJKhGrGZCvjCz