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/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 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