242 lines
9.1 KiB
Python
242 lines
9.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from configs import cfg
|
|
|
|
def checkpoint_rank(path: str) -> tuple[int, int]:
|
|
name = os.path.basename(path)
|
|
prefix, _, raw_value = name.partition("_")
|
|
try:
|
|
value = int(raw_value)
|
|
except ValueError:
|
|
value = -1
|
|
if prefix == "epoch":
|
|
return (2, value)
|
|
if prefix == "step":
|
|
return (1, value)
|
|
return (0, value)
|
|
|
|
def find_latest_training_checkpoint(output_dir: str) -> str | None:
|
|
candidates = []
|
|
for pattern in ("epoch_*", "step_*"):
|
|
candidates.extend(str(path) for path in Path(output_dir).glob(pattern) if path.is_dir())
|
|
if not candidates:
|
|
return None
|
|
return max(candidates, key=checkpoint_rank)
|
|
|
|
def load_trainer_state(checkpoint_dir: str, log) -> dict:
|
|
state_path = os.path.join(checkpoint_dir, "trainer_state.json")
|
|
if os.path.exists(state_path):
|
|
try:
|
|
with open(state_path, "r", encoding="utf-8") as f:
|
|
state = json.load(f)
|
|
if isinstance(state, dict):
|
|
return state
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
log.warning(f"Could not read trainer_state.json from {checkpoint_dir}: {exc}")
|
|
|
|
name = os.path.basename(checkpoint_dir)
|
|
prefix, _, raw_value = name.partition("_")
|
|
try:
|
|
value = int(raw_value)
|
|
except ValueError:
|
|
value = 0
|
|
if prefix == "epoch":
|
|
return {
|
|
"checkpoint_type": "epoch",
|
|
"start_epoch": value,
|
|
"global_step": 0,
|
|
"micro_step_global": 0,
|
|
"next_batch_in_epoch": 0,
|
|
}
|
|
if prefix == "step":
|
|
return {
|
|
"checkpoint_type": "step",
|
|
"start_epoch": 0,
|
|
"global_step": value,
|
|
"micro_step_global": 0,
|
|
"next_batch_in_epoch": 0,
|
|
}
|
|
return {}
|
|
|
|
def packing_checkpoint_metadata(enabled: bool, pack_length: int | None, max_seq_len: int) -> dict[str, int | bool | None]:
|
|
return {
|
|
"sequence_packing_enabled": bool(enabled),
|
|
"sequence_packing_pack_length": int(pack_length) if enabled and pack_length is not None else None,
|
|
"data_max_seq_len": int(max_seq_len),
|
|
}
|
|
|
|
def validate_resume_packing_state(
|
|
trainer_state: dict,
|
|
*,
|
|
enabled: bool,
|
|
pack_length: int,
|
|
max_seq_len: int,
|
|
log,
|
|
) -> None:
|
|
checkpoint_enabled = bool(trainer_state.get("sequence_packing_enabled", False))
|
|
if checkpoint_enabled != bool(enabled):
|
|
log.error(
|
|
"Checkpoint sequence-packing state does not match the current run: "
|
|
f"checkpoint={checkpoint_enabled}, current={bool(enabled)}."
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
if checkpoint_enabled:
|
|
checkpoint_pack_length = trainer_state.get("sequence_packing_pack_length")
|
|
try:
|
|
checkpoint_pack_length = int(checkpoint_pack_length)
|
|
except (TypeError, ValueError):
|
|
log.error("Checkpoint is missing a valid sequence_packing_pack_length value.")
|
|
raise SystemExit(1)
|
|
if checkpoint_pack_length != int(pack_length):
|
|
log.error(
|
|
"Checkpoint pack length does not match the current run: "
|
|
f"checkpoint={checkpoint_pack_length}, current={int(pack_length)}."
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
checkpoint_max_seq_len = trainer_state.get("data_max_seq_len")
|
|
if checkpoint_max_seq_len is not None:
|
|
try:
|
|
checkpoint_max_seq_len = int(checkpoint_max_seq_len)
|
|
except (TypeError, ValueError):
|
|
log.error("Checkpoint is missing a valid data_max_seq_len value.")
|
|
raise SystemExit(1)
|
|
if checkpoint_max_seq_len != int(max_seq_len):
|
|
log.error(
|
|
"Checkpoint max sequence length does not match the current run: "
|
|
f"checkpoint={checkpoint_max_seq_len}, current={int(max_seq_len)}."
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
def save_checkpoint(
|
|
model,
|
|
tokenizer,
|
|
output_dir: str,
|
|
tag: str,
|
|
logger,
|
|
*,
|
|
scheduler=None,
|
|
trainer_state: dict | None = None,
|
|
) -> str:
|
|
save_dir = os.path.join(output_dir, tag)
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
save_start = time.time()
|
|
logger.info(f"[CKPT] Saving {tag} -> {save_dir}/")
|
|
|
|
model_to_save = model.module if hasattr(model, "module") else model
|
|
if hasattr(model_to_save, "_orig_mod"):
|
|
model_to_save = model_to_save._orig_mod
|
|
|
|
model_to_save.config.save_pretrained(save_dir)
|
|
tokenizer.save_pretrained(save_dir)
|
|
|
|
try:
|
|
from safetensors.torch import save_file
|
|
state_dict = {k: v.contiguous().cpu() for k, v in model_to_save.state_dict().items()}
|
|
save_file(state_dict, os.path.join(save_dir, "model.safetensors"))
|
|
logger.info("[CKPT] Saved via safetensors")
|
|
except ImportError:
|
|
torch.save(model_to_save.state_dict(), os.path.join(save_dir, "pytorch_model.bin"))
|
|
logger.info("[CKPT] Saved via torch.save")
|
|
|
|
if scheduler is not None:
|
|
torch.save(scheduler.state_dict(), os.path.join(save_dir, "scheduler.pt"))
|
|
|
|
if trainer_state is not None:
|
|
trainer_state = dict(trainer_state)
|
|
trainer_state.setdefault("tag", tag)
|
|
trainer_state.setdefault("saved_at", time.strftime("%Y-%m-%d %H:%M:%S %Z"))
|
|
with open(os.path.join(save_dir, "trainer_state.json"), "w", encoding="utf-8") as f:
|
|
json.dump(trainer_state, f, indent=2)
|
|
|
|
size_mb = sum(f.stat().st_size for f in Path(save_dir).rglob("*") if f.is_file()) / 1e6
|
|
save_elapsed = time.time() - save_start
|
|
logger.info(f"[CKPT] {tag} -> {save_dir}/ ({size_mb:.0f} MB, {save_elapsed:.1f}s)")
|
|
return save_dir
|
|
|
|
def read_env_flag(name: str, default: bool = False) -> bool:
|
|
raw = os.environ.get(name)
|
|
if raw is None:
|
|
return default
|
|
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
def hub_upload_strict() -> bool:
|
|
strict = getattr(getattr(cfg, "hub", None), "hub_upload_strict", None)
|
|
if strict is None:
|
|
return read_env_flag("QUINTUS_HUB_UPLOAD_STRICT", False)
|
|
return bool(strict)
|
|
|
|
def should_upload_checkpoint_tag(tag: str) -> bool:
|
|
upload_regular = getattr(getattr(cfg, "hub", None), "upload_kd_checkpoints", False) or read_env_flag("QUINTUS_UPLOAD_KD_CHECKPOINTS", False)
|
|
upload_steps = getattr(getattr(cfg, "hub", None), "upload_step_checkpoints", False) or read_env_flag("QUINTUS_UPLOAD_STEP_CHECKPOINTS", False)
|
|
upload_last = getattr(getattr(cfg, "hub", None), "upload_last_checkpoint", False) or read_env_flag("QUINTUS_UPLOAD_LAST_CHECKPOINT", False)
|
|
if tag.startswith("step_"):
|
|
return upload_steps
|
|
if tag.startswith("epoch_"):
|
|
return upload_regular
|
|
if tag == "best":
|
|
return upload_regular
|
|
if tag == "last":
|
|
return upload_last or upload_regular
|
|
return False
|
|
|
|
def maybe_upload_checkpoint(checkpoint_dir: str, tag: str, logger) -> None:
|
|
if not should_upload_checkpoint_tag(tag):
|
|
return
|
|
|
|
token = os.environ.get("HF_TOKEN") or getattr(cfg.hub, "token", None)
|
|
if not token:
|
|
msg = "HF checkpoint upload requested, but HF_TOKEN/cfg.hub.token is missing"
|
|
strict = hub_upload_strict()
|
|
if strict:
|
|
raise RuntimeError(msg)
|
|
logger.warning(f"[CKPT] {msg}; continuing without remote backup")
|
|
return
|
|
|
|
repo_id = getattr(getattr(cfg, "hub", None), "repo_id", None) or os.environ.get("QUINTUS_HUB_REPO_ID") or f"{cfg.hub.username}/{cfg.hub.repo_name}"
|
|
base_path = getattr(getattr(cfg, "hub", None), "ckpt_path_in_repo", None) or os.environ.get("KD_CKPT_PATH_IN_REPO", "models/online_kd_3b_05b_ep3_B200_20260601")
|
|
base_path = base_path.strip("/")
|
|
path_in_repo = f"{base_path}/{tag}"
|
|
commit_prefix = getattr(getattr(cfg, "hub", None), "commit_message_prefix", None) or os.environ.get(
|
|
"KD_COMMIT_MESSAGE_PREFIX",
|
|
"Online KD 8B->1.7B Run",
|
|
)
|
|
commit_message = os.environ.get("KD_COMMIT_MESSAGE") or f"{commit_prefix}: upload {tag}"
|
|
upload_start = time.time()
|
|
size_mb = sum(f.stat().st_size for f in Path(checkpoint_dir).rglob("*") if f.is_file()) / 1e6
|
|
strict = hub_upload_strict()
|
|
logger.info(
|
|
f"[CKPT] Uploading {tag} -> {repo_id}/{path_in_repo} "
|
|
f"({size_mb:.0f} MB, strict={strict})"
|
|
)
|
|
logger.info(f"[CKPT] Commit: {commit_message}")
|
|
|
|
try:
|
|
from huggingface_hub import HfApi
|
|
|
|
api = HfApi(token=token)
|
|
api.create_repo(repo_id=repo_id, repo_type="dataset", private=True, exist_ok=True)
|
|
api.upload_folder(
|
|
folder_path=checkpoint_dir,
|
|
repo_id=repo_id,
|
|
path_in_repo=path_in_repo,
|
|
repo_type="dataset",
|
|
commit_message=commit_message,
|
|
ignore_patterns=["*.tmp", "*.log", "__pycache__/*"],
|
|
)
|
|
upload_elapsed = time.time() - upload_start
|
|
logger.info(f"[CKPT] Uploaded {tag} to HF Hub in {upload_elapsed / 60:.1f}m")
|
|
except Exception as exc:
|
|
msg = f"HF checkpoint upload failed for {tag}: {exc}"
|
|
if hub_upload_strict():
|
|
raise RuntimeError(msg) from exc
|
|
logger.warning(f"[CKPT] {msg}; continuing because hub upload strict mode is disabled")
|