File size: 8,296 Bytes
b6a9d87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from pathlib import Path
import shutil
import torch
import torch.distributed as dist
from transformers import Trainer, TrainerCallback
from transformers.trainer_callback import TrainerControl, TrainerState
from transformers.training_args import TrainingArguments
from gr00t.utils.dist_utils import run_on_rank0, run_or_wait_on_rank0
logger = logging.getLogger(__name__)
def _broadcast_save_decision(save_flag: int, metric_value: float) -> tuple[int, float]:
"""Broadcast rank-0's `(save_flag, metric_value)` decision to every rank.
HF's eval loop populates ``metrics`` only on rank-0 in some configurations,
so the "should we save?" branch must be decided there and synced. Without
this, non-rank-0 ranks would skip the collective ``Trainer.save_model``
call below and rank-0 would deadlock inside the consolidated-gather.
Returns the broadcast pair on every rank. Single-rank or non-distributed
case: returns the input unchanged.
"""
if not (dist.is_available() and dist.is_initialized()):
return save_flag, metric_value
backend = dist.get_backend()
device = (
torch.device(f"cuda:{torch.cuda.current_device()}")
if backend == "nccl"
else torch.device("cpu")
)
# Use float64 so the metric reaches every rank bit-for-bit identical to
# rank-0 — the value is interpolated into the checkpoint directory name
# and must agree across ranks.
payload = torch.tensor(
[float(save_flag), float(metric_value)], device=device, dtype=torch.float64
)
dist.broadcast(payload, src=0)
return int(payload[0].item()), float(payload[1].item())
class CheckpointFormatCallback(TrainerCallback):
"""This callback format checkpoint to make them standalone. For now, it copies all config
files to /checkpoint-{step}/experiment_cfg/:
- conf.yaml
- initial_actions.npz
- metadata.json
"""
def __init__(
self,
run_name: str,
exp_cfg_dir: Path | None = None,
processor_dir: Path | None = None,
):
"""
Args:
run_name: Name of the experiment run
exp_cfg_dir: Path to the directory containing all experiment metadata
"""
self.exp_cfg_dir = exp_cfg_dir
self.processor_dir = processor_dir
def on_save(self, args, state, control, **kwargs):
"""Called after the trainer saves a checkpoint."""
if state.is_world_process_zero:
checkpoint_dir = Path(args.output_dir) / f"checkpoint-{state.global_step}"
# Copy experiment config directory if provided
if self.exp_cfg_dir is not None:
exp_cfg_dst = checkpoint_dir / self.exp_cfg_dir.name
if self.exp_cfg_dir.exists():
print(
f"Copying experiment config directory {self.exp_cfg_dir} to {exp_cfg_dst}"
)
shutil.copytree(self.exp_cfg_dir, exp_cfg_dst, dirs_exist_ok=True)
# Copy processor directory if provided
if self.processor_dir is not None:
if self.processor_dir.exists():
print(f"Copying processor directory {self.processor_dir} to {checkpoint_dir}")
shutil.copytree(self.processor_dir, checkpoint_dir, dirs_exist_ok=True)
# Copy wandb_config.json if provided
wandb_config_src = Path(args.output_dir) / "wandb_config.json"
wandb_config_dst = checkpoint_dir / "wandb_config.json"
if wandb_config_src.exists():
print(f"Copying wandb_config.json from {wandb_config_src} to {wandb_config_dst}")
shutil.copy2(wandb_config_src, wandb_config_dst)
class BestMetricCheckpointCallback(TrainerCallback):
"""Save a copy of the model whenever an evaluation metric improves.
Works under DDP, DeepSpeed ZeRO 1/2/3, and PyTorch FSDP: the save is
delegated to ``Trainer.save_model``, which handles parameter
consolidation under sharded backends. The callback's own job is to
agree across ranks on whether this eval round should save, and to
manage the best-checkpoint directory (mkdir, optional ``exp_cfg_dir``
copy, previous-best cleanup).
"""
def __init__(
self,
metric_name: str,
trainer: Trainer,
*,
greater_is_better: bool = True,
exp_cfg_dir: Path | None = None,
):
"""
Args:
metric_name: Key in the eval ``metrics`` dict to track.
trainer: The owning ``Trainer``; needed to call
``trainer.save_model`` (the sharded-aware save path).
greater_is_better: True when the metric should be maximized.
exp_cfg_dir: Directory copied alongside each best checkpoint.
"""
self.metric_name = metric_name
self.greater_is_better = greater_is_better
self.best_metric = -float("inf") if greater_is_better else float("inf")
self.exp_cfg_dir = exp_cfg_dir
self._best_checkpoint_dir = None
self._trainer = trainer
def on_evaluate(
self,
args: TrainingArguments,
state: TrainerState,
control: TrainerControl,
metrics,
model,
**kwargs,
):
save_flag = 0
metric_value = 0.0
if state.is_world_process_zero and metrics is not None:
current = metrics.get(self.metric_name, None)
if current is not None:
is_better = (
current > self.best_metric
if self.greater_is_better
else current < self.best_metric
)
if is_better:
save_flag = 1
metric_value = float(current)
save_flag, metric_value = _broadcast_save_decision(save_flag, metric_value)
if save_flag == 0:
return
self.best_metric = metric_value
best_checkpoint_dir = (
Path(args.output_dir)
/ f"checkpoint-{state.global_step}-best-{self.metric_name}_{metric_value}"
)
run_on_rank0(best_checkpoint_dir.mkdir, exist_ok=True, label="best_checkpoint.mkdir")
# Collective on every rank: gathers params under ZeRO-3 / FSDP, then
# writes from rank-0. Calling this only on rank-0 would deadlock the
# rest of the process group inside the gather.
self._trainer.save_model(str(best_checkpoint_dir))
with run_or_wait_on_rank0(label="best_checkpoint.copy") as is_rank0:
if is_rank0:
if self.exp_cfg_dir is not None and self.exp_cfg_dir.exists():
exp_cfg_dst = best_checkpoint_dir / self.exp_cfg_dir.name
logger.info(
"Copying experiment config directory %s to %s",
self.exp_cfg_dir,
exp_cfg_dst,
)
shutil.copytree(self.exp_cfg_dir, exp_cfg_dst, dirs_exist_ok=True)
logger.info(
"Best checkpoint saved to %s with metric %s = %s",
best_checkpoint_dir,
self.metric_name,
metric_value,
)
if (
self._best_checkpoint_dir is not None
and Path(self._best_checkpoint_dir).exists()
):
shutil.rmtree(self._best_checkpoint_dir)
self._best_checkpoint_dir = str(best_checkpoint_dir)
|