File size: 5,126 Bytes
df529cc | 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 | #!/usr/bin/env python3
"""Create a self-contained OPD-only model-only recovery recipe.
DeepSpeed ZeRO-2 optimizer partitions are tied to the original world size, so
the four-rank checkpoint cannot be passed to ``Trainer.train(resume...)`` on
another world size. The student model itself is portable. This utility validates a
base recipe, replaces only the student source and output directory, and writes
the fully resolved YAML consumed by the new run.
"""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
from typing import Any
import yaml
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from config.loader import load_config
def _require_complete_student_checkpoint(path: Path) -> Path:
checkpoint = path.expanduser().resolve(strict=False)
if not checkpoint.is_dir():
raise FileNotFoundError(f"Student initialization checkpoint is not a directory: {checkpoint}")
if not (checkpoint / "config.json").is_file():
raise FileNotFoundError(f"Student initialization checkpoint lacks config.json: {checkpoint}")
weight_files = tuple(checkpoint.glob("*.safetensors")) + tuple(checkpoint.glob("*.bin"))
if not weight_files:
raise FileNotFoundError(
f"Student initialization checkpoint lacks model weights (*.safetensors or *.bin): {checkpoint}"
)
return checkpoint
def _materialize(
*,
base_config: Path,
student_checkpoint: Path,
output_dir: Path,
output_config: Path,
num_gpus: int,
gradient_accumulation_steps: int | None,
) -> dict[str, Any]:
"""Return and atomically write the resolved model-only recovery recipe."""
config = load_config(str(base_config))
checkpoint = _require_complete_student_checkpoint(student_checkpoint)
run_output = output_dir.expanduser().resolve(strict=False)
target = output_config.expanduser().resolve(strict=False)
if checkpoint == run_output or checkpoint in run_output.parents:
raise ValueError(
"The new output directory must not contain the model-initialization checkpoint; "
"use a separate output directory for model-only recovery."
)
training = config["training"]
dyme_args = training["dyme_args"]
model = config["model"]
# These are deliberately the only training-recipe changes. All OPD loss,
# filtering, evaluation, and checkpoint-evaluation settings stay exactly
# those of the validated four-rank recipe.
model["pretrained_model_path"] = str(checkpoint)
model["teacher_device_map"] = "local"
training["num_gpus"] = num_gpus
training["num_client"] = num_gpus
dyme_args["output_dir"] = str(run_output)
if gradient_accumulation_steps is not None:
dyme_args["gradient_accumulation_steps"] = gradient_accumulation_steps
# A model-only process must not attempt to restore the four-rank ZeRO-2
# optimizer/RNG partitions even if a future base config supplies one.
dyme_args.pop("resume_from_checkpoint", None)
target.parent.mkdir(parents=True, exist_ok=True)
temporary = target.with_name(f".{target.name}.tmp-{os.getpid()}")
with temporary.open("w", encoding="utf-8") as handle:
yaml.safe_dump(config, handle, allow_unicode=True, sort_keys=False)
os.replace(temporary, target)
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Materialize a self-contained OPD-only model-only recovery YAML."
)
parser.add_argument("--base-config", type=Path, required=True)
parser.add_argument("--student-checkpoint", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--output-config", type=Path, required=True)
parser.add_argument("--num-gpus", type=int, default=3)
parser.add_argument("--gradient-accumulation-steps", type=int, default=None)
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.num_gpus <= 0:
raise ValueError(f"--num-gpus must be positive, got {args.num_gpus}")
if args.gradient_accumulation_steps is not None and args.gradient_accumulation_steps <= 0:
raise ValueError(
"--gradient-accumulation-steps must be positive when provided, "
f"got {args.gradient_accumulation_steps}"
)
config = _materialize(
base_config=args.base_config,
student_checkpoint=args.student_checkpoint,
output_dir=args.output_dir,
output_config=args.output_config,
num_gpus=args.num_gpus,
gradient_accumulation_steps=args.gradient_accumulation_steps,
)
print(
f"Materialized {args.num_gpus}-rank model-only OPD recovery config: "
f"student={config['model']['pretrained_model_path']} "
f"output_dir={config['training']['dyme_args']['output_dir']} "
f"config={args.output_config.expanduser().resolve(strict=False)}",
flush=True,
)
if __name__ == "__main__":
main()
|