feat(automodel): expose activation_checkpointing on customization jobs - #1252
Open
htolentino-nvidia wants to merge 2 commits into
Open
feat(automodel): expose activation_checkpointing on customization jobs#1252htolentino-nvidia wants to merge 2 commits into
htolentino-nvidia wants to merge 2 commits into
Conversation
Automodel supports activation checkpointing end to end — the finetune recipe reads `distributed.activation_checkpointing` off the config (recipes/_dist_utils.py) and normalizes bool | "full" | "selective" — but the platform had no way to set it, and FSDP2Config defaults it to False. Every customization job therefore trained with it off. That blocks porting the NVIDIA Nemotron LoRA cookbooks, which set `activation_checkpointing: true`; the Nemotron 3 Super recipe annotates it "reduces peak memory (avoids OOM on 80GB)". Adds the field alongside expert_parallel_size at each layer it already travels through: plugin schema, adapter, service API schema, compiler, task config, and YAML emission. Typed bool | Literal["full", "selective"] | None to match what Automodel parses, rather than a plain bool, so selective checkpointing — the cheaper mode — stays reachable. Defaults to None and is omitted from the emitted YAML when unset, so existing jobs compile byte-identically; the 15 golden contract configs are unchanged. An explicit false is still emitted, to distinguish "user turned it off" from "not configured". Note the pre-existing embedding branch it now precedes is unreachable: `embedding_config` is a bare EmbeddingConfig() that is never populated, so `do_gradient_checkpointing` is always False. It is left in place as the fallback in case that config is ever wired up. Not included: OpenAPI regeneration (`make refresh-openapi`) and the Stainless SDK sync, which need to run before the field reaches the CLI, Python SDK, or Studio. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Henrique Tolentino <htolentino@nvidia.com>
Contributor
📝 WalkthroughWalkthroughActivation checkpointing now accepts boolean and ChangesActivation checkpointing
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant AutomodelConfig
participant TrainingCompiler
participant BackendConfig
participant DistributedTraining
AutomodelConfig->>TrainingCompiler: activation_checkpointing
TrainingCompiler->>BackendConfig: ParallelismConfig
BackendConfig->>DistributedTraining: explicit checkpointing mode
BackendConfig->>DistributedTraining: fallback when setting is unset
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
services/automodel/src/nmp/automodel/app/jobs/training/compiler.py (1)
173-182: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winOmit unset
activation_checkpointingfromPlatformJobStep.config.
model_dump(mode="json")includesNone, but backend tests require the unset key to be absent. Remove onlyparallelism.activation_checkpointingwhen its value isNone; preserve explicitFalseand string modes. Add compiler coverage for all three cases.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@services/automodel/src/nmp/automodel/app/jobs/training/compiler.py` around lines 173 - 182, Update the compiler flow constructing TrainingStepConfig.ParallelismConfig so parallelism.activation_checkpointing is omitted from the serialized PlatformJobStep.config when unset (None), while preserving explicit False and string modes; add compiler coverage for unset, False, and string values.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@services/automodel/src/nmp/automodel/app/jobs/training/compiler.py`:
- Around line 173-182: Update the compiler flow constructing
TrainingStepConfig.ParallelismConfig so parallelism.activation_checkpointing is
omitted from the serialized PlatformJobStep.config when unset (None), while
preserving explicit False and string modes; add compiler coverage for unset,
False, and string values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4b22ddda-a099-4a61-8e4b-2fdeaca0aabd
📒 Files selected for processing (8)
plugins/nemo-automodel/src/nemo_automodel_plugin/schema.pyservices/automodel/src/nmp/automodel/adapter.pyservices/automodel/src/nmp/automodel/api/v2/jobs/schemas.pyservices/automodel/src/nmp/automodel/app/jobs/training/compiler.pyservices/automodel/src/nmp/automodel/app/jobs/training/schemas.pyservices/automodel/src/nmp/automodel/tasks/training/backends/config.pyservices/automodel/tests/tasks/training/backends/test_config.pyservices/automodel/tests/test_adapter.py
Contributor
|
Runs script/generate-openapi-spec.sh after adding the field, adding activation_checkpointing to AutomodelParallelismSpec in the customizer plugin spec as anyOf[boolean, enum[full, selective]]. Only the plugin spec changes; openapi/openapi.yaml is untouched because the customization API is served through the plugin router. No drift beyond the new field. The Stainless sync for the Python SDK is not included — it needs STAINLESS_API_KEY and is a separate maintainer step. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Henrique Tolentino <htolentino@nvidia.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@plugins/nemo-customizer/openapi/openapi.yaml`:
- Around line 1583-1585: Update the source schema description for
activation_checkpointing to state that omission preserves the embedding-model
fallback, while an explicit false disables activation checkpointing; then
regenerate the OpenAPI file so the generated description matches the backend
contract.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 67900e29-e0fd-4cc2-83ec-20147ffae0d4
📒 Files selected for processing (1)
plugins/nemo-customizer/openapi/openapi.yaml
Comment on lines
+1583
to
+1585
| description: Recompute activations during the backward pass to cut peak | ||
| memory at the cost of speed. 'selective' checkpoints only the most memory-heavy | ||
| ops. Left unset, Automodel defaults to disabled. |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Document the unset fallback correctly.
When activation_checkpointing is omitted, the backend preserves the embedding-model fallback. An explicit false disables activation checkpointing. Update the source schema description before regenerating this file.
Suggested wording
- ops. Left unset, Automodel defaults to disabled.
+ ops. When unset, Automodel uses its model-specific default; set
+ false to disable activation checkpointing explicitly.The supplied backend contract and plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py:122-137 distinguish omitted values from explicit false.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| description: Recompute activations during the backward pass to cut peak | |
| memory at the cost of speed. 'selective' checkpoints only the most memory-heavy | |
| ops. Left unset, Automodel defaults to disabled. | |
| description: Recompute activations during the backward pass to cut peak | |
| memory at the cost of speed. 'selective' checkpoints only the most memory-heavy | |
| ops. When unset, Automodel uses its model-specific default; set | |
| false to disable activation checkpointing explicitly. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-customizer/openapi/openapi.yaml` around lines 1583 - 1585,
Update the source schema description for activation_checkpointing to state that
omission preserves the embedding-model fallback, while an explicit false
disables activation checkpointing; then regenerate the OpenAPI file so the
generated description matches the backend contract.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Automodel supports activation checkpointing end to end — the finetune recipe reads
distributed.activation_checkpointingoff the config (recipes/_dist_utils.py) and normalizes bool | "full" | "selective" — but the platform had no way to set it, and FSDP2Config defaults it to False. Every customization job therefore trained with it off.That blocks porting the NVIDIA Nemotron LoRA cookbooks, which set
activation_checkpointing: true; the Nemotron 3 Super recipe annotates it "reduces peak memory (avoids OOM on 80GB)".Adds the field alongside expert_parallel_size at each layer it already travels through: plugin schema, adapter, service API schema, compiler, task config, and YAML emission.
Typed bool | Literal["full", "selective"] | None to match what Automodel parses, rather than a plain bool, so selective checkpointing — the cheaper mode — stays reachable.
Defaults to None and is omitted from the emitted YAML when unset, so existing jobs compile byte-identically; the 15 golden contract configs are unchanged. An explicit false is still emitted, to distinguish "user turned it off" from "not configured".
Note the pre-existing embedding branch it now precedes is unreachable:
embedding_configis a bare EmbeddingConfig() that is never populated, sodo_gradient_checkpointingis always False. It is left in place as the fallback in case that config is ever wired up.Not included: OpenAPI regeneration (
make refresh-openapi) and the Stainless SDK sync, which need to run before the field reaches the CLI, Python SDK, or Studio.Summary
Related Issue
Changes
Type of Change
Quality Gates
Verification
Signed-off-by:traileruv run pre-commit run -apasses, or any blocked checks are identified belowTargeted validation:
Summary by CodeRabbit
New Features
Bug Fixes
Tests