Make default checkpoint retention conservative

experiments/defaults.py

@@ -116,7 +116,7 @@ def _truncate_wandb_name(name: str) -> str:
     return name
 
 
-def _resolve_hf_export_steps(steps_per_hf_export: int | None, steps_per_export: int) -> int | None:
+def _resolve_hf_export_steps(steps_per_hf_export: int | None, steps_per_export: int | None) -> int | None:
     """Resolve the HF export step interval: None means same as checkpoint, -1 means disabled."""
     if steps_per_hf_export is None:
         return steps_per_export
@@ -125,6 +125,17 @@ def _resolve_hf_export_steps(steps_per_hf_export: int | None, steps_per_export:
     return steps_per_hf_export
 
 
+def _checkpoint_keep(steps_per_export: int | None) -> list[dict]:
+    """Build the `keep` list for `CheckpointerConfig`.
+
+    None means keep no permanent intermediate checkpoints (only the final checkpoint
+    is saved at end-of-training, plus a rolling temporary checkpoint for resumption).
+    """
+    if steps_per_export is None:
+        return []
+    return [dict(every=steps_per_export)]
+
+
 def _validate_train_length(train_seq_len: int | None, model_config: LmConfig) -> int:
     """Resolve and validate the training sequence length against the model's max."""
     actual = unwrap_versioned_value(model_config)
@@ -409,7 +420,7 @@ def default_train(
             steps_per_eval=train_config.steps_per_eval if train_config.steps_per_eval is not None else 1000,
             checkpointer=CheckpointerConfig(
                 save_interval=timedelta(minutes=10),
-                keep=[dict(every=steps_per_export)],
+                keep=_checkpoint_keep(steps_per_export),
             ),
             model_averaging=model_averaging,
             mesh=MeshConfig(
@@ -641,7 +652,7 @@ def default_dpo(
             steps_per_eval=dpo_config.steps_per_eval,
             checkpointer=CheckpointerConfig(
                 save_interval=timedelta(minutes=10),
-                keep=[dict(every=steps_per_export)],
+                keep=_checkpoint_keep(steps_per_export),
             ),
             model_averaging=None,
             mesh=MeshConfig(

experiments/simple_dpo_config.py

@@ -40,7 +40,9 @@ class SimpleDPOConfig:
     max_grad_norm: float | None = 1
 
     steps_per_eval: int = 1000
-    steps_per_checkpoint: int = 1000
+    steps_per_checkpoint: int | None = None
+    """How often to keep a permanent checkpoint. None (default) keeps only the final
+    checkpoint; rolling temporary checkpoints are still written for resumption."""
     steps_per_hf_export: int = 500
     hf_save_dtype: str | None = None
     hf_generation_eos_token_ids: list[int] | None = None

experiments/simple_sft_config.py

@@ -117,8 +117,9 @@ class SimpleSFTConfig:
     steps_per_eval: int = 1000
     """How often to run validation losses."""
 
-    steps_per_checkpoint: int = 1000
-    """How often to save checkpoints."""
+    steps_per_checkpoint: int | None = None
+    """How often to keep a permanent checkpoint. None (default) keeps only the final
+    checkpoint; rolling temporary checkpoints are still written for resumption."""
 
     steps_per_hf_export: int = 500
     """How often to save HuggingFace checkpoints."""

experiments/simple_train_config.py

@@ -45,7 +45,9 @@ class SimpleTrainConfig:
 
     steps_per_eval: int | None = None
     """how often to run validation losses"""
-    steps_per_export: int = 10000
+    steps_per_export: int | None = None
+    """How often to keep a permanent checkpoint. None (default) keeps only the final
+    checkpoint; rolling temporary checkpoints are still written for resumption."""
     steps_per_task_eval: int | None = None
     """how often to run task evaluations"""
     steps_per_hf_export: int | None = None