Share one CUSTOM_REMAT_TENSORS list between MaxTextConfig and RLConfig

[lokic233]

remat_policy: custom filters the tensors the user pinned against a hardcoded
name list. That list is written twice -- once in MaxTextConfig.model_post_init
and once in RLConfig.model_post_init -- and the two copies have drifted:
RLConfig's is missing "kv_proj".

Both classes inherit the same RematAndOffload field set, so under RL a config
with kv_proj: device validates cleanly, is then dropped from
tensors_on_device, and the unified KV projection is silently rematerialised
instead of kept. There is no warning; the only symptom is extra recompute.

Collapse both onto one module-level constant. MaxTextConfig's existing order is
preserved exactly, so the non-RL path is unchanged.

Three tests are added to tests/unit/configs_test.py:

• the constant matches the declared RematLocation fields, minus the one
  documented as non-selectable;
• every name in it has a real checkpoint_name() call in the model code;
• MaxTextConfig and RLConfig honour the same set, for device and offload
  (parametrized, so 4 test cases in total).

Controls that were run, not just asserted:

• Reverting only the two call sites, keeping the constant, turns the parity
  test red with only MaxTextConfig honours: ['kv_proj'] for both device and
  offload. So the test detects the behavioural drift, not the refactor.
• Injecting "engram" into the constant turns both guard tests red.
• The pre-existing 76 tests pass before and after, with the same single
  pre-existing pydantic serializer warning on rl.yml.

A small extraction rides along: the dict normalisation inside
run_config_validation is lifted into normalize_config_dict() so the new tests
run the same preprocessing as the existing 76 rather than a second copy of it.

checkpoint_names_in_tree() parses with ast rather than a regex. This is not
stylistic -- a textual scan also matches checkpoint_name( written inside a
comment, including the new comment in types.py that names "engram" as the
excluded case, which made the check vacuously pass for exactly the name it
exists to catch. That was caught by the negative control above.

Two related problems are deliberately left alone, called out in comments rather
than fixed here, since both need a decision from their owners:

• engram is a declared RematLocation field with no checkpoint_name()
  anywhere in src/maxtext. Measured on this base: 97 checkpoint_name call
  sites, 90 with a literal name, 22 distinct names; the 7 non-literal sites
  are 6 residual_checkpoint_name in the splash-attention kernels and one
  string concatenation in linears.py that yields mlpwi / mlpwi_0 / mlpwi_1,
  all three of which appear as literals elsewhere. engram appears in none
  of them. Listing it would promise control the model code cannot honour, so
  it is excluded and the reason is recorded in the test. Either the layer
  needs a checkpoint_name or the field should be removed.

• configs/pyconfig_deprecated.py:510, validate_and_assign_remat_tensors, holds
  a third and older copy of the same list: 15 names, missing attention_out,
  indexer_cutoff_threshold, kv_proj, mla_kv, mla_q and qkv_proj. Pointing it
  at the constant is mechanically safe -- adding
  from maxtext.configs import types there imports cleanly in both orders,
  checked -- but it is not a no-op: it would make those six names newly
  selectable on that path, which is a behaviour change to a deprecated config
  system and belongs in its own change with its own owner.

One test-only exclusion worth flagging for review: indexer_cutoff_threshold
is in the constant but skipped by the parity test, because setting it requires
use_indexer=True, which in turn requires attention_type='mla', a nonzero
q_lora_rank and a compatible attention kernel. Bolting a whole MLA model spec
onto a list-parity test would make it break for unrelated reasons. It is still
covered by the declared-fields test.

[gemini-code-assist]

Code Review

This pull request centralizes the list of custom rematerialization tensors into a single CUSTOM_REMAT_TENSORS constant in src/maxtext/configs/types.py, resolving a drift between MaxTextConfig and RLConfig where RLConfig was missing the kv_proj tensor. It also refactors configuration preprocessing in tests and adds a comprehensive suite of unit tests to validate custom remat tensors, including an AST-based check to verify that all listed tensors have corresponding checkpoint_name calls in the codebase. Feedback suggests improving the robustness of this AST-based check to handle aliased imports of checkpoint_name.

[gemini-code-assist]

The current logic for identifying checkpoint_name calls is brittle. It correctly handles checkpoint_name(...) and some_module.checkpoint_name(...), but it will fail to detect calls if checkpoint_name is imported with an alias, e.g., from jax.ad_checkpoint import checkpoint_name as cn. This could lead to test failures if developers use aliases, potentially causing them to incorrectly remove valid tensors from CUSTOM_REMAT_TENSORS.

To make this more robust, you could parse import statements (ast.Import and ast.ImportFrom) at the beginning of each file to build a mapping of local names to their canonical function names. This would allow the check to work correctly even with aliased imports.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!