Skip to content

Add ERNIE segmented checkpointing support - #2939

Merged
bghira merged 1 commit into
agent/segmented-checkpointing-cosmos3from
agent/segmented-checkpointing-ernie
Aug 3, 2026
Merged

Add ERNIE segmented checkpointing support#2939
bghira merged 1 commit into
agent/segmented-checkpointing-cosmos3from
agent/segmented-checkpointing-ernie

Conversation

@bghira

@bghira bghira commented Aug 2, 2026

Copy link
Copy Markdown
Owner

Summary

Splits the ERNIE segmented checkpointing support model integration out of #2925.

  • applies the model-family checkpointing hooks and support flags
  • adds this family to the relevant safety-check allow-lists
  • keeps the shared runtime/docs in the base PR so this diff stays model-specific

Stack

Base branch: agent/segmented-checkpointing-cosmos3

Validation

  • .venv/bin/python -m unittest tests.test_segmented_checkpointing_model_support -v
  • commit hooks: Black, isort, flake8, whitespace checks

@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from ce0b3b1 to 647af3e Compare August 3, 2026 08:50
@bghira
bghira requested a review from Copilot August 3, 2026 08:52
@bghira
bghira marked this pull request as ready for review August 3, 2026 08:52

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds ERNIE-family coverage for the segmented checkpointing feature set introduced in the stacked base branch, ensuring the ERNIE transformer exposes the expected control hooks and is permitted by safety checks.

Changes:

  • Add ERNIE to segmented-checkpointing stride allow-lists in safety checks.
  • Implement segmented-stride-aware checkpointing decision logic in the ERNIE transformer wrapper (via should_checkpoint_block) and expose the segment-stride setter.
  • Expand the segmented checkpointing model-support test suite to include ERNIE.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File Description
tests/test_segmented_checkpointing_model_support.py Adds ERNIE to the model-family capability coverage tests (expects backend/interval/stride controls).
simpletuner/helpers/training/default_settings/safety_check.py Allows ERNIE in the segment-stride supported model-family list.
simpletuner/helpers/models/ernie/transformer.py Adds segment-stride setter and uses should_checkpoint_block for stride-aware checkpointing decisions; broadens Unsloth backend detection.
simpletuner/helpers/models/ernie/transformer_diffusers.py Broadens Unsloth backend detection to match backend variants (e.g., unsloth-ffn).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from 647af3e to 2eca909 Compare August 3, 2026 09:02
@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from 2eca909 to 714afdf Compare August 3, 2026 10:45
@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from 714afdf to 725c0ef Compare August 3, 2026 15:30
@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from 725c0ef to e4946b4 Compare August 3, 2026 15:32
@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from e4946b4 to 2d767b7 Compare August 3, 2026 15:33
@bghira
bghira force-pushed the agent/segmented-checkpointing-ernie branch from 2d767b7 to ef6cf14 Compare August 3, 2026 15:35
@bghira
bghira merged commit 3504e55 into main Aug 3, 2026
@bghira
bghira deleted the agent/segmented-checkpointing-ernie branch August 3, 2026 15:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants