scan-specific checkpointing rule lives in shared graph_trainer code #31

Migrated from https://github.com/google-pytorch/torchtitan-archive/issues/410 (opened by jack8558 on 2026-09-15).

https://github.com/google-pytorch/torchtitan-archive/pull/360 adds a scan-aware branch to two files outside /tpu:

• experiments/graph_trainer/memory_policy.py (+13) -- tag_sac_policy tags any torch.ops.higher_order.scan node MUST_SAVE, before policy_fn sees it.
• experiments/graph_trainer/tests/test_passes.py (+20) -- the test for it.

Only the TPU lane builds a scanned decoder today, so this is lane-specific behaviour sitting in a shared pass. Worth moving into /tpu, or giving tag_sac_policy a hook the TPU lane can install the rule through.

The catch: the seams that exist (policy_fn, force_save_nodes, register_memory_policy) are all per-policy, so a /tpu-side version has to shadow every shared policy and will silently miss any policy added later.

Related: #29, #27.