Debug regression

[raytomato]

This PR is created from main branch by dropping PRs merged recently. Trying to see if the CI issue was introduced by the PRs.

[raytomato]

Drop 041050d : https://github.com/apple/axlearn/actions/runs/15154143718/job/42605629824

Peak mem usage is at around 71%, slightly lower.

=== Memory Check: 2025-05-21 05:25:29 ===
Memory: 9.8Gi used, 147Mi free, 15Gi total
Memory usage: 62.8%
=== Memory Check: 2025-05-21 05:25:[59](https://github.com/apple/axlearn/actions/runs/15154143718/job/42605629824#step:6:60) ===
Memory: 10Gi used, 135Mi free, 15Gi total
Memory usage: [67](https://github.com/apple/axlearn/actions/runs/15154143718/job/42605629824#step:6:68).6%
=== Memory Check: 2025-05-21 05:26:29 ===
Memory: 11Gi used, 160Mi free, 15Gi total
Memory usage: [71](https://github.com/apple/axlearn/actions/runs/15154143718/job/42605629824#step:6:72).0%
=== Memory Check: 2025-05-21 05:26:59 ===
Memory: 11Gi used, 107Mi free, 15Gi total
Memory usage: 71.3%
=== Memory Check: 2025-05-21 05:27:29 ===
Memory: 11Gi used, 182Mi free, 15Gi total
Memory usage: 70.9%
=== Memory Check: 2025-05-21 05:27:59 ===
Memory: 8.8Gi used, 2.5Gi free, 15Gi total
Memory usage: 56.4%

[raytomato]

Drop 041050d and 8d4dedf: https://github.com/apple/axlearn/actions/runs/15154303034/job/42606047318?pr=1196

=== Memory Check: 2025-05-21 05:38:29 ===
Memory: 11Gi used, 126Mi free, 15Gi total
Memory usage: 74.0%
=== Memory Check: 2025-05-21 05:38:59 ===
Memory: 11Gi used, 158Mi free, 15Gi total
Memory usage: [76](https://github.com/apple/axlearn/actions/runs/15154303034/job/42606047318?pr=1196#step:6:77).1%

[raytomato]

These are the only two PRs merged yesterday. The mem usage peak difference doesn't seem significant enough to trigger the failure. (And it aligns with the expectation if we check what these two commits are about.)