| |
| """Allège un checkpoint d'entraînement pour le déploiement CPU. |
| |
| Un checkpoint d'entraînement contient souvent l'état de l'optimiseur (Adam = |
| 2 copies des poids en plus) : un modèle de 500M de paramètres peut peser |
| plusieurs Go alors que l'inférence n'a besoin que des poids. Ce script ne garde |
| que `args` + `model_state_dict` (sans le préfixe `_orig_mod.`), en fp16 par |
| défaut : ~2 octets/paramètre, soit ~1 Go pour 500M. Au chargement, PyTorch |
| reconvertit en fp32 (le calcul CPU reste en fp32). |
| |
| Le fichier produit garde le même format que l'original : inference.py le |
| charge sans modification. |
| |
| Usage : |
| python scripts/slim_checkpoint.py chat_model_500M.pt chat_model_500M_slim.pt |
| python scripts/slim_checkpoint.py in.pt out.pt --dtype bf16 |
| """ |
| import argparse |
| import dataclasses |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
| import inference |
|
|
| DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} |
|
|
|
|
| def _args_to_dict(args): |
| if isinstance(args, dict): |
| return dict(args) |
| if dataclasses.is_dataclass(args): |
| return dataclasses.asdict(args) |
| return {k: v for k, v in vars(args).items() if not k.startswith("_")} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("src") |
| ap.add_argument("dst") |
| ap.add_argument("--dtype", choices=DTYPES, default="fp16") |
| opts = ap.parse_args() |
| dtype = DTYPES[opts.dtype] |
|
|
| ckpt = inference._torch_load(opts.src) |
| state = {k.replace("_orig_mod.", ""): v for k, v in ckpt["model_state_dict"].items()} |
|
|
| slim, bad = {}, [] |
| for name, tensor in state.items(): |
| if tensor.is_floating_point(): |
| converted = tensor.detach().to(dtype).contiguous() |
| if dtype == torch.float16 and not torch.isfinite(converted).all(): |
| bad.append(name) |
| slim[name] = converted |
| else: |
| slim[name] = tensor.detach().contiguous() |
| if bad: |
| sys.exit(f"❌ Débordement fp16 sur {len(bad)} tenseur(s) (ex. {bad[0]}). Relance avec --dtype bf16.") |
|
|
| out = {"args": _args_to_dict(ckpt["args"]), "model_state_dict": slim, "val_loss": ckpt.get("val_loss")} |
| torch.save(out, opts.dst) |
|
|
| n_params = sum(t.numel() for t in slim.values()) |
| print(f"✅ {n_params / 1e6:.0f}M paramètres | {os.path.getsize(opts.src) / 1e9:.2f} Go -> " |
| f"{os.path.getsize(opts.dst) / 1e9:.2f} Go ({opts.dtype})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
