Skip to content

fix(train): preflight checkpoint publication on every rank #6089

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
OutisLi wants to merge 2 commits into
base: master
Choose a base branch
from
Open

fix(train): preflight checkpoint publication on every rank #6089

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c90c1ae
fix(train): preflight checkpoint publication on every rank
OutisLi Oct 19, 2026
04f93d9
test(train): remove unused checkpoint path assignments
OutisLi Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 52 additions & 72 deletions deepmd/pt/train/training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1732,6 +1732,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
step_id=_step_id,
display_step=display_step_id,
lr=cur_lr,
prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint,
save_checkpoint=(
self.save_model_merged
if self._lora_enabled
Expand All @@ -1743,6 +1744,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
step_id=_step_id,
display_step=display_step_id,
lr=cur_lr,
prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint,
save_checkpoint=(
self.save_ema_model_merged
if self._lora_enabled
Expand All @@ -1756,27 +1758,9 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
and _step_id != self.start_step
)
if should_save_checkpoint:
# Abort before writing if any gradient norm since the previous
# checkpoint was non-finite.
self.nonfinite_grad_guard.raise_if_nonfinite(
self.wrapper.named_parameters
self._save_training_checkpoint(
step=_step_id, display_step=display_step_id, lr=cur_lr
)
if should_save_checkpoint and (
self.sharding.enabled or self.rank == 0 or dist.get_rank() == 0
):
# Handle the case if rank 0 aborted and re-assigned
self.latest_model = self.ckpt_store.path_for(display_step_id)
self.save_model(self.latest_model, lr=cur_lr, step=_step_id)
if self.rank == 0 or dist.get_rank() == 0:
log.info(f"Saved model to {self.latest_model}")
self.ckpt_store.publish(self.latest_model)
if self.model_ema is not None:
self.latest_ema_model = self.ema_ckpt_store.path_for(
display_step_id
)
self.save_ema_model(self.latest_ema_model, lr=cur_lr, step=_step_id)
if self.rank == 0 or dist.get_rank() == 0:
self.ema_ckpt_store.publish(self.latest_ema_model)

# tensorboard
if self.enable_tensorboard and (
Expand Down Expand Up @@ -1844,62 +1828,32 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
self._discarded_training_batches,
)

if (
self.change_bias_after_training
and self.num_steps > self.start_step
and (self.rank == 0 or dist.get_rank() == 0)
):
if not self.multi_task:
self.model = model_change_out_bias(
self.model,
self.get_sample_func,
_bias_adjust_mode="change-by-statistic",
)
else:
for model_key in self.model_keys:
self.model[model_key] = model_change_out_bias(
self.model[model_key],
self.get_sample_func[model_key],
if self.change_bias_after_training and self.num_steps > self.start_step:
if self.rank == 0 or dist.get_rank() == 0:
if not self.multi_task:
self.model = model_change_out_bias(
self.model,
self.get_sample_func,
_bias_adjust_mode="change-by-statistic",
)
self.latest_model = self.ckpt_store.path_for(self.num_steps)
cur_lr = self.lr_schedule.value(self.num_steps - 1)
self.save_model(self.latest_model, lr=cur_lr, step=self.num_steps - 1)
log.info(f"Saved model to {self.latest_model}")
self.ckpt_store.publish(self.latest_model)
if self.model_ema is not None:
self.latest_ema_model = self.ema_ckpt_store.path_for(self.num_steps)
self.save_ema_model(
self.latest_ema_model,
lr=cur_lr,
step=self.num_steps - 1,
)
self.ema_ckpt_store.publish(self.latest_ema_model)
else:
for model_key in self.model_keys:
self.model[model_key] = model_change_out_bias(
self.model[model_key],
self.get_sample_func[model_key],
_bias_adjust_mode="change-by-statistic",
)
self._save_training_checkpoint(
step=self.num_steps - 1,
display_step=self.num_steps,
lr=self.lr_schedule.value(self.num_steps - 1),
)

if self.num_steps == 0 and self.sharding.enabled:
# ZeRO-1 / FSDP: all ranks participate in save_model (collective op)
self.latest_model = self.ckpt_store.path_for(0)
self.save_model(self.latest_model, lr=0, step=0)
if self.model_ema is not None:
self.latest_ema_model = self.ema_ckpt_store.path_for(0)
self.save_ema_model(self.latest_ema_model, lr=0, step=0)

if (
self.rank == 0 or dist.get_rank() == 0
): # Handle the case if rank 0 aborted and re-assigned
if self.num_steps == 0:
if not self.sharding.enabled:
# When num_steps is 0, the checkpoint is never saved in the loop
self.latest_model = self.ckpt_store.path_for(0)
self.save_model(self.latest_model, lr=0, step=0)
if self.model_ema is not None:
self.latest_ema_model = self.ema_ckpt_store.path_for(0)
self.save_ema_model(self.latest_ema_model, lr=0, step=0)
log.info(f"Saved model to {self.latest_model}")
self.ckpt_store.publish(self.latest_model)
if self.model_ema is not None:
self.ema_ckpt_store.publish(self.latest_ema_model)
if self.num_steps == 0:
self._save_training_checkpoint(step=0, display_step=0, lr=0.0)

if self.rank == 0 or dist.get_rank() == 0:
# Handle the case if rank 0 aborted and was reassigned.
if self.timing_in_training:
average_message = self.step_timer.format_average()
if average_message is not None:
Expand Down Expand Up @@ -1951,6 +1905,32 @@ def _close_lmdb_loaders(self) -> None:
for dataset in datasets.values():
dataset.close()

def _save_training_checkpoint(
self, *, step: int, display_step: int, lr: float
) -> None:
"""Publish live and optional EMA checkpoints after one all-rank check."""
self.ensure_finite_gradients_for_checkpoint()
# Preserve chief ownership if rank 0 was reassigned.
is_chief = self.rank == 0 or dist.get_rank() == 0
if not is_chief and not self.sharding.enabled:
return
self.latest_model = self.ckpt_store.path_for(display_step)
self.save_model(self.latest_model, lr=lr, step=step)
if is_chief:
log.info(f"Saved model to {self.latest_model}")
self.ckpt_store.publish(self.latest_model)
if self.model_ema is not None:
self.latest_ema_model = self.ema_ckpt_store.path_for(display_step)
self.save_ema_model(self.latest_ema_model, lr=lr, step=step)
if is_chief:
self.ema_ckpt_store.publish(self.latest_ema_model)

def ensure_finite_gradients_for_checkpoint(self) -> None:
"""Check the accumulated gradient norm on all ranks before publication."""
self.nonfinite_grad_guard.raise_if_nonfinite(
self.wrapper.named_parameters, synchronize=self.is_distributed
)

def _collect_checkpoint_states(
self,
*,
Expand Down
16 changes: 13 additions & 3 deletions deepmd/pt_expt/train/gradient.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
)

import torch
import torch.distributed as dist

if TYPE_CHECKING:
from collections.abc import (
Expand Down Expand Up @@ -123,6 +124,8 @@ def update(self, total_norm: torch.Tensor) -> None:
def raise_if_nonfinite(
self,
named_parameters: Callable[[], Iterable[tuple[str, torch.nn.Parameter]]],
*,
synchronize: bool = False,
) -> None:
"""
Raise if any norm accumulated since the previous call was non-finite.
Expand All @@ -134,15 +137,22 @@ def raise_if_nonfinite(
----------
named_parameters : Callable[[], Iterable[tuple[str, torch.nn.Parameter]]]
Accessor for the model's named parameters, consulted only when raising.
synchronize : bool, optional
Agree on divergence across all ranks before raising or entering
checkpoint collectives. Every rank must participate when enabled.

Raises
------
RuntimeError
If a non-finite gradient norm was recorded.
"""
if self._nonfinite is None:
return
diverged = bool(self._nonfinite)
diverged = self._nonfinite is not None and bool(self._nonfinite)
if synchronize:
# Agree before diagnostics too: inspecting sharded gradients can
# itself require collectives, so all ranks must take that path.
gathered = [False] * dist.get_world_size()
dist.all_gather_object(gathered, diverged)
diverged = any(gathered)
self._nonfinite = None
if diverged:
raise_nonfinite_gradient_norm(named_parameters())
Expand Down
20 changes: 13 additions & 7 deletions deepmd/pt_expt/train/training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3055,13 +3055,19 @@ def _broadcast_value_from_rank0(self, value: Any) -> Any:

@property
def checkpoint_is_collective(self) -> bool:
"""Whether assembling a checkpoint needs every rank."""
return self.sharding.enabled
"""Whether checkpoint preflight needs every rank."""
return self.is_distributed

def ensure_finite_gradients_for_checkpoint(self) -> None:
"""Check the accumulated gradient norm on all ranks before publication."""
self.nonfinite_grad_guard.raise_if_nonfinite(
self.wrapper.named_parameters, synchronize=self.is_distributed
)

def save_checkpoint(self, step: int) -> None:
# Abort before writing if any gradient norm since the previous
# checkpoint was non-finite, so a diverged interval is not persisted.
self.nonfinite_grad_guard.raise_if_nonfinite(self.wrapper.named_parameters)
self.ensure_finite_gradients_for_checkpoint()
if self.rank != 0 and not self.sharding.enabled:
return
ckpt_path = self.ckpt_store.path_for(step)
self._save_checkpoint_to_path(ckpt_path, step=step)
if self.rank == 0:
Expand Down Expand Up @@ -3309,8 +3315,7 @@ def run(self) -> None:
super().run(self.training_tasks)
if self.change_bias_after_training and self.num_steps > self.start_step:
self._change_bias_after_training()
if self.rank_context.is_chief:
self.save_checkpoint(self.num_steps)
self.save_checkpoint(self.num_steps)
finally:
self._close_data_systems()
if self.rank_context.is_chief:
Expand Down Expand Up @@ -3363,6 +3368,7 @@ def run_full_validation(
display_step=display_step,
lr=learning_rate,
save_checkpoint=save_checkpoint,
prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint,
)

def select_task(self, tasks: TrainingTaskCollection) -> TrainingTask:
Expand Down
59 changes: 39 additions & 20 deletions deepmd/pt_expt/train/validation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
)
from dataclasses import (
dataclass,
replace,
)
from pathlib import (
Path,
Expand Down Expand Up @@ -315,8 +316,14 @@ def run(
display_step: int,
lr: float,
save_checkpoint: Any,
prepare_checkpoint: Callable[[], None] | None = None,
) -> FullValidationResult | None:
"""Run full validation if the current step is due."""
"""Run full validation if the current step is due.

``prepare_checkpoint`` runs on every rank only for a selected checkpoint.
It must agree on distributed failures before returning or raising, so
no peer enters checkpoint collectives after another rejects the save.
"""
if not self.should_run(display_step):
return None

Expand All @@ -327,10 +334,15 @@ def run(
caught_exception: Exception | None = None
error_message = None
save_path = [None]
proposed_records = self.topk_records
if self.rank == 0:
try:
result = self._evaluate(display_step)
save_path[0] = result.saved_best_path
proposed_records, save_path[0] = self._propose_best_state(
display_step=display_step,
selected_metric_value=result.selected_metric_value,
)
result = replace(result, saved_best_path=save_path[0])
except Exception as exc:
caught_exception = exc
error_message = (
Expand All @@ -343,6 +355,15 @@ def run(
if self.is_distributed:
dist.broadcast_object_list(save_path, src=0)

if save_path[0] is not None and prepare_checkpoint is not None:
# All ranks agree before mutating top-K metadata or entering a
# serializer that may gather sharded model/optimizer state.
prepare_checkpoint()

if self.rank == 0:
self.topk_records = proposed_records
self._sync_state_store()

if save_path[0] is not None:
try:
# Assembling a checkpoint from shards is collective, so every
Expand Down Expand Up @@ -384,7 +405,7 @@ def run(
return result if self.rank == 0 else None

def _evaluate(self, display_step: int) -> FullValidationResult:
"""Evaluate all validation systems and update best state."""
"""Evaluate all validation systems without publishing checkpoint state."""
# === Step 1. Switch to Evaluation Mode ===
was_training = bool(getattr(self.model, "training", True))
self.model.eval()
Expand All @@ -408,18 +429,13 @@ def _evaluate(self, display_step: int) -> FullValidationResult:
f"validation dataset: {self.metric_name.upper()}."
)

# === Step 3. Update Best Tracking ===
selected_metric_value = float(metrics[self.metric_key])
saved_best_path = self._update_best_state(
display_step=display_step,
selected_metric_value=selected_metric_value,
)
return FullValidationResult(
display_step=display_step,
metrics=metrics,
selected_metric_key=self.metric_key,
selected_metric_value=selected_metric_value,
saved_best_path=saved_best_path,
saved_best_path=None,
)

def evaluate_all_systems(self) -> dict[str, float]:
Expand Down Expand Up @@ -696,13 +712,13 @@ def predict_batch(
prediction["virial"] = batch_prediction["virial"]
return prediction

def _update_best_state(
def _propose_best_state(
self,
*,
display_step: int,
selected_metric_value: float,
) -> str | None:
"""Update the top-K records and return the checkpoint path to save."""
) -> tuple[list[BestCheckpointRecord], str | None]:
"""Propose top-K records and a save path without changing published state."""
candidate = BestCheckpointRecord(
metric=selected_metric_value,
step=display_step,
Expand All @@ -714,14 +730,17 @@ def _update_best_state(
updated_records.sort()
updated_records = updated_records[: self.max_best_ckpt]
if candidate not in updated_records:
return None

self.topk_records = updated_records
self._sync_state_store()
if not self.save_best:
return None
candidate_rank = self.topk_records.index(candidate) + 1
return str(self._best_checkpoint_path(display_step, candidate_rank))
return self.topk_records, None
save_path = (
str(
self._best_checkpoint_path(
display_step, updated_records.index(candidate) + 1
)
)
if self.save_best
else None
)
return updated_records, save_path

def _sync_state_store(self) -> None:
"""Synchronize top-K validation state into the configured state store."""
Expand Down
Loading
Loading