Training engine: buffer metrics on the host by igorts-git · Pull Request #6017 · AI-Hypercomputer/maxtext

igorts-git · 2026-10-28T22:54:44Z

Description

Changes the training engine's MetricsRecorder from appending to JAX arrays on every micro-batch to buffering metrics in host-side Python lists and building the arrays once per step, at flush.

Why: the eager jnp.append path dispatched one device op per metric per micro-batch (about 100 metrics for GRPO). The new flush is one jnp.stack per metric per step, independent of the micro-batch count.

Changes:

MetricsRecorder keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup.
buffer_metrics eagerly validates metric types and shapes host-side without dispatching device operations, preventing malformed metrics from entering pending buffers.
record_metrics skips intermediate_outputs and tuple/list aux values.
Adds documentation rule in AGENTS.md requiring comments and docstrings to explain current code state from first principles rather than comparing against deleted or legacy implementations.

Differences from `atwigg/mlperf`

Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: MetricsRecorder.restore_buffers, weighted_metric_bounds and the metric_bounds checkpoint metadata, the engine's _restore_metrics_buffer and the intra-step loss-cache rebuild are not ported, nor are the tests for them.
MLPerf's update(self, **kwargs) is not adopted.
Added _validate_metric to eagerly validate metric types and shapes host-side in MetricsRecorder.buffer_metrics, rather than deferring failures to flush time.

Attribution

Ported from Param Bole's #5568 on atwigg/mlperf.

Co-authored-by: Param Bole parambole@google.com

Fixes #6017

Tests

cd tests/post_training/unit
python -m pytest maxtext_engine_test.py -q

Result: 71 passed. Covers Python-list buffering, eager host-side type/shape validation, the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, and the single-stack flush. Pre-commit (codespell, pylint, pyink) passed.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

codecov · 2026-10-29T00:03:43Z

Codecov Report

❌ Patch coverage is 85.13514% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/training_engine/metrics.py	84.93%	5 Missing and 6 partials ⚠️

📢 Thoughts on this report? Let us know!

- MetricsRecorder keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - update() flushes the step's metrics before the optimizer step, so an unconvertible value raises with nothing applied. - record_metrics skips intermediate_outputs and tuple/list aux values. Differences from atwigg/mlperf: - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: MetricsRecorder.restore_buffers, weighted_metric_bounds and the metric_bounds checkpoint metadata, the engine's _restore_metrics_buffer and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - The pre-step flush is added to PR 5999's explicit update(self); MLPerf's update(self, **kwargs) is not adopted. Co-authored-by: Param Bole <parambole@google.com> Fixes AI-Hypercomputer#6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 70 passed. Covers Python-list buffering and the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, the single-stack flush, and the pre-update flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files).

Imported from GitHub PR #6017 Stacked on top of #6038 (PR 0). # Description Changes the training engine's `MetricsRecorder` from appending to JAX arrays on every micro-batch to buffering metrics in host-side Python lists and building the arrays once per step, at flush. Why: the eager `jnp.append` path dispatched one device op per metric per micro-batch (about 100 metrics for GRPO). The new flush is one `jnp.stack` per metric per step, independent of the micro-batch count. Changes: - `MetricsRecorder` keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - `update()` flushes the step's metrics before the optimizer step, so an unconvertible value raises with nothing applied. - `record_metrics` skips `intermediate_outputs` and tuple/list aux values. ## Differences from `atwigg/mlperf` - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: `MetricsRecorder.restore_buffers`, `weighted_metric_bounds` and the `metric_bounds` checkpoint metadata, the engine's `_restore_metrics_buffer` and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - The pre-step flush is added to PR 5999's explicit `update(self)`; MLPerf's `update(self, **kwargs)` is not adopted. ## Attribution Ported from Param Bole's #5568 on `atwigg/mlperf`. Co-authored-by: Param Bole <parambole@google.com> Fixes #6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 70 passed. Covers Python-list buffering and the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, the single-stack flush, and the pre-update flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files). PiperOrigin-RevId: 1009954644

- MetricsRecorder keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - buffer_metrics eagerly validates metric types and shapes host-side without dispatching device operations, preventing malformed metrics from entering pending buffers. - record_metrics skips intermediate_outputs and tuple/list aux values. Differences from atwigg/mlperf: - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: MetricsRecorder.restore_buffers, weighted_metric_bounds and the metric_bounds checkpoint metadata, the engine's _restore_metrics_buffer and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - MLPerf's update(self, **kwargs) is not adopted. Co-authored-by: Param Bole <parambole@google.com> Fixes AI-Hypercomputer#6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 71 passed. Covers Python-list buffering, eager host-side type/shape validation, the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, and the single-stack flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files).

- MetricsRecorder keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - buffer_metrics eagerly validates metric types and shapes host-side without dispatching device operations, preventing malformed metrics from entering pending buffers. - record_metrics skips intermediate_outputs and tuple/list aux values. - Adds documentation rule in AGENTS.md requiring comments and docstrings to explain current code state from first principles rather than comparing against deleted or legacy implementations. Differences from atwigg/mlperf: - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: MetricsRecorder.restore_buffers, weighted_metric_bounds and the metric_bounds checkpoint metadata, the engine's _restore_metrics_buffer and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - MLPerf's update(self, **kwargs) is not adopted. Co-authored-by: Param Bole <parambole@google.com> Fixes AI-Hypercomputer#6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 71 passed. Covers Python-list buffering, eager host-side type/shape validation, the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, and the single-stack flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files).

- MetricsRecorder keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - buffer_metrics eagerly validates metric types and shapes host-side without dispatching device operations, preventing malformed metrics from entering pending buffers. - record_metrics skips intermediate_outputs and tuple/list aux values. - Adds documentation rule in AGENTS.md requiring comments and docstrings to explain current code state from first principles rather than comparing against deleted or legacy implementations. Differences from atwigg/mlperf: - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: MetricsRecorder.restore_buffers, weighted_metric_bounds and the metric_bounds checkpoint metadata, the engine's _restore_metrics_buffer and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - MLPerf's update(self, **kwargs) is not adopted. - Added _validate_metric to eagerly validate metric types and shapes host-side in MetricsRecorder.buffer_metrics, rather than deferring failures to flush time. Co-authored-by: Param Bole <parambole@google.com> Fixes AI-Hypercomputer#6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 71 passed. Covers Python-list buffering, eager host-side type/shape validation, the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, and the single-stack flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files).

Imported from GitHub PR #6017 # Description Changes the training engine's `MetricsRecorder` from appending to JAX arrays on every micro-batch to buffering metrics in host-side Python lists and building the arrays once per step, at flush. Why: the eager `jnp.append` path dispatched one device op per metric per micro-batch (about 100 metrics for GRPO). The new flush is one `jnp.stack` per metric per step, independent of the micro-batch count. Changes: - `MetricsRecorder` keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - `buffer_metrics` eagerly validates metric types and shapes host-side without dispatching device operations, preventing malformed metrics from entering pending buffers. - `record_metrics` skips `intermediate_outputs` and tuple/list aux values. - Adds documentation rule in `AGENTS.md` requiring comments and docstrings to explain current code state from first principles rather than comparing against deleted or legacy implementations. ## Differences from `atwigg/mlperf` - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: `MetricsRecorder.restore_buffers`, `weighted_metric_bounds` and the `metric_bounds` checkpoint metadata, the engine's `_restore_metrics_buffer` and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - MLPerf's `update(self, **kwargs)` is not adopted. ## Attribution Ported from Param Bole's #5568 on `atwigg/mlperf`. Co-authored-by: Param Bole <parambole@google.com> Fixes #6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 71 passed. Covers Python-list buffering, eager host-side type/shape validation, the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, and the single-stack flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files). PiperOrigin-RevId: 1009954644

Imported from GitHub PR #6017 # Description Changes the training engine's `MetricsRecorder` from appending to JAX arrays on every micro-batch to buffering metrics in host-side Python lists and building the arrays once per step, at flush. Why: the eager `jnp.append` path dispatched one device op per metric per micro-batch (about 100 metrics for GRPO). The new flush is one `jnp.stack` per metric per step, independent of the micro-batch count. Changes: - `MetricsRecorder` keeps each step's metrics in Python lists and materializes them into 1D arrays once, at flush. The flush is all-or-nothing, and pending items are keyed by buffer identity and dropped on eviction and cleanup. - `buffer_metrics` eagerly validates metric types and shapes host-side without dispatching device operations, preventing malformed metrics from entering pending buffers. - `record_metrics` skips `intermediate_outputs` and tuple/list aux values. - Adds documentation rule in `AGENTS.md` requiring comments and docstrings to explain current code state from first principles rather than comparing against deleted or legacy implementations. ## Differences from `atwigg/mlperf` - Checkpoints are taken only between steps (PR 0), so the metrics are never checkpointed: `MetricsRecorder.restore_buffers`, `weighted_metric_bounds` and the `metric_bounds` checkpoint metadata, the engine's `_restore_metrics_buffer` and the intra-step loss-cache rebuild are not ported, nor are the tests for them. - MLPerf's `update(self, **kwargs)` is not adopted. - Added `_validate_metric` to eagerly validate metric types and shapes host-side in `MetricsRecorder.buffer_metrics`, rather than deferring failures to flush time. ## Attribution Ported from Param Bole's #5568 on `atwigg/mlperf`. Co-authored-by: Param Bole <parambole@google.com> Fixes #6017 # Tests ```bash cd tests/post_training/unit python -m pytest maxtext_engine_test.py -q ``` Result: 71 passed. Covers Python-list buffering, eager host-side type/shape validation, the flush layout (including more micro-batches after a mid-step flush), same-id buffers kept apart, all-or-nothing flush, eviction and cleanup of pending items, and the single-stack flush. Pre-commit (codespell, pylint, pyink) passed. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files). PiperOrigin-RevId: 1009954644

github-actions Bot added the skip-ci-tests Add this label to the PR to skip running CI tests label Oct 28, 2026

igorts-git force-pushed the port/pr1-metrics-checkpointing branch 2 times, most recently from b413bef to dd375e2 Compare October 28, 2026 23:50

igorts-git changed the title ~~Training engine: buffer metrics on the host and harden checkpoint save/restore~~ Training engine: buffer metrics on the host Oct 28, 2026

igorts-git mentioned this pull request Oct 28, 2026

Training engine: harden checkpoint save/restore #6019

Closed

4 tasks done

SurbhiJainUSC force-pushed the igorts/tunix-dist-rl-api branch 2 times, most recently from 3e61343 to 9cf01be Compare October 29, 2026 16:35

igorts-git force-pushed the port/pr1-metrics-checkpointing branch from dd375e2 to 0053443 Compare October 29, 2026 18:10

copybara-service Bot deleted the branch AI-Hypercomputer:igorts/tunix-dist-rl-api October 29, 2026 21:48

copybara-service Bot closed this Oct 29, 2026

igorts-git reopened this Oct 30, 2026

igorts-git changed the base branch from igorts/tunix-dist-rl-api to main October 30, 2026 03:44

igorts-git force-pushed the port/pr1-metrics-checkpointing branch from 0053443 to 49e5c89 Compare October 30, 2026 03:53

copybara-service Bot mentioned this pull request Oct 30, 2026

PR #6017: Training engine: buffer metrics on the host #6051

Merged

4 tasks done

igorts-git force-pushed the port/pr1-metrics-checkpointing branch from 49e5c89 to db6c0a4 Compare October 30, 2026 16:38

igorts-git force-pushed the port/pr1-metrics-checkpointing branch from db6c0a4 to ae6ee1a Compare October 30, 2026 17:09

igorts-git force-pushed the port/pr1-metrics-checkpointing branch from ae6ee1a to a092171 Compare October 30, 2026 17:25

igorts-git force-pushed the port/pr1-metrics-checkpointing branch from a092171 to 99673a5 Compare October 30, 2026 17:29

igorts-git mentioned this pull request Oct 30, 2026

Training engine: harden checkpoint save/restore #6058

Open

4 tasks done

copybara-service Bot closed this in 647d453 Oct 30, 2026

github-actions Bot mentioned this pull request Oct 31, 2026

Failed build: MaxText Package Tests #6052

Open

Navigation Menu

Repository navigation

Training engine: buffer metrics on the host - #6017

Training engine: buffer metrics on the host#6017
igorts-git wants to merge 1 commit into
AI-Hypercomputer:mainfrom
igorts-git:port/pr1-metrics-checkpointing

igorts-git commented Oct 28, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Oct 29, 2026 •

edited

Loading

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

igorts-git commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Differences from atwigg/mlperf

Attribution

Tests

Checklist

Uh oh!

codecov Bot commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

igorts-git commented Oct 28, 2026 •

edited

Loading

Differences from `atwigg/mlperf`

codecov Bot commented Oct 29, 2026 •

edited

Loading