| |
| """Пересохраняет чекпоинт FANTA с разбиением весов на шарды. |
| |
| Примеры: |
| python3 scripts/reshard_checkpoint.py runs/checkpoint-5000 runs/checkpoint-5000-sharded |
| python3 scripts/reshard_checkpoint.py runs/checkpoint-5000 out --max_shard_size 1GB |
| |
| Результат: |
| out/ |
| config.json |
| model-00001-of-00008.safetensors |
| ... |
| model-00008-of-00008.safetensors |
| model.safetensors.index.json |
| |
| HF (`from_pretrained`) автоматически читает index.json и собирает шарды — склеивать вручную не нужно. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| |
| _PROJECT_ROOT = Path(__file__).resolve().parent.parent |
| if str(_PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(_PROJECT_ROOT)) |
|
|
| from fanta.modeling_fanta import FantaForPredictionAndAnomaly |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("src", help="Папка исходного чекпоинта (с model.safetensors + config.json)") |
| ap.add_argument("dst", help="Папка куда сохранить шардированный чекпоинт") |
| ap.add_argument( |
| "--max_shard_size", |
| default="1GB", |
| help="Размер одного шарда (строка для HF). ~1GB → ~8 шардов для 8GB модели.", |
| ) |
| args = ap.parse_args() |
|
|
| src = Path(args.src) |
| dst = Path(args.dst) |
| if not (src / "config.json").is_file(): |
| print(f"Нет config.json в {src}", file=sys.stderr) |
| sys.exit(1) |
|
|
| print(f"Загружаю {src} …") |
| model = FantaForPredictionAndAnomaly.from_pretrained(src) |
|
|
| dst.mkdir(parents=True, exist_ok=True) |
| print(f"Сохраняю шарды max_shard_size={args.max_shard_size} → {dst}") |
| model.save_pretrained( |
| dst, |
| safe_serialization=True, |
| max_shard_size=args.max_shard_size, |
| ) |
|
|
| shards = sorted(dst.glob("model-*.safetensors")) |
| if not shards: |
| shards = sorted(dst.glob("*.safetensors")) |
| total = sum(p.stat().st_size for p in shards) |
| print(f"Готово. Файлов весов: {len(shards)}, суммарно {total / 1024**3:.2f} GiB") |
| for p in shards: |
| print(f" {p.name} {p.stat().st_size / 1024**3:.2f} GiB") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
