From bebd07dcb0fc730377bafb4f1fda51df725c0285 Mon Sep 17 00:00:00 2001 From: Henrique Tolentino Date: Wed, 12 Aug 2026 11:24:51 -0400 Subject: [PATCH 1/2] feat(automodel): expose activation_checkpointing on customization jobs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Signed-off-by: Henrique Tolentino --- .../src/nemo_automodel_plugin/schema.py | 8 +++ .../automodel/src/nmp/automodel/adapter.py | 1 + .../src/nmp/automodel/api/v2/jobs/schemas.py | 8 +++ .../automodel/app/jobs/training/compiler.py | 1 + .../automodel/app/jobs/training/schemas.py | 3 +- .../tasks/training/backends/config.py | 6 +- .../tasks/training/backends/test_config.py | 56 +++++++++++++++++++ services/automodel/tests/test_adapter.py | 44 +++++++++++++++ 8 files changed, 125 insertions(+), 2 deletions(-) diff --git a/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py b/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py index 740787c4f5..70d75cb322 100644 --- a/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py +++ b/plugins/nemo-automodel/src/nemo_automodel_plugin/schema.py @@ -127,6 +127,14 @@ class ParallelismSpec(AutomodelSchema): context_parallel_size: int = Field(default=1, gt=0) expert_parallel_size: int | None = Field(default=None, gt=0) sequence_parallel: bool = Field(default=False, description="Enable sequence parallelism.") + activation_checkpointing: bool | Literal["full", "selective"] | None = Field( + default=None, + 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." + ), + ) class OutputRequest(AutomodelSchema): diff --git a/services/automodel/src/nmp/automodel/adapter.py b/services/automodel/src/nmp/automodel/adapter.py index aa4999828c..e82fa4f479 100644 --- a/services/automodel/src/nmp/automodel/adapter.py +++ b/services/automodel/src/nmp/automodel/adapter.py @@ -80,6 +80,7 @@ def _build_training_block(spec: dict[str, Any]) -> SFTTraining | DistillationTra context_parallel_size=parallelism.get("context_parallel_size", 1), expert_parallel_size=parallelism.get("expert_parallel_size"), sequence_parallel=parallelism.get("sequence_parallel", False), + activation_checkpointing=parallelism.get("activation_checkpointing"), ), "execution_profile": training.get("execution_profile"), } diff --git a/services/automodel/src/nmp/automodel/api/v2/jobs/schemas.py b/services/automodel/src/nmp/automodel/api/v2/jobs/schemas.py index fea67a374e..43329a4cfa 100644 --- a/services/automodel/src/nmp/automodel/api/v2/jobs/schemas.py +++ b/services/automodel/src/nmp/automodel/api/v2/jobs/schemas.py @@ -136,6 +136,14 @@ class ParallelismParams(BaseModel): context_parallel_size: int = Field(default=1, gt=0, description="Context parallel size.") expert_parallel_size: Optional[int] = Field(default=None, gt=0, description="Expert parallel size (MoE models).") sequence_parallel: bool = Field(default=False, description="Enable sequence parallelism.") + activation_checkpointing: Optional[Union[bool, Literal["full", "selective"]]] = Field( + default=None, + description=( + "Recompute activations during the backward pass to reduce peak memory at the cost of " + "speed. 'selective' checkpoints only the most memory-heavy ops. Unset leaves Automodel's " + "default (disabled)." + ), + ) # ============================================================ diff --git a/services/automodel/src/nmp/automodel/app/jobs/training/compiler.py b/services/automodel/src/nmp/automodel/app/jobs/training/compiler.py index 011c92219e..10d107fcaa 100644 --- a/services/automodel/src/nmp/automodel/app/jobs/training/compiler.py +++ b/services/automodel/src/nmp/automodel/app/jobs/training/compiler.py @@ -178,6 +178,7 @@ def compile_training_step( context_parallel_size=p.context_parallel_size, expert_parallel_size=p.expert_parallel_size, sequence_parallel=p.sequence_parallel, + activation_checkpointing=p.activation_checkpointing, ), integrations=job_spec.integrations, output_model=job_spec.output.name, diff --git a/services/automodel/src/nmp/automodel/app/jobs/training/schemas.py b/services/automodel/src/nmp/automodel/app/jobs/training/schemas.py index 40d1550576..7f83fca128 100644 --- a/services/automodel/src/nmp/automodel/app/jobs/training/schemas.py +++ b/services/automodel/src/nmp/automodel/app/jobs/training/schemas.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from enum import Enum -from typing import Optional +from typing import Literal, Optional, Union from nemo_platform_plugin.integrations import IntegrationsSpec from nmp.automodel.app.constants import ( @@ -204,6 +204,7 @@ class ParallelismConfig(BaseModel): context_parallel_size: int = 1 expert_parallel_size: Optional[int] = None sequence_parallel: bool = False + activation_checkpointing: Optional[Union[bool, Literal["full", "selective"]]] = None # === Main Config Fields === model: ModelConfig diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/config.py b/services/automodel/src/nmp/automodel/tasks/training/backends/config.py index d4f31b4b59..240d1d274f 100644 --- a/services/automodel/src/nmp/automodel/tasks/training/backends/config.py +++ b/services/automodel/src/nmp/automodel/tasks/training/backends/config.py @@ -146,7 +146,11 @@ def compile_automodel_config( "ep_size": p.expert_parallel_size, "sequence_parallel": p.sequence_parallel, } - if _is_embedding_model and embedding_config.do_gradient_checkpointing: + # Explicit parallelism setting wins; embedding jobs keep their own toggle as the + # fallback so `do_gradient_checkpointing` behaves as before when nothing is set. + if p.activation_checkpointing is not None: + cfg["distributed"]["activation_checkpointing"] = p.activation_checkpointing + elif _is_embedding_model and embedding_config.do_gradient_checkpointing: cfg["distributed"]["activation_checkpointing"] = True if p.pipeline_parallel_size > 1: cfg["distributed"]["pipeline"] = { diff --git a/services/automodel/tests/tasks/training/backends/test_config.py b/services/automodel/tests/tasks/training/backends/test_config.py index 77391e4fe9..7458811b45 100644 --- a/services/automodel/tests/tasks/training/backends/test_config.py +++ b/services/automodel/tests/tasks/training/backends/test_config.py @@ -414,3 +414,59 @@ def test_compile_uses_fallback_packing_factor_for_schedule(tmp_path: Path) -> No assert compiled["step_scheduler"]["max_steps"] == 2 assert compiled["step_scheduler"]["val_every_steps"] == 1 assert compiled["lr_scheduler"]["lr_warmup_steps"] == 1 + + +def _compile_with_parallelism(tmp_path: Path, **parallelism: Any) -> dict[str, Any]: + """Compile a known-good contract fixture with parallelism overrides applied.""" + fixture = Path(__file__).parents[3] / "contract" / "input_configs" / "llama-3.2-1b" / "llama_3_2_1b_lora.json" + raw = json.loads(fixture.read_text()) + raw.pop("backend", None) + config = TrainingStepConfig.model_validate(raw) + for key, value in parallelism.items(): + setattr(config.parallelism, key, value) + + prepared = PreparedDataset( + merged_dir=tmp_path, + train_file=tmp_path / "train.jsonl", + validation_file=tmp_path / "validation.jsonl", + train_samples=100, + validation_samples=10, + ) + + with ( + patch(f"{CONFIG_MODULE}.prepare_dataset", return_value=prepared), + patch(f"{CONFIG_MODULE}.DatasetValidator"), + patch(f"{CONFIG_MODULE}.estimate_dataset_sequence_lengths", return_value=None), + patch(f"{CONFIG_MODULE}._configure_datasets"), + patch(f"{CONFIG_MODULE}._configure_moe_backend"), + patch(f"{CONFIG_MODULE}.build_wandb_config", return_value=None), + patch(f"{CONFIG_MODULE}.build_mlflow_config", return_value=None), + ): + return compile_automodel_config(config, tmp_path, MagicMock()) + + +class TestActivationCheckpointing: + """Emission of `distributed.activation_checkpointing` for causal-LM jobs. + + Automodel's FSDP2Config defaults this to False, so a key we never emit means + activation checkpointing is off — which is what every non-embedding job used to get. + """ + + def test_omitted_when_unset(self, tmp_path: Path) -> None: + compiled = _compile_with_parallelism(tmp_path) + assert "activation_checkpointing" not in compiled["distributed"] + + def test_emitted_when_enabled(self, tmp_path: Path) -> None: + compiled = _compile_with_parallelism(tmp_path, activation_checkpointing=True) + assert compiled["distributed"]["activation_checkpointing"] is True + + @pytest.mark.parametrize("mode", ["full", "selective"]) + def test_string_modes_pass_through(self, tmp_path: Path, mode: str) -> None: + # Automodel accepts bool | "full" | "selective"; the strings must not be coerced. + compiled = _compile_with_parallelism(tmp_path, activation_checkpointing=mode) + assert compiled["distributed"]["activation_checkpointing"] == mode + + def test_explicit_false_is_emitted(self, tmp_path: Path) -> None: + # Distinct from unset: the user asked for it off, so say so rather than omit. + compiled = _compile_with_parallelism(tmp_path, activation_checkpointing=False) + assert compiled["distributed"]["activation_checkpointing"] is False diff --git a/services/automodel/tests/test_adapter.py b/services/automodel/tests/test_adapter.py index 4305c20d85..f9467c87b6 100644 --- a/services/automodel/tests/test_adapter.py +++ b/services/automodel/tests/test_adapter.py @@ -225,3 +225,47 @@ def test_adapter_integrations_from_automodel_job_output() -> None: assert spec.integrations is not None assert spec.integrations.wandb is not None assert spec.integrations.wandb.project == "plugin-project" + + +def test_adapter_plumbs_activation_checkpointing() -> None: + """`parallelism.activation_checkpointing` must survive into the v2 training spec.""" + spec = automodel_spec_to_compiler_output( + { + "model": "meta/llama", + "dataset": {"training": "default/train"}, + "training": {"training_type": "sft", "finetuning_type": "lora"}, + "parallelism": {"num_gpus_per_node": 8, "activation_checkpointing": True}, + "output": {"name": "out", "type": "adapter", "fileset": "out-fs"}, + }, + ) + assert isinstance(spec.training, SFTTraining) + assert spec.training.parallelism.activation_checkpointing is True + + +def test_adapter_activation_checkpointing_accepts_selective() -> None: + """Automodel takes bool | 'full' | 'selective'; the string modes must pass through.""" + spec = automodel_spec_to_compiler_output( + { + "model": "meta/llama", + "dataset": {"training": "default/train"}, + "training": {"training_type": "sft", "finetuning_type": "lora"}, + "parallelism": {"num_gpus_per_node": 8, "activation_checkpointing": "selective"}, + "output": {"name": "out", "type": "adapter", "fileset": "out-fs"}, + }, + ) + assert isinstance(spec.training, SFTTraining) + assert spec.training.parallelism.activation_checkpointing == "selective" + + +def test_adapter_activation_checkpointing_defaults_to_none() -> None: + """Unset means "don't emit", preserving Automodel's own default of disabled.""" + spec = automodel_spec_to_compiler_output( + { + "model": "meta/llama", + "dataset": {"training": "default/train"}, + "training": {"training_type": "sft", "finetuning_type": "lora"}, + "output": {"name": "out", "type": "adapter", "fileset": "out-fs"}, + }, + ) + assert isinstance(spec.training, SFTTraining) + assert spec.training.parallelism.activation_checkpointing is None From f0d540e5b5f5334287ca91ba7915a92d8dba28a2 Mon Sep 17 00:00:00 2001 From: Henrique Tolentino Date: Wed, 12 Aug 2026 13:05:19 -0400 Subject: [PATCH 2/2] chore(customizer): regenerate OpenAPI spec for activation_checkpointing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Signed-off-by: Henrique Tolentino --- plugins/nemo-customizer/openapi/openapi.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/plugins/nemo-customizer/openapi/openapi.yaml b/plugins/nemo-customizer/openapi/openapi.yaml index 50421c9a38..6cef93cbf8 100644 --- a/plugins/nemo-customizer/openapi/openapi.yaml +++ b/plugins/nemo-customizer/openapi/openapi.yaml @@ -1572,6 +1572,17 @@ components: title: Sequence Parallel description: Enable sequence parallelism. default: false + activation_checkpointing: + anyOf: + - type: boolean + - type: string + enum: + - full + - selective + title: Activation Checkpointing + 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. additionalProperties: false type: object title: AutomodelParallelismSpec