Skip to content

support epoch in hf pipeline#3088

Merged
copybara-service[bot] merged 1 commit intomainfrom
aireen/hf_epoch
Feb 6, 2026
Merged

support epoch in hf pipeline#3088
copybara-service[bot] merged 1 commit intomainfrom
aireen/hf_epoch

Conversation

@aireenmei
Copy link
Collaborator

@aireenmei aireenmei commented Feb 4, 2026

Description

  • support multi epoch training in hf pipeline. Multi-epoch is already supported in other pipelines . b/482070325
  • migrate data input validations from pyconfig_deprecated.py to types.py
  • Fix mdformat for correct numbering and fix numbering in some docs

Tests

Test on a small dataset with 1 epoch, data run out after 21 steps: https://paste.googleplex.com/5202487038574592
Test on the same dataset with 2 epochs, data run out after 42 steps: https://paste.googleplex.com/5516188497084416
Also used this script to test the effects of .repeat() and .shuffle(). The output is as expected.

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 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
Copy link

codecov bot commented Feb 4, 2026

Codecov Report

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

Files with missing lines Patch % Lines
src/MaxText/input_pipeline/_hf_data_processing.py 8.33% 9 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

@aireenmei aireenmei force-pushed the aireen/hf_epoch branch 3 times, most recently from 6670227 to fe0d2d5 Compare February 5, 2026 19:59
@aireenmei
Copy link
Collaborator Author

@SamuelMarks could you review the change in .pre-commit-config.yaml?

@copybara-service copybara-service bot merged commit a54b374 into main Feb 6, 2026
39 of 42 checks passed
@copybara-service copybara-service bot deleted the aireen/hf_epoch branch February 6, 2026 01:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants